{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is part of the SelfiSys project and is licensed under the GNU General Public License v3.0 or later (GPL-3.0-or-later). See the accompanying LICENSE file or visit [GNU GPL](https://www.gnu.org/licenses/gpl-3.0.en.html) for details.\n",
    "\n",
    "Tristan Hoellinger<br/>\n",
    "Institut d'Astrophysique de Paris</br>\n",
    "tristan.hoellinger@iap.fr\n",
    "\n",
    "# Investigate the binning and interpolation strategy\n",
    "\n",
    "Validate the interpolation strategy for representing power spectra."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the environment and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from selfisys.global_parameters import *\n",
    "\n",
    "wd = Path(OUTPUT_PATH) / \"expl_notebooks/interpolation/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from pathlib import Path\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from selfisys.setup_model import *\n",
    "from selfisys.utils.tools import get_k_max\n",
    "from selfisys.utils.plot_params import *\n",
    "from selfisys.utils.plot_examples import *\n",
    "\n",
    "# Configure plotting aesthetics for consistent visualisation\n",
    "setup_plotting()\n",
    "\n",
    "modeldir = Path(wd) / \"model\"\n",
    "modeldir.mkdir(parents=True, exist_ok=True)\n",
    "(Path(wd) / \"Figures\").mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 512\n",
    "L = 3600\n",
    "Pinit = 50\n",
    "S = 64"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute a power spectrum and plot it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k_max = 0.774\n"
     ]
    }
   ],
   "source": [
    "k_max = get_k_max(L, size)  # k_max in h/Mpc\n",
    "print(f\"k_max = {k_max}\")\n",
    "# Cosmo at the expansion point:\n",
    "params_planck = params_planck_kmax_missing.copy()\n",
    "params_planck[\"k_max\"] = k_max\n",
    "# Cosmo for BBKS spectrum with fiducial cosmology (for normalisation):\n",
    "params_BBKS = params_BBKS_kmax_missing.copy()\n",
    "params_BBKS[\"k_max\"] = k_max\n",
    "# Groundtruth cosmology:\n",
    "params_cosmo_obs = params_cosmo_obs_kmax_missing.copy()\n",
    "params_cosmo_obs[\"k_max\"] = k_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20:59:12|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.setup_model)\u001b[00m Computing Fourier grid...\n",
      "[20:59:12|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid...\n",
      "[20:59:19|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid done.\n",
      "[20:59:19|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/G_sim.h5'...\n",
      "[20:59:19|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=3600, L1=3600, L2=3600\u001b[00m\n",
      "[20:59:19|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=512, N1=512, N2=512, N2_HC=257, N_HC=67371008, NUM_MODES=140945\u001b[00m\n",
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/G_sim.h5' done.\n",
      "[20:59:20|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.setup_model)\u001b[00m Computing Fourier grid done.\n",
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid...\n",
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid done.\n",
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/G_ss.h5'...\n",
      "[20:59:20|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=3600, L1=3600, L2=3600\u001b[00m\n",
      "[20:59:20|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=512, N1=512, N2=512, N2_HC=257, N_HC=67371008, NUM_MODES=37\u001b[00m\n",
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/G_ss.h5' done.\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum...\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum done.\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum...\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum done.\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum...\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum done.\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum...\n",
      "[20:59:20\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum done.\n",
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Write power spectrum in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/P_ss_obj.h5'...\n",
      "[20:59:20|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=3600, L1=3600, L2=3600\u001b[00m\n",
      "[20:59:20|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=512, N1=512, N2=512, N2_HC=257, N_HC=67371008, NUM_MODES=37\u001b[00m\n",
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Write power spectrum in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/P_ss_obj.h5' done.\n"
     ]
    }
   ],
   "source": [
    "params = setup_model(\n",
    "    workdir=modeldir,\n",
    "    params_planck=params_planck,\n",
    "    params_P0=params_BBKS,\n",
    "    size=size,\n",
    "    L=L,\n",
    "    S=S,\n",
    "    Pinit=Pinit,\n",
    "    force=True,\n",
    ")\n",
    "(\n",
    "    size,\n",
    "    L,\n",
    "    P,\n",
    "    S,\n",
    "    G_sim_path,\n",
    "    G_ss_path,\n",
    "    Pbins_bnd,\n",
    "    Pbins,\n",
    "    k_s,\n",
    "    P_ss_obj_path,\n",
    "    P_0,\n",
    "    planck_Pk,\n",
    ") = params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20:59:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid...\n",
      "[20:59:28|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid done.\n",
      "[20:59:28|\u001b[38;5;113mSTATUS    \u001b[00m]|Read Fourier grid in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/G_sim.h5'...\n",
      "[20:59:28|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=3600.0, L1=3600.0, L2=3600.0\u001b[00m\n",
      "[20:59:28|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=512, N1=512, N2=512, N2_HC=257, N_HC=67371008, NUM_MODES=140945\u001b[00m\n",
      "[20:59:28|\u001b[38;5;113mSTATUS    \u001b[00m]|Read Fourier grid in data file '/Users/hoellinger/SELFI/expl_notebooks/interpolation/model/G_sim.h5' done.\n"
     ]
    }
   ],
   "source": [
    "from pysbmy.power import PowerSpectrum, FourierGrid\n",
    "\n",
    "true_P = PowerSpectrum(L, L, L, size, size, size, params_planck)\n",
    "G_sim = FourierGrid.read(G_sim_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20:59:29|\u001b[1;36mINFO      \u001b[00m]|Figure saved to: /Users/hoellinger/SELFI/expl_notebooks/interpolation/Figures/summary.pdf\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 425936 characters]",
      "text/plain": [
       "<Figure size 4500x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_power_spectrum(\n",
    "    G_sim,\n",
    "    true_P,\n",
    "    k_s,\n",
    "    planck_Pk,\n",
    "    Pbins,\n",
    "    Pbins_bnd,\n",
    "    size,\n",
    "    L,\n",
    "    wd,\n",
    "    \"Power spectrum with Planck 2018 cosmological parameters\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20:59:30|\u001b[1;36mINFO      \u001b[00m]|Figure saved to: /Users/hoellinger/SELFI/expl_notebooks/interpolation/Figures/summary_relerr.pdf\n",
      "[20:59:30|\u001b[1;36mINFO      \u001b[00m]|Max relative error over all support points: 0.663%\n",
      "[20:59:30|\u001b[1;36mINFO      \u001b[00m]|Max relative error up to 1D Nyquist frequency: 0.134%\n",
      "[20:59:30|\u001b[1;36mINFO      \u001b[00m]|Max relative error up to k = 2e-1: 0.069%\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 322092 characters]",
      "text/plain": [
       "<Figure size 4500x1500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "relative_error_analysis(G_sim, true_P, k_s, planck_Pk, Pbins, Pbins_bnd, size, L, wd)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "simbel",
   "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.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
