csiborgtools/notebooks/flow_mapping.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using a calibrated flow model to predict $z_{\\rm cosmo}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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",
    "\n",
    "import csiborgtools\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "\n",
    "paths = csiborgtools.read.Paths(**csiborgtools.paths_glamdring)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_calibration(catalogue, simname, nsim, ksmooth):\n",
    "    fname = f\"/mnt/extraspace/rstiskalek/csiborg_postprocessing/peculiar_velocity/flow_samples_{catalogue}_{simname}_smooth_{ksmooth}.hdf5\"  # noqa\n",
    "    keys = [\"Vext_x\", \"Vext_y\", \"Vext_z\", \"alpha\", \"beta\", \"sigma_v\"]\n",
    "\n",
    "    # SN_keys = ['mag_cal', 'alpha_cal', 'beta_cal']\n",
    "    SN_keys = []\n",
    "    calibration_samples = {}\n",
    "    with File(fname, 'r') as f:\n",
    "        for key in keys:\n",
    "            calibration_samples[key] = f[f\"sim_{nsim}/{key}\"][:]\n",
    "\n",
    "        for key in SN_keys:\n",
    "            calibration_samples[key] = f[f\"sim_{nsim}/{key}\"][:]\n",
    "\n",
    "    return calibration_samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test running a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23:16:25: reading the catalogue.\n",
      "23:16:25: reading the interpolated field.\n",
      "23:16:25: calculating the radial velocity.\n",
      "Selected 100/100 galaxies.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/users/rstiskalek/csiborgtools/csiborgtools/flow/flow_model.py:113: UserWarning: The number of radial steps is even. Skipping the first step at 0.0 because Simpson's rule requires an odd number of steps.\n",
      "  warn(f\"The number of radial steps is even. Skipping the first \"\n"
     ]
    }
   ],
   "source": [
    "# fpath_data = \"/mnt/extraspace/rstiskalek/catalogs/PV_compilation_Supranta2019.hdf5\"\n",
    "fpath_data = \"/mnt/extraspace/rstiskalek/catalogs/PV_mock_CB2_17417_small.hdf5\"\n",
    "\n",
    "simname = \"csiborg2_main\"\n",
    "catalogue = \"CB2_small\"\n",
    "# nsim = 0\n",
    "loader = csiborgtools.flow.DataLoader(simname, 19 - 5, catalogue, fpath_data, paths, ksmooth=1)\n",
    "\n",
    "calibration_samples = load_calibration(catalogue, simname, 17417 - 500, 1)\n",
    "flow_model = csiborgtools.flow.get_model(loader, zcmb_max=0.07)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = csiborgtools.flow.Observed2CosmologicalRedshift(calibration_samples, loader.rdist, loader._Omega_m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [],
   "source": [
    "# zcosmo_mean, zcosmo_std = flow_model.predict_zcosmo_from_calibration(\n",
    "    # calibration_samples[\"mag_cal\"], calibration_samples[\"alpha_cal\"], calibration_samples[\"beta_cal\"])\n",
    "zcosmo_mean = loader.cat[\"zcosmo\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Marginalizing:   0%|          | 0/5000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Marginalizing: 100%|██████████| 5000/5000 [00:02<00:00, 2262.93it/s]\n"
     ]
    }
   ],
   "source": [
    "# n = 2 is a very good test\n",
    "# CONVERT DEGREES TO RADIANS!!\n",
    "n = 40\n",
    "zcos, post = model.posterior_zcosmo(\n",
    "    loader.cat[\"zobs\"][n], np.deg2rad(loader.cat[\"RA\"][n]), np.deg2rad(loader.cat[\"DEC\"][n]),\n",
    "    loader.los_density[n], loader.los_radial_velocity[n])\n",
    "# zcos, post_bad = model.posterior_zcosmo(\n",
    "#     loader.cat[\"zobs\"][n], loader.cat[\"RA\"][n], loader.cat[\"DEC\"][n],\n",
    "#     loader.los_density[n], loader.los_radial_velocity[n])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAnYAAAHWCAYAAAD6oMSKAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMywgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/H5lhTAAAACXBIWXMAAA9hAAAPYQGoP6dpAABfLElEQVR4nO3deVxU5eIG8GdmWEZANlEWxVBERRFRQFwyKylod8vletPIrDS9FqVFN6VSL2pmlJqW95prLm1mZmiXxFzAFVRwSU1FlAGFYARkcTi/P/jNXElQloH3zMzz/XzORxzOnHmOp/DxvOc9RyFJkgQiIiIiMnlK0QGIiIiIyDhY7IiIiIjMBIsdERERkZlgsSMiIiIyEyx2RERERGaCxY6IiIjITLDYEREREZkJFjsiIiIiM2ElOoApqKysxNWrV9GyZUsoFArRcYiIiMiCSJKEGzduwMvLC0rl3c/JsdjVwdWrV+Ht7S06BhEREVmwy5cvo127dnddh8WuDlq2bAmg6g/U0dFRcBoiInGKy4vh9ZEXAODqG1dhb2MvOBGhuBjwqjomuHoVsOcxMTdarRbe3t6GPnI3LHZ1oB9+dXR0ZLEjIoumKlcB6qqvHR0dWezkQKX639eOjix2Zqwul4Nx8gQRERGRmWCxIyIiIjITLHZEREREZoLX2BEREZkpnU6HiooK0THoHqytraG6/VrJRmCxIyIiMjOSJEGj0aCgoEB0FKojZ2dneHh4NPp+uSx2REREZkZf6tq0aQM7OzveXF/GJElCSUkJcnNzAQCenp6N2h6LHRERkRnR6XSGUteqVSvRcagOWrRoAQDIzc1FmzZtGjUsy8kTREREZkR/TZ2dnZ3gJFQf+uPV2GsiWeyIiIjMEIdfTYuxjheLHREREZGZYLEjIiIiMhMsdkRERGRS8vLy0KZNG1y8eFF0lDoZPXo0Pvroo2b5LFkWu6VLl8LHxwdqtRphYWE4ePBgret+9913CAkJgbOzM+zt7REUFIS1a9dWW+f555+HQqGotkRGRjb1bhAREVETmDt3Lp555hn4+PiIjlIn7777LubOnYvCwsIm/yzZFbtNmzYhOjoasbGxOHr0KHr27ImIiAjD/V3+ytXVFf/85z+RnJyM48ePIyoqClFRUdixY0e19SIjI5GdnW1YNmzY0By7Q0REREZUUlKC//znP5gwYYLoKHUWEBAAX19frFu3rsk/S3bFbtGiRZg4cSKioqLQrVs3LF++HHZ2dli5cmWN6z/44IMYOnQo/P394evri2nTpiEwMBB79+6ttp6trS08PDwMi4uLS3PsDhEREdWDRqOBQqHAJ598gl69ekGtVqN79+6Gv9e3b98OW1tb9O3bt9r7/vWvf90xOqdQKBAfHy9gL+701FNPYePGjU3+ObIqduXl5Thy5AjCw8MNrymVSoSHhyM5Ofme75ckCYmJiThz5gweeOCBat9LSkpCmzZt0KVLF0yaNAl5eXlGz09E5uHmzZt466238OOPP4qOQmRx0tLSAAArV65EfHw80tLS0L59e4wdOxaVlZXYs2cPgoOD73jf1KlTq43MTZw4Effddx9GjBjRzHtQsz59+uDgwYMoKytr0s+R1ZMnrl+/Dp1OB3d392qvu7u74/Tp07W+r7CwEG3btkVZWRlUKhU+++wzPPLII4bvR0ZGYtiwYejQoQPOnz+Pd955B4899hiSk5NrvLtzWVlZtT94rVZrhL0jIlPxj3/8A//+97+xYsUKaDQa2NjYiI5E1Cj6x1aJUN9Hmh07dgzW1tb44YcfDNfQzZkzByEhIbhy5QouXboELy+vO97XsmVLtGzZEgAwc+ZM7Ny5E0lJSWjXrp1R9qOxvLy8UF5eDo1Gg/vuu6/JPkdWxa6hWrZsibS0NBQVFSExMRHR0dHo2LEjHnzwQQBVs1H0evTogcDAQPj6+iIpKQmDBw++Y3txcXF4//33mys+EcnI+vXr8e9//xsA8Oeff2LHjh146qmnBKciapySkhI4ODgI+eyioiLY29vXef20tDQMGzas2sQIR0dHw9c3b96EWq2u9f2zZs3C2rVrkZSUJKvJFfrHhjV1wZbVUKybmxtUKhVycnKqvZ6TkwMPD49a36dUKtGpUycEBQXhjTfewIgRIxAXF1fr+h07doSbmxvOnTtX4/djYmJQWFhoWC5fvtywHSIik3LmzBm8/PLLAIC2bdsCAL766iuRkYgsTlpaGoKCgqq9lpycDDc3N7Rt2xZubm74888/a3xvbGws1qxZc0epW7lyJQIDA9GzZ0+8+eabAID58+cjICAAPXr0wPr16wFUldDIyEj06NEDPXr0wI4dO3Dx4kX07NkTY8eOhZ+fHyZNmoQtW7YgLCwMAQEBOHv2rOFzatqmXn5+PgCgdevWjf0juitZnbGzsbFBcHAwEhMTMWTIEABAZWUlEhMTMWXKlDpvp7Ky8q5j2FlZWcjLy4Onp2eN37e1tYWtrW29shORabt58yZGjhyJ4uJiPPTQQ4iLi0Pfvn3xww8/4MaNG4YhHiJTZGdnh6KiImGfXVc3b97E2bNnodPpDK9VVlYiPj4e48ePh1KpRK9evWqcXRobG4vVq1ffUepOnDiBjz/+GHv27IGzszPy8/Nx6NAhbN68GYcPH0ZJSQlCQ0Px0EMPITk5Ga1atUJCQgIkScKNGzeQn5+PU6dOYfPmzejUqRMCAgLg4OCAAwcO4PPPP8eSJUvwySef1LpN/bBxeno62rVrBzc3t4b/YdaBrM7YAUB0dDRWrFiB1atX49SpU5g0aRKKi4sRFRUFABg3bhxiYmIM68fFxeGXX37BH3/8gVOnTuGjjz7C2rVr8fe//x1AVfuePn06UlJScPHiRSQmJuKZZ55Bp06dEBERIWQfiUh+XnvtNRw/fhxt2rTB+vXr0adPH/j5+eHmzZv44YcfRMcjahSFQgF7e3shS32urztx4gQUCgXWrVuH5ORknDp1CqNGjUJBQQHeffddAEBERAQyMjKqnbWbM2cOli1bho0bN0KtVkOj0UCj0aCsrAy7du3CqFGj4OzsDKDqNmn79u3D8OHDoVar4erqisGDB+PQoUPo0aMHfvvtN8yYMQMpKSmGIeAuXbqgS5cuUKlU8Pf3N0zy7NGjh+EmybVtU2/Pnj149NFHG3MY60R2xW7UqFFYuHAhZs2ahaCgIKSlpSEhIcEwoSIzMxPZ2dmG9YuLizF58mR0794dAwYMwLfffot169bhxRdfBACoVCocP34cTz/9NDp37owJEyYgODgYe/bs4Vk5IgIAbNiwAV988QUUCgXWr18PT09PKBQK/O1vfwPA4Vii5pKWloauXbvinXfewfDhwxESEgKdTofdu3cbilmPHj3Qu3dvbN68GUDVxJAPP/wQ165dQ79+/eDp6WlYjh8/Xq/P79y5M9LS0tC9e3dER0djyZIlAFCtLyiVSsPvlUpltbOLtSktLcWWLVswceLEeuVpEInuqbCwUAIgFRYWio5CREZ25swZycHBQQIgzZw5847vAZBUKpWUm5srKKG8FJUVSXgPEt6DVFRWJDoOSZIkFRVJElC1FBVJN2/elE6ePCndvHlTdLJ6mzx5sjRmzJh7rrdt2zbJ399f0ul091z3xIkTUkBAgFRQUCBJkiTl5eVJhw4dkoKDg6XS0lIpPz9f8vX1la5evSpduXLF8Oe2ceNGacKECdKFCxek4OBgw/aGDx8u7dq1S5IkSUpOTpaeeOIJSZKkWrcpSZL02WefSY888shdc97tuNWnh8jqGjsiouZUWlqKkSNHoqioCIMGDUJsbGy173fu3BkhISE4fPgwvv76a0yePFlQUiLLkJaWVqdZ6E888QTOnj2LK1euwNvb+67rBgQEYNq0aRgwYACsrKzw6KOPYsGCBXj22WcRHBwMhUKB999/H56entixYwfefPNNqFQqtGjRAv/5z3/qnD0kJKTGbQKAtbU1Fi9eXOdtNYZCkiSpWT7JhGm1Wjg5OaGwsLDalGsiMm2TJk3C8uXL0bp1a6SlpdV4b6yPP/4Y0dHR6N+/P/bt2ycgpbwUlxf
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "\n",
    "mask = post > 1e-5\n",
    "plt.plot(zcos[mask], post[mask], color=\"black\", label=r\"$p(z_{\\rm cosmo})$\")\n",
    "# plt.plot(zcos[mask], post_bad[mask], color=\"black\", ls=\"dashed\")\n",
    "\n",
    "\n",
    "plt.ylim(0)\n",
    "plt.axvline(zcosmo_mean[n], color=\"green\", label=r\"$z_{\\rm cal}$\")\n",
    "# plt.fill_betweenx([0, plt.ylim()[1]], zcosmo_mean[n] - zcosmo_std[n],\n",
    "                #   zcosmo_mean[n] + zcosmo_std[n], color=\"green\", alpha=0.2)\n",
    "plt.axvline(loader.cat[\"zobs\"][n], color=\"red\", label=r\"$z_{\\rm CMB}$\")\n",
    "\n",
    "plt.xlabel(r\"$z$\")\n",
    "plt.ylabel(r\"$p(z)$\")\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "# plt.savefig(\"../plots/zcosmo_posterior.png\", dpi=450)\n",
    "plt.show()"
   ]
  },
  {
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}