{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Quijote $C_{\\ell}$ of the LOS pecululiar velocity covariance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import healpy\n",
    "from h5py import File\n",
    "\n",
    "import csiborgtools\n",
    "import healpy as hp\n",
    "from scipy.special import spherical_jn\n",
    "from scipy.interpolate import interp1d\n",
    "from scipy.integrate import quad\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "k, Pk = np.load(\"/mnt/users/rstiskalek/BayesianBulkFlows/data/pk_fiducial.npy\").T\n",
    "loglog_Pk = interp1d(np.log(k), np.log(Pk))\n",
    "\n",
    "def Pk(k):\n",
    "    return np.exp(loglog_Pk(np.log(k)))\n",
    "\n",
    "\n",
    "def djn(n, x):\n",
    "    return (n / x) * spherical_jn(n, x) - spherical_jn(n + 1, x)\n",
    "\n",
    "\n",
    "def C_ell_theory_equal_radius(ell, r):\n",
    "    dDdtau = 35.5\n",
    "    \n",
    "    norm = dDdtau**2 \n",
    "    integrand = lambda k: norm * Pk(k) * djn(ell, k * r)**2\n",
    "    return quad(integrand, 0.0001, 10, limit=1000, limlst=5)[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num. simulations:   50\n",
      "Num. observers:     8\n",
      "Num. radii:         5\n",
      "Radii:              [100.  112.5 125.  137.5 150. ]\n",
      "Dividing the dipole by a factor of 9:\n"
     ]
    }
   ],
   "source": [
    "num_simulations = None\n",
    "load_skymaps = False\n",
    "\n",
    "with File(\"/mnt/extraspace/rstiskalek/BBF/Quijote_Cell/C_ell_fiducial.h5\", 'r') as f:\n",
    "    num_observers = f.attrs[\"num_observers\"]\n",
    "    num_radii = f.attrs[\"num_radii\"]\n",
    "    npix_per_map = f.attrs[\"npix_per_map\"]\n",
    "    if num_simulations is None:\n",
    "        num_simulations = f.attrs[\"num_simulations\"]\n",
    "\n",
    "    ell = f[\"ell\"][:]\n",
    "    radii = f[\"radii\"][:]\n",
    "\n",
    "    Cell = []\n",
    "    if load_skymaps:\n",
    "        skymap = []\n",
    "\n",
    "    for n in range(num_simulations):\n",
    "        grp_simulation = f[f\"nsim_{n}\"]\n",
    "\n",
    "        for i in range(num_observers):\n",
    "            grp_observer = grp_simulation[f\"observer_{i}\"]\n",
    "\n",
    "            Cell_ = np.full((num_radii, len(ell)), np.nan)\n",
    "            if load_skymaps:\n",
    "                skymap_ = np.full((num_radii, npix_per_map), np.nan)\n",
    "            else:\n",
    "                skymap_ = None\n",
    "\n",
    "            for j in range(num_radii):\n",
    "                Cell_[j] = grp_observer[f\"C_ell_{j}\"][...]\n",
    "                if load_skymaps:\n",
    "                    skymap_[j] = grp_observer[f\"skymap_{j}\"][...]\n",
    "\n",
    "\n",
    "            Cell.append(Cell_)\n",
    "            if load_skymaps:\n",
    "                skymap.append(skymap_)\n",
    "\n",
    "\n",
    "Cell = np.stack(Cell)\n",
    "if load_skymaps:\n",
    "    skymap = np.stack(skymap)\n",
    "\n",
    "\n",
    "print(f\"Num. simulations:   {num_simulations}\")\n",
    "print(f\"Num. observers:     {num_observers}\")\n",
    "print(f\"Num. radii:         {num_radii}\")\n",
    "print(f\"Radii:              {radii}\")\n",
    "\n",
    "\n",
    "\n",
    "print(\"Dividing the dipole by a factor of 9:\")\n",
    "Cell[..., 1] /= 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# healpy.mollview(skymap[2, 4], title=None, cmap=\"coolwarm\", unit=r\"$V_{\\rm rad} ~ [{\\rm km/s}]$\", )\n",
    "# plt.savefig(\"../../plots/skymap.png\", dpi=450)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1113107/3790696556.py:17: IntegrationWarning: The occurrence of roundoff error is detected, which prevents \n",
      "  the requested tolerance from being achieved.  The error may be \n",
      "  underestimated.\n",
      "  return quad(integrand, 0.0001, 10, limit=1000, limlst=5)[0]\n"
     ]
    }
   ],
   "source": [
    "k = -1\n",
    "\n",
    "\n",
    "ytheory = [C_ell_theory_equal_radius(i, radii[k]) for i in ell]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 45072 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure()\n",
    "\n",
    "plt.title(r\"$C_{\\ell}$ but $C_{1}$ divided by 9\")\n",
    "ylow, ymed, yhigh = np.percentile(Cell[:, k, :], [16, 50, 84], axis=0)\n",
    "plt.plot(ell, ymed, label=\"Quijote\")\n",
    "plt.fill_between(ell, ylow, yhigh, alpha=0.3)\n",
    "plt.plot(ell, ytheory, label=\"Linear theory\", color=\"black\", linestyle=\"--\")\n",
    "\n",
    "plt.legend()\n",
    "plt.xlabel(r\"$\\ell$\")\n",
    "plt.ylabel(r\"$C_{\\rm ell}$\")\n",
    "plt.xlim(ell.min(), ell.max())\n",
    "plt.yscale(\"log\")\n",
    "\n",
    "plt.savefig(\"../../plots/C_ell.png\", dpi=450)\n",
    "\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": [
    "\n",
    "\n",
    "fig, axs = plt.subplots(len(radii), 1, figsize=(5, 4 * len(radii)), sharex=True)\n",
    "\n",
    "for k in range(len(radii)):\n",
    "\n",
    "    ylow, ymed, yhigh = np.percentile(C_ell[:, k], [16, 50, 84], axis=0)\n",
    "    axs[k].plot(ell, ymed, label=rf\"$R = {radii[k]} ~ [\\mathrm{{Mpc}} / h]$\")\n",
    "    axs[k].fill_between(ell, ylow, yhigh, alpha=0.1)\n",
    "\n",
    "    Cell_theory = [C_ell_theory_equal_radius(ell_, radii[k]) for ell_ in ell]\n",
    "    axs[k].plot(ell, Cell_theory, color=\"black\", label=\"Theory\")\n",
    "\n",
    "    axs[k].set_yscale(\"log\")\n",
    "    axs[k].legend()\n",
    "\n",
    "    axs[k].set_ylabel(r\"$C_{\\ell}$\")\n",
    "    axs[k].set_xlabel(r\"$\\ell$\")\n",
    "\n",
    "axs[-1].set_xlim(0, ell.max())\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"../../plots/C_ell_fiducial.png\", dpi=450)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hp.mollview(skymap[4], unit=r\"$V_{\\rm rad} ~ [\\mathrm{km} / \\mathrm{s}]$\",  cmap=\"coolwarm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 4\n",
    "\n",
    "Cell_theory = [C_ell_theory_equal_radius(ell_, radii[k]) for ell_ in ell]\n",
    "\n",
    "plt.figure()\n",
    "plt.title(f\"R = {radii[k]}\")\n",
    "\n",
    "\n",
    "for n in range(50):\n",
    "    plt.plot(ell, C_ell[n, k], alpha=0.1, color=\"black\")\n",
    "\n",
    "\n",
    "plt.plot(ell, C_ell[0, k], color=\"red\", label=\"Simulated\")\n",
    "plt.plot(ell, Cell_theory, color=\"blue\", label=\"Theory\")\n",
    "\n",
    "plt.yscale(\"log\")\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
}
