{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process PV catalogues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import join\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from h5py import File\n",
    "from astropy.coordinates import match_coordinates_sky, SkyCoord\n",
    "from astropy import units as u\n",
    "\n",
    "%matplotlib inline\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "SPEED_OF_LIGHT = 299_792.458"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Supernovae data\n",
    "\n",
    "### LOSS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "a2dir = \"/mnt/extraspace/rstiskalek/catalogs/PV/PV_Supranta/A2\"\n",
    "\n",
    "names = [\"z_CMB\", \"mB\", \"x1\", \"c\", \"e_mB\", \"e_x1\", \"e_c\", \"RA\", \"DEC\"]\n",
    "dtype = [(n, np.float32) for n in names]\n",
    "data = np.genfromtxt(join(a2dir, \"loss.csv\"), delimiter=\",\", skip_header=1,\n",
    "                     usecols=[5 + n for n in range(len(names))])\n",
    "\n",
    "loss_data = np.empty(len(data), dtype=dtype)\n",
    "for i, n in enumerate(names):\n",
    "    loss_data[n] = data[:, i]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Foundation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "names = [\"z_CMB\", \"RA\", \"DEC\", \"x1\", \"mB\", \"c\", \"peak\", \"e_peak\", \"e_x1\", \"e_mB\", \"e_c\"]\n",
    "dtype = [(n, np.float32) for n in names]\n",
    "data = np.genfromtxt(join(a2dir, \"foundation.csv\"), delimiter=\",\", skip_header=1,\n",
    "                     usecols=[3 + n for n in range(len(names))])\n",
    "\n",
    "foundation_data = np.empty(len(data), dtype=dtype)\n",
    "for i, n in enumerate(names):\n",
    "    foundation_data[n] = data[:, i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pantheon+, all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<KeysViewHDF5 ['DEGdeg', 'DEHdeg', 'DEJ2000', 'Host', 'IAUC', 'RAGdeg', 'RAHdeg', 'RAJ2000', 'SNID', 'SimbadName', 'e_vpec', 'e_zHD', 'e_zhel', 'recno', 'vpec', 'vpecG', 'zGHD', 'zGcmb', 'zGhel', 'zHD', 'zcmb', 'zhel']>\n"
     ]
    }
   ],
   "source": [
    "fpath_full = \"/mnt/extraspace/rstiskalek/catalogs/PV/Pantheon+SH0ES.dat\"\n",
    "fpath_group = \"/mnt/extraspace/rstiskalek/catalogs/PV/pantheon+_groups.hdf5\"\n",
    "\n",
    "data_full = np.genfromtxt(fpath_full, names=True, dtype=None, encoding=None)\n",
    "data_group = {}\n",
    "\n",
    "# Read in the groups\n",
    "with File(fpath_group, \"r\") as f:\n",
    "    print(f.keys())\n",
    "    for key in f.keys():\n",
    "        try:\n",
    "            data_group[key] = f[key][...]\n",
    "        except IndexError:\n",
    "            print(f\"Failed to read {key}\")\n",
    "        # data[key] = f[key][...]\n",
    "\n",
    "# data = data[data[\"zCMB\"] < 0.1]\n",
    "\n",
    "keys = [\"zCMB\", \"zCMBERR\", \"mB\", \"mBERR\", \"x1\", \"x1ERR\", \"c\", \"cERR\", \"RA\",\n",
    "        \"DEC\", \"VPEC\", \"VPECERR\", \"biasCor_m_b\", \"biasCorErr_m_b\"]\n",
    "pantheon_data = {}\n",
    "for key in keys:\n",
    "    pantheon_data[key] = data_full[key]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Maximum 2D separation is: 0.0005450609539575544 deg\n"
     ]
    }
   ],
   "source": [
    "match_coord = SkyCoord(data_full[\"RA\"], data_full[\"DEC\"], unit=\"deg\")\n",
    "catalogue_coord = SkyCoord(data_group[\"RAJ2000\"], data_group[\"DEJ2000\"],\n",
    "                           unit=\"deg\")\n",
    "\n",
    "idxs, sep2d, sep3d = match_coordinates_sky(match_coord, catalogue_coord, nthneighbor=1)\n",
    "print(f\"Maximum 2D separation is: {sep2d.max()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "zcmb_supernovae = data_group[\"zcmb\"][idxs]\n",
    "zcmb_group = data_group[\"zGcmb\"][idxs]\n",
    "\n",
    "pantheon_data[\"zCMB_SN\"] = zcmb_supernovae\n",
    "pantheon_data[\"zCMB_Group\"] = zcmb_group\n",
    "\n",
    "mask = pantheon_data[\"zCMB\"] < 0.1\n",
    "for key in pantheon_data.keys():\n",
    "    pantheon_data[key] = pantheon_data[key][mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 14036 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.scatter(pantheon_data[\"zCMB_SN\"], pantheon_data[\"zCMB_Group\"], s=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tully-Fisher galaxies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SFI++ galaxies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_folder = \"/mnt/extraspace/rstiskalek/catalogs/PV/PV_Supranta/tf\"\n",
    "\n",
    "names = [\"RA\", \"DEC\", \"z_CMB\", \"mag\", \"eta\", \"e_mag\", \"e_eta\"]\n",
    "\n",
    "dtype = [(n, np.float32) for n in names]\n",
    "data = np.genfromtxt(join(tf_folder, \"sfi_gals_tf.csv\"), delimiter=\",\", skip_header=1,\n",
    "                     usecols=[2 + n for n in range(len(names))])\n",
    "\n",
    "sfi_gals = np.empty(len(data), dtype=dtype)\n",
    "for i, n in enumerate(names):\n",
    "    sfi_gals[n] = data[:, i]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SFI++ galaxies masked"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "names = [\"RA\", \"DEC\", \"z_CMB\", \"mag\", \"eta\", \"e_mag\", \"e_eta\"]\n",
    "\n",
    "dtype = [(n, np.float32) for n in names]\n",
    "data = np.genfromtxt(join(tf_folder, \"sfi_gals_tf_masked.csv\"), delimiter=\",\", skip_header=1,\n",
    "                     usecols=[2 + n for n in range(len(names))])\n",
    "\n",
    "sfi_gals_masked = np.empty(len(data), dtype=dtype)\n",
    "for i, n in enumerate(names):\n",
    "    sfi_gals_masked[n] = data[:, i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### SFI++ groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "names = [\"RA\", \"DEC\", \"z_CMB\", \"r_hMpc\", \"e_r_hMpc\"]\n",
    "\n",
    "dtype = [(n, np.float32) for n in names]\n",
    "data = np.genfromtxt(join(tf_folder, \"sfi_grps.csv\"), delimiter=\",\", skip_header=1,\n",
    "                     usecols=[1 + n for n in range(len(names))])\n",
    "\n",
    "sfi_groups = np.empty(len(data), dtype=dtype)\n",
    "for i, n in enumerate(names):\n",
    "    sfi_groups[n] = data[:, i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Cross $\\texttt{SFI++ galaxies}$ and $\\texttt{SFI++ groups}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "[$3^\\circ40{}^\\prime57.6210022{}^{\\prime\\prime}$\n",
       " $0^\\circ45{}^\\prime09.98210907{}^{\\prime\\prime}$\n",
       " $1^\\circ47{}^\\prime48.02945137{}^{\\prime\\prime}$ ...\n",
       " $0^\\circ59{}^\\prime48.39719296{}^{\\prime\\prime}$\n",
       " $1^\\circ27{}^\\prime20.76476097{}^{\\prime\\prime}$\n",
       " $2^\\circ01{}^\\prime03.90094757{}^{\\prime\\prime}$]"
      ],
      "text/plain": [
       "<Angle [3.6826725, 0.7527728, 1.7966748, ..., 0.996777 , 1.455768 ,\n",
       "        2.0177503] deg>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "match_coord = SkyCoord(sfi_gals[\"RA\"], sfi_gals[\"DEC\"], unit=\"deg\")\n",
    "catalogue_coord = SkyCoord(sfi_groups[\"RA\"], sfi_groups[\"DEC\"],\n",
    "                           unit=\"deg\")\n",
    "\n",
    "idxs, sep2d, sep3d = match_coordinates_sky(match_coord, catalogue_coord, nthneighbor=1)\n",
    "sep2d.to(u.degree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15200 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "m = sep2d.value < 0.5\n",
    "plt.figure()\n",
    "plt.hist(sep2d, bins=\"auto\")\n",
    "plt.xlabel(\"Separation [deg]\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/sfi_gals_to_sfi_groups.png\", bbox_inches=\"tight\",\n",
    "            dpi=450)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2MTF "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "names = [\"RA\", \"DEC\", \"mag\", \"e_mag\", \"z_CMB\", \"r_hMpc\", \"e_rhMpc\", \"M\", \"eta\", \"e_eta\"]\n",
    "\n",
    "dtype = [(n, np.float32) for n in names]\n",
    "data = np.genfromtxt(join(tf_folder, \"twomtf_k.csv\"), delimiter=\",\", skip_header=1,\n",
    "                     usecols=[2 + n for n in range(len(names))])\n",
    "\n",
    "twomtf_gals = np.empty(len(data), dtype=dtype)\n",
    "for i, n in enumerate(names):\n",
    "    twomtf_gals[n] = data[:, i]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write to HDF5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "outdir = \"/mnt/extraspace/rstiskalek/catalogs\"\n",
    "fname = \"PV_compilation.hdf5\"\n",
    "\n",
    "with File(join(outdir, fname), 'w') as f:\n",
    "    # Write LOSS\n",
    "    grp = f.create_group(\"LOSS\")\n",
    "    for name in loss_data.dtype.names:\n",
    "        grp.create_dataset(name, data=loss_data[name])\n",
    "\n",
    "    # Write Foundation\n",
    "    grp = f.create_group(\"Foundation\")\n",
    "    for name in foundation_data.dtype.names:\n",
    "        grp.create_dataset(name, data=foundation_data[name])\n",
    "\n",
    "    # Write SFI gals\n",
    "    grp = f.create_group(\"SFI_gals\")\n",
    "    for name in sfi_gals.dtype.names:\n",
    "        grp.create_dataset(name, data=sfi_gals[name])\n",
    "\n",
    "    # Write SFI gals masked\n",
    "    grp = f.create_group(\"SFI_gals_masked\")\n",
    "    for name in sfi_gals_masked.dtype.names:\n",
    "        grp.create_dataset(name, data=sfi_gals_masked[name])\n",
    "\n",
    "    # Write SFI groups\n",
    "    grp = f.create_group(\"SFI_groups\")\n",
    "    for name in sfi_groups.dtype.names:\n",
    "        grp.create_dataset(name, data=sfi_groups[name])\n",
    "\n",
    "    # Write 2MTF gals\n",
    "    grp = f.create_group(\"2MTF\")\n",
    "    for name in twomtf_gals.dtype.names:\n",
    "        grp.create_dataset(name, data=twomtf_gals[name])\n",
    "\n",
    "    # Write Pantheon\n",
    "    grp = f.create_group(\"Pantheon+\")\n",
    "    for name in pantheon_data.keys():\n",
    "        grp.create_dataset(name, data=pantheon_data[name])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv_csiborg",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
