{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Density & velocity fields alond a LOS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The autoreload extension is already loaded. To reload it, use:\n",
      "  %reload_ext autoreload\n"
     ]
    }
   ],
   "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",
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from numpyro.infer import MCMC, NUTS, init_to_median\n",
    "\n",
    "import csiborgtools\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": [
    "## LOS density & radial velocity plots "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fpath = \"/mnt/extraspace/rstiskalek/catalogs/PV_compilation.hdf5\"\n",
    "\n",
    "loader_carrick = csiborgtools.flow.DataLoader(\"Carrick2015\", [0], \"LOSS\", fpath, paths, ksmooth=0, )\n",
    "# loaders_csiborg2X = [csiborgtools.flow.DataLoader(\"csiborg2X\", i, \"LOSS\", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]\n",
    "# loaders_csiborg2 = [csiborgtools.flow.DataLoader(\"csiborg2_main\", i, \"LOSS\", fpath, paths, ksmooth=1, verbose=False) for i in range(20)]\n",
    "\n",
    "loader_CF4 = csiborgtools.flow.DataLoader(\"CF4gp\", [0], \"LOSS\", fpath, paths, ksmooth=0, )\n",
    "loader_lilow = csiborgtools.flow.DataLoader(\"Lilow2024\", [0], \"LOSS\", fpath, paths, ksmooth=0, )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ks = [i for i in range(50)]\n",
    "ks = [0,]\n",
    "\n",
    "for k in ks:\n",
    "    fig, axs = plt.subplots(2, 1, figsize=(7, 7), sharex=True)\n",
    "    fig.subplots_adjust(wspace=0)\n",
    "    cols = plt.rcParams['axes.prop_cycle'].by_key()['color']\n",
    "\n",
    "    # # CSiBORG2\n",
    "    # x = loaders_csiborg2X[0].rdist\n",
    "    # y = np.asarray([loaders_csiborg2[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[0].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)\n",
    "    # axs[0].plot(x, ymed, color=cols[0], label=\"CSiBORG2\")\n",
    "\n",
    "    # y = np.asarray([loaders_csiborg2[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[1].fill_between(x, ylow, yhigh, color=cols[0], alpha=0.25)\n",
    "    # axs[1].plot(x, ymed, color=cols[0], label=\"CSiBORG2\")\n",
    "\n",
    "    # # CSiBORG2X\n",
    "    # x = loaders_csiborg2X[0].rdist\n",
    "    # y = np.asarray([loaders_csiborg2X[i].los_density[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[0].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)\n",
    "    # axs[0].plot(x, ymed, color=cols[1], label=\"CSiBORG2X\")\n",
    "\n",
    "    # y = np.asarray([loaders_csiborg2X[i].los_radial_velocity[k, :] for i in range(len(loaders_csiborg2X))])\n",
    "    # ylow, ymed, yhigh = np.percentile(y, [16, 50, 84], axis=0)\n",
    "    # axs[1].fill_between(x, ylow, yhigh, color=cols[1], alpha=0.25)\n",
    "    # axs[1].plot(x, ymed, color=cols[1], label=\"CSiBORG2X\")\n",
    "\n",
    "    # Plot Carrick+2015\n",
    "    axs[0].plot(loader_carrick.rdist, loader_carrick.los_density[0, k, :], color=\"red\", label=\"Carrick+2015\")\n",
    "    axs[1].plot(loader_carrick.rdist, loader_carrick.los_radial_velocity[0, k, :] * 0.43, color=\"red\")\n",
    "\n",
    "    # Plot CF4\n",
    "    c = cols[4]\n",
    "    axs[0].plot(loader_CF4.rdist, loader_CF4.los_density[0, k, :], color=c, label=\"CF4\")\n",
    "    axs[1].plot(loader_CF4.rdist, loader_CF4.los_radial_velocity[0, k, :], color=c)\n",
    "\n",
    "    # Plot Lilow2024\n",
    "    c = cols[5]\n",
    "    axs[0].plot(loader_lilow.rdist, loader_lilow.los_density[0, k, :], color=c, label=\"Lilow+2024\")\n",
    "    axs[1].plot(loader_lilow.rdist, loader_lilow.los_radial_velocity[0, k, :], color=c)\n",
    "\n",
    "\n",
    "    axs[1].set_xlabel(r\"$r ~ [\\mathrm{Mpc} / h]$\")\n",
    "    axs[0].set_ylabel(r\"$\\rho_{\\rm LOS} / \\langle \\rho_{\\rm matter} \\rangle$\")\n",
    "    axs[1].set_ylabel(r\"$v_{\\rm LOS} ~ [\\mathrm{km/s}]$\")\n",
    "    axs[0].set_yscale(\"log\")\n",
    "\n",
    "    axs[0].legend(loc=\"upper right\")\n",
    "    axs[0].set_xlim(0, 200)\n",
    "\n",
    "    fig.tight_layout(w_pad=0, h_pad=0)\n",
    "    fig.savefig(f\"../../plots/LOSS_los_{k}.png\", dpi=500, bbox_inches=\"tight\")\n",
    "\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test running a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-06-29 19:40:25.229961:   reading the catalogue,\n",
      "2024-06-29 19:40:25.243502:   reading the interpolated field,\n",
      "2024-06-29 19:40:25.261423:   calculating the radial velocity.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/users/rstiskalek/csiborgtools/csiborgtools/flow/flow_model.py:91: 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.hdf5\"\n",
    "\n",
    "simname = \"Carrick2015\"\n",
    "catalogue = \"LOSS\"\n",
    "loader = csiborgtools.flow.DataLoader(simname, [0, 0], catalogue, fpath_data, paths, ksmooth=0, )\n",
    "\n",
    "SN_hyperparams = {\"e_mu_mean\": 0.1, \"e_mu_std\": 0.05,\n",
    "                  \"mag_cal_mean\": -18.25, \"mag_cal_std\": 0.5,\n",
    "                  \"alpha_cal_mean\": 0.148, \"alpha_cal_std\": 0.05,\n",
    "                  \"beta_cal_mean\": 3.112, \"beta_cal_std\": 1.0,\n",
    "                  }\n",
    "calibration_hyperparams = {\"Vext_std\": 250,\n",
    "                           \"alpha_mean\": 1.0, \"alpha_std\": 0.5,\n",
    "                           \"beta_mean\": 1.0, \"beta_std\": 0.5,\n",
    "                           \"sigma_v_mean\": 150., \"sigma_v_std\": 100.,\n",
    "                           \"sample_alpha\": True, \"sample_beta\": True,\n",
    "                           }\n",
    "get_model_kwargs = {\"zcmb_max\": 0.05}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Running HMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Selected 50/50 galaxies.\n"
     ]
    }
   ],
   "source": [
    "model = csiborgtools.flow.get_model(loader, **get_model_kwargs)\n",
    "model_kwargs = {\"distmod_hyperparams\": SN_hyperparams, \"calibration_hyperparams\": calibration_hyperparams,}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [02:10<00:00,  7.68it/s, 7 steps of size 4.49e-01. acc. prob=0.90]  \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "                 mean       std    median      5.0%     95.0%     n_eff     r_hat\n",
      "    Vext[0]     -3.71     69.92     -3.04   -123.73    103.87    469.72      1.00\n",
      "    Vext[1]    -27.47     95.52    -30.48   -151.20    172.63    308.02      1.00\n",
      "    Vext[2]    -59.27    131.26    -57.79   -273.64    137.55    456.29      1.00\n",
      "      alpha      1.09      0.38      1.10      0.50      1.69    400.05      1.00\n",
      "  alpha_cal      0.13      0.03      0.13      0.09      0.17    558.81      1.00\n",
      "       beta      0.43      0.11      0.44      0.27      0.61    341.86      1.00\n",
      "   beta_cal      3.54      0.18      3.54      3.23      3.81    606.77      1.00\n",
      "       e_mu      0.08      0.03      0.08      0.04      0.12    330.71      1.00\n",
      "    mag_cal    -18.19      0.04    -18.19    -18.25    -18.13    389.94      1.00\n",
      "    sigma_v    176.93     52.05    169.93    102.74    267.56    315.30      1.00\n",
      "\n",
      "Number of divergences: 0\n"
     ]
    }
   ],
   "source": [
    "kernel = NUTS(model, init_strategy=init_to_median(num_samples=100))\n",
    "mcmc = MCMC(kernel, num_warmup=500, num_samples=500)\n",
    "\n",
    "rng_key = jax.random.PRNGKey(5)\n",
    "mcmc.run(rng_key, extra_fields=(\"potential_energy\",), **model_kwargs)\n",
    "mcmc.print_summary()\n",
    "samples = mcmc.get_samples()"
   ]
  },
  {
   "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
}