{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scienceplots\n",
    "import seaborn as sns\n",
    "import healpy as hp\n",
    "from h5py import File\n",
    "\n",
    "\n",
    "from reconstruction_comparison import *\n",
    "SPEED_OF_LIGHT = 299_792.458\n",
    "from csiborgtools import radec_to_galactic\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paper plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with File(\"/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5\", 'r') as f:\n",
    "    SFI = {}\n",
    "    for key in f[\"SFI_gals\"].keys():\n",
    "        SFI[key] = f[\"SFI_gals\"][key][...]\n",
    "    \n",
    "    TWOMTF = {}\n",
    "    for key in f[\"2MTF\"].keys():\n",
    "        TWOMTF[key] = f[\"2MTF\"][key][...]\n",
    "    \n",
    "    FOUNDATION = {}\n",
    "    for key in f[\"Foundation\"].keys():\n",
    "        FOUNDATION[key] = f[\"Foundation\"][key][...]\n",
    "    \n",
    "    LOSS = {}\n",
    "    for key in f[\"LOSS\"].keys():\n",
    "        LOSS[key] = f[\"LOSS\"][key][...]\n",
    "\n",
    "\n",
    "\n",
    "with File(\"/mnt/extraspace/rstiskalek/catalogs/PV/CF4/CF4_TF-distances.hdf5\", 'r') as f:\n",
    "    CF4_TFR = {}\n",
    "    for key in f.keys():\n",
    "        CF4_TFR[key] = f[key][...]\n",
    "    \n",
    "    RA = f[\"RA\"][...] * 360 / 24\n",
    "    CF4_TFR[\"RA\"] = RA\n",
    "    \n",
    "    CF4_TFR[\"DEC\"] = f[\"DE\"][...]\n",
    "\n",
    "    l, b = radec_to_galactic(CF4_TFR[\"RA\"], CF4_TFR[\"DE\"])\n",
    "    m = np.abs(b) > 7.5\n",
    "    for key in CF4_TFR.keys():\n",
    "        CF4_TFR[key] = CF4_TFR[key][m]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Redshift distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 43816 characters]",
      "text/plain": [
       "<Figure size 350x262.5 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with plt.style.context('science'):\n",
    "    plt.rcParams.update({'font.size': 9})\n",
    "    plt.figure()\n",
    "\n",
    "    sns.kdeplot(LOSS[\"z_CMB\"], label=\"LOSS\", fill=True, bw_adjust=0.75)\n",
    "    sns.kdeplot(FOUNDATION[\"z_CMB\"], label=\"Foundation\", fill=True, bw_adjust=0.75)\n",
    "    sns.kdeplot(SFI[\"z_CMB\"], label=\"SFI\", fill=True, bw_adjust=0.75)\n",
    "    sns.kdeplot(TWOMTF[\"z_CMB\"], label=\"2MTF\", fill=True, bw_adjust=0.75)\n",
    "\n",
    "    m = CF4_TFR[\"i\"] > 5\n",
    "    sns.kdeplot(CF4_TFR[\"Vcmb\"][m] / SPEED_OF_LIGHT, label=catalogue_to_pretty(\"CF4_TFR_i\"), fill=True)\n",
    "\n",
    "    m = CF4_TFR[\"w1\"] > 5\n",
    "    sns.kdeplot(CF4_TFR[\"Vcmb\"][m] / SPEED_OF_LIGHT, label=catalogue_to_pretty(\"CF4_TFR_w1\"), fill=True)\n",
    "\n",
    "    plt.xlabel(r\"$z_{\\rm CMB}$\")\n",
    "    plt.ylabel(\"Normalised PDF\")\n",
    "\n",
    "    plt.legend()\n",
    "    plt.xlim(0, 0.07)\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(f\"../../plots/zcmb_dist.pdf\", dpi=450)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Sky distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 133420 characters]",
      "text/plain": [
       "<Figure size 830x332 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with plt.style.context('science'):\n",
    "    plt.rcParams.update({'font.size': 9})\n",
    "    figwidth = 8.3\n",
    "    fig, axs = plt.subplots(2, 3, figsize=(figwidth, 0.4 * figwidth),\n",
    "                            sharex=True, sharey=True)\n",
    "\n",
    "\n",
    "    names = [\"LOSS\", \"Foundation\", \"SFI\", \"2MTF\", \"CF4_TFR_i\", \"CF4_TFR_w1\"]\n",
    "    vals = [LOSS, FOUNDATION, SFI, TWOMTF, CF4_TFR, CF4_TFR]\n",
    "    for i, ax in enumerate(axs.flat):\n",
    "        RA = vals[i][\"RA\"]\n",
    "        DEC = vals[i][\"DEC\"]\n",
    "\n",
    "        if names[i] == \"CF4_TFR_i\":\n",
    "            m = vals[i][\"i\"] > 5\n",
    "            RA = RA[m]\n",
    "            DEC = DEC[m]\n",
    "\n",
    "        if names[i] == \"CF4_TFR_w1\":\n",
    "            m = vals[i][\"w1\"] > 5\n",
    "            RA = RA[m]\n",
    "            DEC = DEC[m]\n",
    "\n",
    "        sns.scatterplot(x=RA, y=DEC, ax=ax, s=1.5, rasterized=True)\n",
    "\n",
    "        # Add panel name in top-left corner of each subplot\n",
    "        ax.text(0.04, 0.16, catalogue_to_pretty(names[i]), transform=ax.transAxes, \n",
    "                fontsize=9, verticalalignment='top', \n",
    "                bbox=dict(boxstyle=\"round\", facecolor='lightgray', alpha=1),\n",
    "                zorder=10)  # Ensure text is on top of other plot elements\n",
    "\n",
    "    axs[0, 0].set_xlim(0, 360)\n",
    "    axs[0, 0].set_ylim(-90, 90)\n",
    "\n",
    "    for i in range(3):\n",
    "        axs[-1, i].set_xlabel(r\"$\\mathrm{RA} ~ [\\mathrm{deg}]$\")\n",
    "\n",
    "    for i in range(2):\n",
    "        axs[i, 0].set_ylabel(r\"$\\delta ~ [\\mathrm{deg]}$\")\n",
    "\n",
    "    fig.tight_layout()\n",
    "    plt.savefig(f\"../../plots/sky_dist.pdf\", dpi=350)\n",
    "    fig.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 118196 characters]",
      "text/plain": [
       "<Figure size 350x367.5 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "with plt.style.context('science'):\n",
    "    plt.rcParams.update({'font.size': 9})\n",
    "\n",
    "    RA = CF4_TFR[\"RA\"]\n",
    "    DEC = CF4_TFR[\"DEC\"]\n",
    "\n",
    "    fig, axs = plt.subplots(2, 1, figsize=(3.5, 1.5 * 2.45), sharex=True, sharey=True)\n",
    "    fig.subplots_adjust(hspace=0.0, wspace=0.0)\n",
    "\n",
    "    m = CF4_TFR[\"w1\"] > 5\n",
    "    sns.scatterplot(x=RA[m], y=DEC[m], s=5, rasterized=True, ax=axs[0])\n",
    "    axs[0].text(0.15, 0.15, catalogue_to_pretty(\"CF4_TFR_w1\"),\n",
    "                transform=axs[0].transAxes, fontsize=\"small\",\n",
    "                verticalalignment='center', horizontalalignment='center',\n",
    "                bbox=dict(facecolor='white', alpha=0.75)\n",
    "                )\n",
    "\n",
    "    m = CF4_TFR[\"i\"] > 5\n",
    "    sns.scatterplot(x=RA[m], y=DEC[m], s=5, rasterized=True, ax=axs[1])\n",
    "    axs[1].text(0.15, 0.15, catalogue_to_pretty(\"CF4_TFR_i\"),\n",
    "                transform=axs[1].transAxes, fontsize=\"small\",\n",
    "                verticalalignment='center', horizontalalignment='center',\n",
    "                bbox=dict(facecolor='white', alpha=0.75))\n",
    "\n",
    "    axs[0].set_xlim(0, 360)\n",
    "    axs[0].set_ylim(-90, 90)\n",
    "\n",
    "    for i in range(2):\n",
    "        axs[i].set_ylabel(r\"$\\delta ~ [\\mathrm{deg}]$\")\n",
    "    axs[-1].set_xlabel(r\"$\\mathrm{RA} ~ [\\mathrm{deg}]$\")\n",
    "\n",
    "    fig.tight_layout()\n",
    "    fig.savefig(f\"../../plots/sky_dist_CF4.pdf\", dpi=350)\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Experimenting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "l, b = radec_to_galactic(CF4_TFR[\"RA\"], CF4_TFR[\"DE\"])\n",
    "# l, b = radec_to_galactic(SFI[\"RA\"], SFI[\"DEC\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.06692355049302931"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.abs(b).min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "232"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(np.abs(b) < 10).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11204 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.hist(np.abs(b), bins=\"auto\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<KeysViewHDF5 ['DEC', 'RA', 'e_eta', 'e_mag', 'eta', 'mag', 'z_CMB']>\n"
     ]
    }
   ],
   "source": [
    "with File(\"/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5\", 'r') as f:\n",
    "    print(f[\"SFI_gals\"].keys())\n",
    "\n",
    "    zCMB_sfi = f[\"SFI_gals\"][\"z_CMB\"][...]\n",
    "    mag_sfi = f[\"SFI_gals\"][\"mag\"][...]\n",
    "    eta_sfi = f[\"SFI_gals\"][\"eta\"][...]\n",
    "    e_mag_sfi = f[\"SFI_gals\"][\"e_mag\"][...]\n",
    "    e_eta_sfi = f[\"SFI_gals\"][\"e_eta\"][...]\n",
    "\n",
    "    zCMB_2MTF = f[\"2MTF\"][\"z_CMB\"][...]\n",
    "    mag_2MTF = f[\"2MTF\"][\"mag\"][...]\n",
    "    eta_2MTF = f[\"2MTF\"][\"eta\"][...]\n",
    "    e_mag_2MTF = f[\"2MTF\"][\"e_mag\"][...]\n",
    "    e_eta_2MTF = f[\"2MTF\"][\"e_eta\"][...]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "with File(\"/mnt/extraspace/rstiskalek/catalogs/PV/CF4/CF4_TF-distances.hdf5\", 'r') as f:\n",
    "    zCMB = f[\"Vcmb\"][...] / 3e5\n",
    "    rband = f[\"r\"][...]\n",
    "    iband = f[\"i\"][...]\n",
    "    w1band = f[\"w1\"][...]\n",
    "    RA = f[\"RA\"][...] * 360 / 24\n",
    "    dec = f[\"DE\"][...]\n",
    "\n",
    "# m  = (iband > 5) #& (w1band > 5)\n",
    "m  = (w1band> 5) #& (w1band > 5)\n",
    "w1band = w1band[m]\n",
    "iband = iband[m]\n",
    "RA = RA[m]\n",
    "dec = dec[m]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "dmag = iband - w1band"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 205012 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.scatter(RA, dec, c=dmag, s=1)\n",
    "plt.colorbar()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 11836 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.hist(iband - w1band, bins=\"auto\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 10952 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "m = rband > 0\n",
    "\n",
    "plt.figure()\n",
    "plt.hist(rband[m], bins=\"auto\")\n",
    "plt.axvline(17.6, c=\"red\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 16080 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "plt.hist(zCMB, bins=\"auto\")\n",
    "plt.xlabel(r\"$z_{\\rm CMB}$\")\n",
    "plt.ylabel(\"Counts\")\n",
    "plt.xlim(0)\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/zCMB_CF4.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 394564 characters]",
      "text/plain": [
       "<Figure size 1337.25x1250 with 20 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import seaborn as sns\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "df_sfi = pd.DataFrame({\n",
    "    r\"$z_{\\rm CMB}$\": zCMB_sfi,\n",
    "    r\"$m$\": mag_sfi,\n",
    "    r\"$\\eta$\": eta_sfi,\n",
    "    r\"$\\sigma_{m}$\": e_mag_sfi,\n",
    "    r\"$\\sigma_\\eta$\": e_eta_sfi,\n",
    "    \"Dataset\": \"SFI\"\n",
    "})\n",
    "\n",
    "df_2mtf = pd.DataFrame({\n",
    "    r\"$z_{\\rm CMB}$\": zCMB_2MTF,\n",
    "    r\"$m$\": mag_2MTF,\n",
    "    r\"$\\eta$\": eta_2MTF,\n",
    "    r\"$\\sigma_{m}$\": e_mag_2MTF,\n",
    "    r\"$\\sigma_\\eta$\": e_eta_2MTF,\n",
    "    \"Dataset\": \"2MTF\"\n",
    "})\n",
    "\n",
    "df_combined = pd.concat([df_sfi, df_2mtf])\n",
    "g = sns.pairplot(df_combined, hue=\"Dataset\", kind=\"kde\", diag_kind=\"kde\",\n",
    "             palette=\"Set1\", corner=True, plot_kws={\"alpha\": 0.75},\n",
    "             diag_kws={\"fill\": True})\n",
    "sns.move_legend(g, \"upper right\", bbox_to_anchor=(0.7, 0.7))\n",
    "\n",
    "# Show the plot\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/TFR_data.pdf\", dpi=450)\n",
    "plt.show()\n"
   ]
  },
  {
   "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
}
