{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (C) 2024 Richard Stiskalek\n",
    "# This program is free software; you can redistribute it and/or modify it\n",
    "# under the terms of the GNU General Public License as published by the\n",
    "# Free Software Foundation; either version 3 of the License, or (at your\n",
    "# option) any later version.\n",
    "#\n",
    "# This program is distributed in the hope that it will be useful, but\n",
    "# WITHOUT ANY WARRANTY; without even the implied warranty of\n",
    "# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General\n",
    "# Public License for more details.\n",
    "#\n",
    "# You should have received a copy of the GNU General Public License along\n",
    "# with this program; if not, write to the Free Software Foundation, Inc.,\n",
    "# 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from h5py import File\n",
    "import csiborgtools\n",
    "\n",
    "SPEED_OF_LIGHT = 299792.458  # km / s\n",
    "\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read in upGLADE\n",
    "\n",
    "- Mask out galaxies with bad redshifts\n",
    "- Convert heliocentric redshifts to the CMB frame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/mnt/users/rstiskalek/csiborgtools/data/upglade_all_z0p05_new.h5\"\n",
    "data = {}\n",
    "with File(fname, \"r\") as f:\n",
    "    for i, key in enumerate([\"RA\", \"dec\", \"zhelio\", \"e_zhelio\"]):\n",
    "        data[key] = f[\"data\"][\"block0_values\"][:, i]\n",
    "data[\"DEC\"] = data.pop(\"dec\")\n",
    "\n",
    "print(f\"Initially, we have {data['zhelio'].size} objects.\")\n",
    "\n",
    "# Ask about this\n",
    "mask = (data[\"e_zhelio\"] < 0) | (data[\"zhelio\"] < 0)\n",
    "print(f\"Masking {mask.sum()} objects that have `e_zhelio` < 0 or `zhelio` < 0.\")\n",
    "for key in data.keys():\n",
    "    data[key][mask] = np.nan\n",
    "\n",
    "mask = (data[\"e_zhelio\"] / data[\"zhelio\"] > 0.15)\n",
    "print(f\"Masking {mask.sum()} objects that have `e_zhelio` / `zhelio` > 0.15.\")\n",
    "for key in data.keys():\n",
    "    data[key][mask] = np.nan\n",
    "\n",
    "print(f\"Finally, we have {np.sum(np.isfinite(data['zhelio']))} objects.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(data[\"zhelio\"], data[\"e_zhelio\"], s=0.1)\n",
    "\n",
    "plt.yscale(\"log\")\n",
    "plt.xscale(\"log\")\n",
    "plt.xlabel(r\"$z_{\\rm helio}$\")\n",
    "plt.ylabel(r\"$\\sigma_{z_{\\rm helio}}$\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "zcmb, e_zcmb = csiborgtools.heliocentric_to_cmb(data[\"zhelio\"], data[\"RA\"], data[\"DEC\"], data[\"e_zhelio\"])\n",
    "data[\"zcmb\"] = zcmb\n",
    "data[\"e_zcmb\"] = e_zcmb\n",
    "\n",
    "\n",
    "mask = (data[\"zcmb\"] > 0.06) #& ~np.isnan(data[\"zhelio\"])\n",
    "print(f\"Masking {mask.sum()} objects that have `zcmb` > 0.06.\")\n",
    "for key in data.keys():\n",
    "    data[key][mask] = np.nan\n",
    "\n",
    "print(f\"Finally, we have {np.sum(np.isfinite(data['zhelio']))} objects.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.scatter(data[\"zcmb\"], data[\"e_zcmb\"], s=0.001)\n",
    "plt.xlabel(r\"$z_{\\rm CMB}$\")\n",
    "plt.ylabel(r\"$\\sigma_{z_{\\rm CMB}}$\")\n",
    "\n",
    "plt.xlim(0, 0.05)\n",
    "plt.ylim(0, 0.008)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/UPGLADE_zcmb_vs_ezcmb.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Write only masked galaxies to this file, but also save the mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname_here = fname.replace(\".h5\", \"_PROCESSED.h5\")\n",
    "mask = np.isfinite(data[\"RA\"])\n",
    "print(f\"Writing {mask.sum()} objects to `{fname_here}`.\")\n",
    "\n",
    "with File(fname_here, \"w\") as f:\n",
    "    for key in data.keys():\n",
    "        f[key] = data[key][mask]\n",
    "\n",
    "    f[\"mask\"] = mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Having generated this file, next step is to run `field_los.py` to evaluate the density and velocity field along the LOS of each object that is not masked.\n",
    "- Then, the next step is to run `post_upglade.py` to calculate the cosmological redshift of each object in UPGLADE.\n",
    "    - Based on Carrick2015 samples calibrated against Pantheon+"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results verification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/UPGLADE/zcosmo_UPGLADE.hdf5\"\n",
    "\n",
    "with File(fname, \"r\") as f:\n",
    "    indxs = f[\"indxs\"][:]\n",
    "    mean_zcosmo = f[\"mean_zcosmo\"][:]\n",
    "    std_zcosmo = f[\"std_zcosmo\"][:]\n",
    "\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(mean_zcosmo, std_zcosmo, s=0.001)\n",
    "\n",
    "plt.xlabel(r\"$z_{\\rm cosmo}$\")\n",
    "plt.ylabel(r\"$\\sigma_{z_{\\rm cosmo}}$\")\n",
    "\n",
    "plt.xlim(0, 0.05)\n",
    "plt.ylim(0, 0.006)\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/UPGLADE_zcosmo_vs_ezcosmo.png\")\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Combine datasets\n",
    "i.e. match to the original UPGLADE file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data[\"zcosmo\"] = np.zeros_like(data[\"RA\"])\n",
    "data[\"zcosmo\"][mask] = mean_zcosmo\n",
    "\n",
    "data[\"e_zcosmo\"] = np.zeros_like(data[\"RA\"])\n",
    "data[\"e_zcosmo\"][mask] = std_zcosmo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(data[\"zcmb\"], data[\"zcosmo\"], s=0.01)\n",
    "plt.axline([0, 0], slope=1, color=\"red\", ls=\"--\")\n",
    "\n",
    "plt.xlabel(r\"$z_{\\rm CMB}$\")\n",
    "plt.ylabel(r\"$z_{\\rm cosmo}$\")\n",
    "plt.xlim(0)\n",
    "plt.ylim(0)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/UPGLADE_zcmb_vs_zcosmo.png\")\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(data[\"zcmb\"], (data[\"zcosmo\"] - data[\"zcmb\"]) * SPEED_OF_LIGHT, s=0.001)\n",
    "\n",
    "plt.xlabel(r\"$z_{\\rm CMB}$\")\n",
    "plt.ylabel(r\"$c (z_{\\rm cosmo} - z_{\\rm CMB}) ~ [\\mathrm{km} / \\mathrm{s}]$\")\n",
    "plt.xlim(0)\n",
    "plt.ylim(-1000, 1000)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"../../plots/UPGLADE_zcmb_vs_dzcosmo.png\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Save the data\n",
    "- In a format matching what Gergely shared."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = \"/mnt/users/rstiskalek/csiborgtools/data/upglade_z_0p05_all_WZCOSMO.h5\"\n",
    "print(f\"Writing to `{fname}`.\")\n",
    "\n",
    "with File(fname, \"w\") as f:\n",
    "    for key in data.keys():\n",
    "        f.create_dataset(key, data=data[key])"
   ]
  },
  {
   "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
}