{
 "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",
    "# Explore the forward model of galaxy surveys\n",
    "\n",
    "Implement a hidden-box forward model of a spectroscopic galaxy survey."
   ]
  },
  {
   "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 *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "wd_ext = \"demo/\"  # Extension for the output directory\n",
    "name = \"ref_n01\"  # Name for this run\n",
    "# Set simulation parameters. Examples:\n",
    "#  - custom19COLA20: COLA, z_ini = 19, 20 time steps, non-linear RSDs\n",
    "#  - custom19PM20: PM, z_ini = 19, 20 time steps, non-linear RSDs\n",
    "#  - custom19COLA20lin: COLA, z_ini = 19, 20 time steps, linear RSDs\n",
    "sim_params = \"custom19COLA20\"  # Simulation parameters\n",
    "force = True  # Overwrite existing files?\n",
    "force_obs = True  # Overwrite existing observations?\n",
    "RSDs = True  # Use the redshift-space fields for the observations?\n",
    "id_obs = \"obs_rawmask_n01_wRSDs\"  # Identifier for the observations\n",
    "size = 128  # Grid size\n",
    "Np0 = 128  # Number of dark matter particles per spatial dimension.\n",
    "Npm0 = 128  # Side length of the particle-mesh grid in voxels.\n",
    "L = 3600  # Side length of the simulation box in Mpc/h.\n",
    "S = 64  # Number of support wavenumbers for the input power spectrum.\n",
    "Pinit = 50  # Maximum number of bins for the summaries.\n",
    "Nnorm = 2  # Number of simulations to compute the normalisation factor.\n",
    "# Time stepping:\n",
    "aa = [0.05, 0.55, 0.67, 0.8696]  # Initial and snapshot scale factors\n",
    "total_steps = 10  # Total number of time steps for the gravity solver\n",
    "# Path to the survey mask (optional):\n",
    "survey_mask_path = OUTPUT_PATH + \"expl_notebooks/surveymask/raw_mask_N128.npy\"\n",
    "# Radial selection parameters:\n",
    "radial_selection = \"multiple_lognormal\"\n",
    "selection_params = [\n",
    "    0.1150,\n",
    "    0.1492,\n",
    "    0.1818,\n",
    "    0.1500,\n",
    "    0.4925,\n",
    "    0.8182,\n",
    "    1,\n",
    "    1,\n",
    "    1,\n",
    "]\n",
    "lin_bias = [1.47, 1.99, 2.32]  # Linear galaxy biases\n",
    "obs_density = None  # Observed galaxy density\n",
    "noise = 0.1  # Noise level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Automatic reloading of modules\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from os.path import exists\n",
    "from pathlib import Path\n",
    "\n",
    "from pysbmy.timestepping import *\n",
    "from selfisys.setup_model import *\n",
    "from selfisys.hiddenbox import HiddenBox\n",
    "from selfisys.utils.tools import get_k_max\n",
    "from selfisys.utils.timestepping import merge_nTS\n",
    "from selfisys.utils.plot_utils import *\n",
    "from selfisys.utils.plot_examples import plot_galaxy_field_slice\n",
    "from selfisys.utils.examples_utils import clear_large_plot\n",
    "\n",
    "seedphase = BASELINE_SEEDPHASE\n",
    "seednoise = BASELINE_SEEDNOISE\n",
    "seednorm = BASELINE_SEEDNORM\n",
    "seedphase_obs = SEEDPHASE_OBS\n",
    "seednoise_obs = SEEDNOISE_OBS\n",
    "baseid_obs = BASEID_OBS  # Warning: never prepend a prefix to BASEID_OBS\n",
    "\n",
    "# Configure plotting aesthetics for consistent visualisation\n",
    "setup_plotting()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k_max = 0.194\n"
     ]
    }
   ],
   "source": [
    "k_max = get_k_max(L, size)  # k_max in h/Mpc\n",
    "print(f\"k_max = {k_max}\")\n",
    "Npop = len(lin_bias) if not isinstance(lin_bias, float) else 1\n",
    "gravity_on = sim_params[:6] != \"nograv\"\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\n",
    "\n",
    "if radial_selection == \"multiple_lognormal\":\n",
    "    selection_params = np.reshape(np.array(selection_params), (3, -1))\n",
    "\n",
    "# Create output directories:\n",
    "wd = f\"{OUTPUT_PATH}{wd_ext}{size}{int(L)}{Pinit}{Nnorm}/{name}/\"\n",
    "modeldir = wd + \"model/\"\n",
    "\n",
    "Path(modeldir).mkdir(parents=True, exist_ok=True)\n",
    "Path(wd, \"data\").mkdir(parents=True, exist_ok=True)\n",
    "Path(wd, \"Figures\").mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[21:58:51|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.setup_model)\u001b[00m Computing Fourier grid...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/G_sim.h5'...\n",
      "[21:58:51|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=3600, L1=3600, L2=3600\u001b[00m\n",
      "[21:58:51|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=128, N1=128, N2=128, N2_HC=65, N_HC=1064960, NUM_MODES=8041\u001b[00m\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/G_sim.h5' done.\n",
      "[21:58:51|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.setup_model)\u001b[00m Computing Fourier grid done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/G_ss.h5'...\n",
      "[21:58:51|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=3600, L1=3600, L2=3600\u001b[00m\n",
      "[21:58:51|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=128, N1=128, N2=128, N2_HC=65, N_HC=1064960, NUM_MODES=32\u001b[00m\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/G_ss.h5' done.\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum...\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum done.\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum...\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum done.\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum...\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum done.\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum...\n",
      "[21:58:51\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write power spectrum in data file '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/P_ss_obj.h5'...\n",
      "[21:58:51|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=3600, L1=3600, L2=3600\u001b[00m\n",
      "[21:58:51|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=128, N1=128, N2=128, N2_HC=65, N_HC=1064960, NUM_MODES=32\u001b[00m\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write power spectrum in data file '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/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=force,\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_EH,\n",
    ") = params\n",
    "\n",
    "\n",
    "def theta2P(theta):\n",
    "    return theta * P_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts1.h5'...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts1.h5' done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts1.h5'...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts1.h5' done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts2.h5'...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts2.h5' done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts2.h5'...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts2.h5' done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts3.h5'...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts3.h5' done.\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts3.h5'...\n",
      "[21:58:51|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts3.h5' done.\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts1.h5'...\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts1.h5' done.\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts2.h5'...\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts2.h5' done.\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts3.h5'...\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/ts3.h5' done.\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/merged.h5'...\n",
      "[21:58:52|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/merged.h5' done.\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 11668 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 9768 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 10216 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 16760 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "if gravity_on:\n",
    "    reset_plotting()  # Default style for Simbelmynë\n",
    "    merged_path = f\"{modeldir}merged.h5\"\n",
    "    if not exists(merged_path) or force:\n",
    "        # Determine the number of steps between consecutive scale factors\n",
    "        nsteps = [\n",
    "            round((aa[i + 1] - aa[i]) / (aa[-1] - aa[0]) * total_steps) for i in range(len(aa) - 1)\n",
    "        ]\n",
    "        if sum(nsteps) != total_steps:\n",
    "            nsteps[nsteps.index(max(nsteps))] += total_steps - sum(nsteps)\n",
    "        indices_steps_cumul = list(np.cumsum(nsteps) - 1)\n",
    "        np.save(f\"{modeldir}indices_steps_cumul.npy\", indices_steps_cumul)\n",
    "\n",
    "        # Generate and save individual time-stepping files for Simbelmynë\n",
    "        TSs = []\n",
    "        TS_paths = []\n",
    "        for i, (ai, af) in enumerate(zip(aa[:-1], aa[1:])):\n",
    "            snapshots = np.full((nsteps[i]), False)\n",
    "            snapshots[-1] = True  # Mark the last step as a snapshot\n",
    "            TS = StandardTimeStepping(ai, af, snapshots, 0)\n",
    "            TS_path = f\"{modeldir}ts{i + 1}.h5\"\n",
    "            TS.write(TS_path)\n",
    "            TSs.append(read_timestepping(TS_path))\n",
    "            TS_paths.append(TS_path)\n",
    "\n",
    "        # Plot individual time-stepping\n",
    "        for i, TS in enumerate(TSs):\n",
    "            TS.plot(path=modeldir + \"TS\" + str(i) + \".png\")\n",
    "\n",
    "        # Merge all time-stepping files into a single file\n",
    "        merge_nTS(TS_paths, merged_path)\n",
    "    else:\n",
    "        indices_steps_cumul = np.load(f\"{modeldir}indices_steps_cumul.npy\")\n",
    "\n",
    "    TS_merged = read_timestepping(merged_path)\n",
    "    TS_merged.plot(path=modeldir + \"TS_merged.png\")  # Plot the merged time-stepping\n",
    "    setup_plotting()  # Reset plotting style for this project\n",
    "else:\n",
    "    merged_path = None\n",
    "    indices_steps_cumul = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "if (not exists(modeldir + \"theta_gt.npy\")) or force:\n",
    "    from pysbmy.power import get_Pk\n",
    "\n",
    "    theta_gt = get_Pk(k_s, params_cosmo_obs)\n",
    "    np.save(modeldir + \"theta_gt\", theta_gt)\n",
    "else:\n",
    "    theta_gt = np.load(modeldir + \"theta_gt.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Instantiating the HiddenBox...\n",
      "[21:58:53|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.utils.plot_utils)\u001b[00m Plotting selection functions...\n",
      "[21:58:53|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(selfisys.utils.plot_utils)\u001b[00m Figure saved to /Users/hoellinger/SELFI/demo/1283600502/ref_n01/model/obs_rawmask_n01_wRSDs_select_fct.png\n",
      "[21:58:54|\u001b[38;5;113mSTATUS    \u001b[00m]|Defining normalisation constants...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77f29ca12bce4255b0d09e362f5f5312",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(\"Instantiating the HiddenBox...\")\n",
    "HB_selfi = HiddenBox(\n",
    "    k_s=k_s,\n",
    "    P_ss_path=P_ss_obj_path,\n",
    "    Pbins_bnd=Pbins_bnd,\n",
    "    theta2P=theta2P,\n",
    "    P=P * Npop,\n",
    "    size=size,\n",
    "    L=L,\n",
    "    G_sim_path=G_sim_path,\n",
    "    G_ss_path=G_ss_path,\n",
    "    Np0=Np0,\n",
    "    Npm0=Npm0,\n",
    "    fsimdir=wd[:-1],\n",
    "    noise_std=noise,\n",
    "    radial_selection=radial_selection,\n",
    "    selection_params=selection_params,\n",
    "    observed_density=obs_density,\n",
    "    linear_bias=lin_bias,\n",
    "    norm_csts=None,\n",
    "    survey_mask_path=survey_mask_path,\n",
    "    local_mask_prefix=id_obs,\n",
    "    sim_params=sim_params,\n",
    "    TimeStepDistribution=merged_path,\n",
    "    TimeSteps=indices_steps_cumul,  \n",
    "    eff_redshifts=1 / aa[-1] - 1,\n",
    "    seedphase=seedphase,  # Default, global value for the class instance\n",
    "    seednoise=seednoise,  # Default, global value for the class instance\n",
    "    fixnoise=False,\n",
    "    seednorm=seednorm,\n",
    "    reset=True,  # Reset the survey mask at instantiation\n",
    "    save_frequency=5,\n",
    "    verbosity=0,\n",
    ")\n",
    "\n",
    "if not exists(modeldir + \"norm_csts.npy\") or force:\n",
    "    from selfisys.normalise_hb import define_normalisation\n",
    "\n",
    "    HB_selfi.switch_recompute_pool()\n",
    "    norm_csts = define_normalisation(\n",
    "        HB_selfi,\n",
    "        Pbins,\n",
    "        params_planck,\n",
    "        Nnorm,\n",
    "        min_k_norma=MIN_K_NORMALISATION,\n",
    "        npar=1,\n",
    "        force=force,\n",
    "    )\n",
    "    HB_selfi.switch_recompute_pool()\n",
    "    np.save(modeldir + \"norm_csts.npy\", norm_csts)\n",
    "else:\n",
    "    print(\"> Normalisation constants already exist.\")\n",
    "    norm_csts = np.load(modeldir + \"norm_csts.npy\")\n",
    "HB_selfi.update(norm_csts=norm_csts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if not exists(modeldir + f\"phi_{id_obs}.npy\") or force or force_obs:\n",
    "    d_obs = -1\n",
    "    HB_selfi.switch_recompute_pool()\n",
    "    phi_obs, g_obs = HB_selfi.make_data(\n",
    "        cosmo=params_cosmo_obs,\n",
    "        id=baseid_obs,\n",
    "        seedphase=seedphase_obs,\n",
    "        seednoise=seednoise_obs,\n",
    "        d=d_obs,\n",
    "        force_powerspectrum=force_obs,\n",
    "        force_parfiles=force_obs,\n",
    "        force_sim=force_obs,\n",
    "        force_cosmo=force_obs,\n",
    "        return_g=True,\n",
    "        RSDs=RSDs,\n",
    "        prefix_mocks=id_obs,\n",
    "    )\n",
    "    HB_selfi.switch_recompute_pool()\n",
    "    g_obs = np.array(g_obs)\n",
    "    np.save(modeldir + f\"phi_{id_obs}.npy\", phi_obs)\n",
    "    np.save(modeldir + f\"g_{id_obs}.npy\", g_obs)\n",
    "else:\n",
    "    print(\"Observed data already exist.\")\n",
    "    phi_obs = np.load(modeldir + f\"phi_{id_obs}.npy\")\n",
    "    g_obs = np.load(modeldir + f\"g_{id_obs}.npy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 173008 characters]",
      "text/plain": [
       "<Figure size 1800x1800 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Extract and display one slice through the observed field for population i\n",
    "i = 1\n",
    "g = g_obs[i][:, :, int(size // 12)]\n",
    "\n",
    "# Save slices through the observed field for all populations\n",
    "plot_galaxy_field_slice(g, size, L, wd, id_obs)\n",
    "for i in range(3):\n",
    "    plot_galaxy_field_slice(\n",
    "        g_obs[i][:, :, int(size // 12)], size, L, wd, id_obs + f\"pop{i+1}\", display=False\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualise the 3D grid mask\n",
    "pop_idx = 2\n",
    "if size <= 128:\n",
    "    fig = plotly_3d(g_obs[pop_idx], size, colormap=\"Blues\", limits=\"truncate\")\n",
    "else:\n",
    "    # Downsample the grid for visualisation\n",
    "    downsample_factor = size // 128\n",
    "    downsampled_field = g_obs[\n",
    "        pop_idx, ::downsample_factor, ::downsample_factor, ::downsample_factor\n",
    "    ]\n",
    "    fig = plotly_3d(downsampled_field, L=L, colormap=\"Blues\", limits=\"truncate\")\n",
    "\n",
    "# fig.show()\n",
    "# # clear_large_plot(fig)  # Uncomment to clear the Plotly figure to avoid memory issues"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 525220 characters]",
      "text/plain": [
       "<Figure size 3600x1500 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_observations(k_s, theta_gt, planck_Pk_EH, P_0, Pbins, phi_obs, Npop,\n",
    "                  path=wd + f\"Figures/observations_{id_obs}.png\", force_plot=True)"
   ]
  }
 ],
 "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
}
