{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tristan Hoellinger<br/>\n",
    "Institut d'Astrophysique de Paris</br>\n",
    "tristan.hoellinger@iap.fr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Non-regression tests towards implementing P3M gravity"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the environment and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pyright: reportWildcardImportFromLibrary=false\n",
    "from wip3m import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "workdir = ROOT_PATH\n",
    "output_path = OUTPUT_PATH\n",
    "\n",
    "# L = 64  # Box size in Mpc/h\n",
    "# L = 32  # Box size in Mpc/h\n",
    "L = 16  # Box size in Mpc/h\n",
    "# L = 8  # Box size in Mpc/h\n",
    "N = 32  # Density grid size\n",
    "Np = 32  # Number of dark matter particles per spatial dimension\n",
    "Npm = 64  # PM grid size\n",
    "n_Tiles = 8\n",
    "\n",
    "go_beyond_Nyquist_ss = True  # for the summary statistics\n",
    "\n",
    "force = True\n",
    "force_hard = True\n",
    "\n",
    "run_id = \"run13\"\n",
    "\n",
    "# Good set of parameters for the force diagnostic\n",
    "# nPairsForceDiagnostic_pm = nPairsForceDiagnostic_spm = nPairsForceDiagnostic_p3m = 5\n",
    "# nBinsForceDiagnostic = 30\n",
    "# maxTrialsForceDiagnostic = int(2e9)\n",
    "\n",
    "# Faster force diagnostic\n",
    "nPairsForceDiagnostic_pm = nPairsForceDiagnostic_spm = nPairsForceDiagnostic_p3m = 3\n",
    "nBinsForceDiagnostic = 20\n",
    "maxTrialsForceDiagnostic = int(1e8)\n",
    "\n",
    "# Simulation parameters\n",
    "# nsteps_cola = 20\n",
    "# nsteps_pm = 200\n",
    "# nsteps_spm = 200\n",
    "# nsteps_p3m = 200\n",
    "nsteps_cola = 20\n",
    "nsteps_pm = 50\n",
    "nsteps_spm = 50\n",
    "nsteps_p3m = 50\n",
    "# nsteps_cola = 10\n",
    "# nsteps_pm = 20\n",
    "# nsteps_spm = 20\n",
    "# nsteps_p3m = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In principle nothing needs to be changed below this cell."
   ]
  },
  {
   "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 isfile\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "\n",
    "from pysbmy.power import PowerSpectrum\n",
    "from pysbmy.fft import FourierGrid, read_FourierGrid\n",
    "from pysbmy.field import read_field\n",
    "from pysbmy.correlations import get_autocorrelation\n",
    "\n",
    "from wip3m.tools import get_k_max, generate_sim_params, generate_white_noise_Field\n",
    "from wip3m.params import params_planck_kmax_missing, cosmo_small_to_full_dict, z2a, BASELINE_SEEDPHASE\n",
    "from wip3m.plot_utils import *  # type: ignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k_max = 10.883000000000001\n"
     ]
    }
   ],
   "source": [
    "corner = -L / 2.0\n",
    "RedshiftLPT = 19.0\n",
    "RedshiftFCs = 0.0\n",
    "ai = z2a(RedshiftLPT)\n",
    "af = z2a(RedshiftFCs)\n",
    "k_max = get_k_max(L, N)  # k_max in h/Mpc\n",
    "print(f\"k_max = {k_max}\")\n",
    "cosmo = params_planck_kmax_missing.copy()\n",
    "cosmo[\"k_max\"] = k_max\n",
    "\n",
    "TimeStepDistribution_pm = 0\n",
    "TimeStepDistribution_cola = 0\n",
    "TimeStepDistribution_spm = 0\n",
    "TimeStepDistribution_p3m = 0\n",
    "\n",
    "wd = workdir + run_id + \"/\"\n",
    "simdir = output_path + run_id + \"/\"\n",
    "logdir = simdir + \"logs/\"\n",
    "if force_hard:\n",
    "    import shutil\n",
    "    if Path(simdir).exists():\n",
    "        shutil.rmtree(simdir)\n",
    "    if Path(wd).exists():\n",
    "        shutil.rmtree(wd)\n",
    "Path(wd).mkdir(parents=True, exist_ok=True)\n",
    "Path(logdir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "input_white_noise_file = simdir + \"input_white_noise.h5\"\n",
    "input_seed_phase_file = simdir + \"seed\"\n",
    "ICs_path = simdir + \"initial_density.h5\"\n",
    "simpath = simdir\n",
    "\n",
    "# Path to the input matter power spectrum (generated later)\n",
    "input_power_file = simdir + \"input_power.h5\"\n",
    "\n",
    "# Paths the the force diagnostic CSVs\n",
    "OutputForceDiagnostic_pm = simdir + \"force_diagnostic_pm.txt\"\n",
    "OutputForceDiagnostic_cola = simdir + \"force_diagnostic_cola.txt\"\n",
    "OutputForceDiagnostic_spm = simdir + \"force_diagnostic_spm.txt\"\n",
    "OutputForceDiagnostic_p3m = simdir + \"force_diagnostic_p3m.txt\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate the parameter files"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first preparatory step is to generate all the parameter files required for all the simulations.\n",
    "\n",
    "To this end we use the `generate_sim_params` function defined in `params.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "common_params = {\n",
    "    \"Np\": Np,\n",
    "    \"N\": N,\n",
    "    \"L\": L,\n",
    "    \"corner0\": corner,\n",
    "    \"corner1\": corner,\n",
    "    \"corner2\": corner,\n",
    "    \"h\": cosmo[\"h\"],\n",
    "    \"Omega_m\": cosmo[\"Omega_m\"],\n",
    "    \"Omega_b\": cosmo[\"Omega_b\"],\n",
    "    \"n_s\": cosmo[\"n_s\"],\n",
    "    \"sigma8\": cosmo[\"sigma8\"],\n",
    "}\n",
    "\n",
    "lpt_params = common_params.copy()\n",
    "lpt_params[\"method\"] = \"lpt\"\n",
    "lpt_params[\"InputPowerSpectrum\"] = input_power_file\n",
    "lpt_params[\"ICsMode\"] = 1\n",
    "# 0 : the codes generates white noise, then initial conditions\n",
    "# 1 : external white noise specified, the code multiplies by the power spectrum\n",
    "# 2 : external initial conditions specified\n",
    "lpt_params[\"InputWhiteNoise\"] = input_white_noise_file\n",
    "\n",
    "pm_params = common_params.copy()\n",
    "pm_params[\"method\"] = \"pm\"\n",
    "pm_params[\"TimeStepDistribution\"] = TimeStepDistribution_pm\n",
    "pm_params[\"ai\"] = ai\n",
    "pm_params[\"af\"] = af\n",
    "pm_params[\"RedshiftLPT\"] = RedshiftLPT\n",
    "pm_params[\"RedshiftFCs\"] = RedshiftFCs\n",
    "pm_params[\"Npm\"] = Npm\n",
    "pm_params[\"nsteps\"] = nsteps_pm\n",
    "pm_params[\"RunForceDiagnostic\"] = True\n",
    "pm_params[\"nPairsForceDiagnostic\"] = nPairsForceDiagnostic_pm\n",
    "pm_params[\"nBinsForceDiagnostic\"] = nBinsForceDiagnostic\n",
    "pm_params[\"OutputForceDiagnostic\"] = OutputForceDiagnostic_pm\n",
    "pm_params[\"maxTrialsForceDiagnostic\"] = maxTrialsForceDiagnostic\n",
    "\n",
    "cola_params = common_params.copy()\n",
    "cola_params[\"method\"] = \"cola\"\n",
    "cola_params[\"TimeStepDistribution\"] = TimeStepDistribution_cola\n",
    "cola_params[\"ai\"] = ai\n",
    "cola_params[\"af\"] = af\n",
    "cola_params[\"RedshiftLPT\"] = RedshiftLPT\n",
    "cola_params[\"RedshiftFCs\"] = RedshiftFCs\n",
    "cola_params[\"Npm\"] = Npm\n",
    "cola_params[\"nsteps\"] = nsteps_cola\n",
    "cola_params[\"RunForceDiagnostic\"] = False\n",
    "cola_params[\"RunForceDiagnostic\"] = True\n",
    "cola_params[\"nPairsForceDiagnostic\"] = nPairsForceDiagnostic_pm\n",
    "cola_params[\"nBinsForceDiagnostic\"] = nBinsForceDiagnostic\n",
    "cola_params[\"OutputForceDiagnostic\"] = OutputForceDiagnostic_cola\n",
    "cola_params[\"maxTrialsForceDiagnostic\"] = maxTrialsForceDiagnostic\n",
    "\n",
    "spm_params = common_params.copy()\n",
    "spm_params[\"method\"] = \"spm\"\n",
    "spm_params[\"EvolutionMode\"] = 5\n",
    "spm_params[\"TimeStepDistribution\"] = TimeStepDistribution_spm\n",
    "spm_params[\"ai\"] = ai\n",
    "spm_params[\"af\"] = af\n",
    "spm_params[\"RedshiftLPT\"] = RedshiftLPT\n",
    "spm_params[\"RedshiftFCs\"] = RedshiftFCs\n",
    "spm_params[\"Npm\"] = Npm\n",
    "spm_params[\"nsteps\"] = nsteps_spm\n",
    "spm_params[\"n_Tiles\"] = n_Tiles\n",
    "spm_params[\"RunForceDiagnostic\"] = True\n",
    "spm_params[\"nPairsForceDiagnostic\"] = nPairsForceDiagnostic_spm\n",
    "spm_params[\"nBinsForceDiagnostic\"] = nBinsForceDiagnostic\n",
    "spm_params[\"OutputForceDiagnostic\"] = OutputForceDiagnostic_spm\n",
    "spm_params[\"maxTrialsForceDiagnostic\"] = maxTrialsForceDiagnostic\n",
    "\n",
    "p3m_params = common_params.copy()\n",
    "p3m_params[\"method\"] = \"p3m\"\n",
    "p3m_params[\"EvolutionMode\"] = 4\n",
    "p3m_params[\"TimeStepDistribution\"] = TimeStepDistribution_p3m\n",
    "p3m_params[\"ai\"] = ai\n",
    "p3m_params[\"af\"] = af\n",
    "p3m_params[\"RedshiftLPT\"] = RedshiftLPT\n",
    "p3m_params[\"RedshiftFCs\"] = RedshiftFCs\n",
    "p3m_params[\"Npm\"] = Npm\n",
    "p3m_params[\"nsteps\"] = nsteps_p3m\n",
    "p3m_params[\"n_Tiles\"] = n_Tiles\n",
    "p3m_params[\"RunForceDiagnostic\"] = True\n",
    "p3m_params[\"nPairsForceDiagnostic\"] = nPairsForceDiagnostic_p3m\n",
    "p3m_params[\"nBinsForceDiagnostic\"] = nBinsForceDiagnostic\n",
    "p3m_params[\"OutputForceDiagnostic\"] = OutputForceDiagnostic_p3m\n",
    "p3m_params[\"maxTrialsForceDiagnostic\"] = maxTrialsForceDiagnostic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14:42:20|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Generating parameter file...\n",
      "[14:42:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/example_lpt.sbmy'...\n",
      "[14:42:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/example_lpt.sbmy' done.\n",
      "[14:42:20|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Parameter file written to /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/example_lpt.sbmy\n",
      "PM nsteps = 50:\n",
      "[14:42:20|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Time-stepping distribution file: /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_pm.h5\n",
      "[14:42:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_pm.h5'...\n",
      "[14:42:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_pm.h5' done.\n",
      "[14:42:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_pm.h5'...\n",
      "[14:42:20|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_pm.h5' done.\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Generating parameter file...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_pm.sbmy'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_pm.sbmy' done.\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Parameter file written to /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_pm.sbmy\n",
      "COLA nsteps = 20:\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Time-stepping distribution file: /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_ts_cola.h5\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_ts_cola.h5'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_ts_cola.h5' done.\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_ts_cola.h5'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_ts_cola.h5' done.\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Generating parameter file...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_example_cola.sbmy'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_example_cola.sbmy' done.\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Parameter file written to /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps20_example_cola.sbmy\n",
      "SPM nsteps = 50:\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Time-stepping distribution file: /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_spm.h5\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_spm.h5'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_spm.h5' done.\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_spm.h5'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_spm.h5' done.\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Generating parameter file...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_spm.sbmy'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_spm.sbmy' done.\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Parameter file written to /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_spm.sbmy\n",
      "P3M nsteps = 50:\n",
      "[14:42:21|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Time-stepping distribution file: /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_p3m.h5\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_p3m.h5'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Write timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_p3m.h5' done.\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_p3m.h5'...\n",
      "[14:42:21|\u001b[38;5;113mSTATUS    \u001b[00m]|Read timestepping configuration in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_ts_p3m.h5' done.\n",
      "[14:42:22|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Generating parameter file...\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_p3m.sbmy'...\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Writing parameter file in '/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_p3m.sbmy' done.\n",
      "[14:42:22|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.tools)\u001b[00m Parameter file written to /Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/simbelmyne/simbelmyne2025/WIP_P3M/run13/nsteps50_example_p3m.sbmy\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 30240 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 18356 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 30240 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 30240 characters]",
      "text/plain": [
       "<Figure size 500x100 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "reset_plotting()  # Default style for Simbelmynë\n",
    "generate_sim_params(lpt_params, ICs_path, wd, simdir, None, force)\n",
    "\n",
    "print(f\"PM nsteps = {nsteps_pm}:\")\n",
    "file_ext = f\"nsteps{nsteps_pm}\"  # \"pm\" is already in the filename\n",
    "generate_sim_params(pm_params, ICs_path, wd, simdir, file_ext, force)\n",
    "\n",
    "print(f\"COLA nsteps = {nsteps_cola}:\")\n",
    "file_ext = f\"nsteps{nsteps_cola}\"  # \"cola\" is already in the filename\n",
    "generate_sim_params(cola_params, ICs_path, wd, simdir, file_ext, force)\n",
    "\n",
    "print(f\"SPM nsteps = {nsteps_spm}:\")\n",
    "file_ext = f\"nsteps{nsteps_spm}\"  # \"spm\" is already in the filename\n",
    "generate_sim_params(spm_params, ICs_path, wd, simdir, file_ext, force)\n",
    "\n",
    "print(f\"P3M nsteps = {nsteps_p3m}:\")\n",
    "file_ext = f\"nsteps{nsteps_p3m}\"  # \"p3m\" is already in the filename\n",
    "generate_sim_params(p3m_params, ICs_path, wd, simdir, file_ext, force)\n",
    "setup_plotting()  # Reset plotting style for this project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate the initial phase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "generate_white_noise_Field(\n",
    "    L=L,\n",
    "    size=N,\n",
    "    corner=corner,\n",
    "    seedphase=BASELINE_SEEDPHASE,\n",
    "    fname_whitenoise=input_white_noise_file,\n",
    "    seedname_whitenoise=input_seed_phase_file,\n",
    "    force_phase=force,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generating the input power spectrum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second preparatory step is to compute the initial power spectrum to be used in the simulations, given the cosmological parameters and prescription specified in ``params.py``. The power spectrum is saved in `input_power_file`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid...\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid done.\n",
      "[14:42:22\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum...\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Write power spectrum in data file '/Users/hoellinger/WIP3M/run13/input_power.h5'...\n",
      "[14:42:22\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing normalization of the power spectrum done.\n",
      "[14:42:22\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum...\n",
      "[14:42:22\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Computing power spectrum done.\n",
      "[14:42:22|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=16, L1=16, L2=16\u001b[00m\n",
      "[14:42:22|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=32, N1=32, N2=32, N2_HC=17, N_HC=17408, NUM_MODES=464\u001b[00m\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Write power spectrum in data file '/Users/hoellinger/WIP3M/run13/input_power.h5' done.\n"
     ]
    }
   ],
   "source": [
    "# If cosmo[\"WhichSpectrum\"] == \"class\", then classy is required.\n",
    "if not isfile(input_power_file) or force:\n",
    "    Pk = PowerSpectrum(L, L, L, N, N, N, cosmo_small_to_full_dict(cosmo))\n",
    "    Pk.write(input_power_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid...\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Setting up Fourier grid done.\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/WIP3M/run13/input_ss_k_grid.h5'...\n",
      "[14:42:22|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=16, L1=16, L2=16\u001b[00m\n",
      "[14:42:22|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=32, N1=32, N2=32, N2_HC=17, N_HC=17408, NUM_MODES=40\u001b[00m\n",
      "[14:42:22|\u001b[38;5;113mSTATUS    \u001b[00m]|Write Fourier grid in data file '/Users/hoellinger/WIP3M/run13/input_ss_k_grid.h5' done.\n"
     ]
    }
   ],
   "source": [
    "# k grid used to compute the final overdensity power spectrum\n",
    "Pinit = 100\n",
    "trim_threshold = 100  # Merge bins until this minimum number of modes per bin is reached\n",
    "log_kmin = np.log10(2 * np.pi / (np.sqrt(3) * L))  # Minimum non-zero k in h/Mpc\n",
    "if go_beyond_Nyquist_ss:\n",
    "    k_max_ss = get_k_max(L, N)\n",
    "else:\n",
    "    k_max_ss = get_k_max(L, N) / np.sqrt(3)  # 1D Nyquist frequency\n",
    "Pbins_left_bnds = np.logspace(log_kmin, np.log10(k_max_ss), Pinit + 1, dtype=np.float32)\n",
    "Pbins_left_bnds = Pbins_left_bnds[:-1]\n",
    "input_ss_file = simdir + \"input_ss_k_grid.h5\"\n",
    "Gk = FourierGrid(\n",
    "    L,\n",
    "    L,\n",
    "    L,\n",
    "    N,\n",
    "    N,\n",
    "    N,\n",
    "    k_modes=Pbins_left_bnds,\n",
    "    kmax=k_max_ss,\n",
    "    trim_bins=True,\n",
    "    trim_threshold=trim_threshold,\n",
    ")\n",
    "Gk.write(input_ss_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to run the actual simulations using the Simbelmynë executable. Warning: the following may take some time, even in relatively low dimension, and should not be run on a login node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "if not isfile(ICs_path) or not isfile(simdir + \"lpt_density.h5\") or not isfile(simdir + \"lpt_particles.gadget3\") or force:\n",
    "    !simbelmyne {wd}example_lpt.sbmy {logdir}lpt.txt\n",
    "\n",
    "file_ext = f\"nsteps{nsteps_pm}\"  # \"pm\" is already in the filename\n",
    "if not isfile(simdir + f\"{file_ext}_final_density_pm.h5\") or force:\n",
    "    !simbelmyne {wd}{file_ext}_example_pm.sbmy {logdir}{file_ext}_pm.txt\n",
    "\n",
    "file_ext = f\"nsteps{nsteps_cola}\"  # \"cola\" is already in the filename\n",
    "if not isfile(simdir + f\"{file_ext}_final_density_cola.h5\") or force:\n",
    "    !simbelmyne {wd}{file_ext}_example_cola.sbmy {logdir}{file_ext}_cola.txt\n",
    "\n",
    "file_ext = f\"nsteps{nsteps_spm}\"  # \"spm\" is already in the filename\n",
    "if not isfile(simdir + f\"{file_ext}_final_density_spm.h5\") or force:\n",
    "    !simbelmyne {wd}{file_ext}_example_spm.sbmy {logdir}{file_ext}_spm.txt\n",
    "\n",
    "file_ext = f\"nsteps{nsteps_p3m}\"  # \"p3m\" is already in the filename\n",
    "if not isfile(simdir + f\"{file_ext}_final_density_p3m.h5\") or force:\n",
    "    !simbelmyne {wd}{file_ext}_example_p3m.sbmy {logdir}{file_ext}_p3m.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "source": [
    "The logs can be monitored in the corresponding files in the `logdir` directory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the evolved dark matter density fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/lpt_density.h5'...\n",
      "[14:43:07|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/lpt_density.h5' done.\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps20_final_density_cola.h5'...\n",
      "[14:43:07|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps20_final_density_cola.h5' done.\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_pm.h5'...\n",
      "[14:43:07|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_pm.h5' done.\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_spm.h5'...\n",
      "[14:43:07|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_spm.h5' done.\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_p3m.h5'...\n",
      "[14:43:07|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:07|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_p3m.h5' done.\n"
     ]
    }
   ],
   "source": [
    "slice_ijk = (N // 2, slice(None), slice(None))\n",
    "DELTA_LPT = read_field(simdir + \"lpt_density.h5\").data[slice_ijk]\n",
    "DELTA_COLA = read_field(simdir + f\"nsteps{nsteps_cola}_final_density_cola.h5\").data[slice_ijk]\n",
    "DELTA_PM = read_field(simdir + f\"nsteps{nsteps_pm}_final_density_pm.h5\").data[slice_ijk]\n",
    "DELTA_SPM = read_field(simdir + f\"nsteps{nsteps_spm}_final_density_spm.h5\").data[slice_ijk]\n",
    "DELTA_P3M = read_field(simdir + f\"nsteps{nsteps_p3m}_final_density_p3m.h5\").data[slice_ijk]\n",
    "diff_p3m_pm = DELTA_P3M - DELTA_PM\n",
    "diff_p3m_spm = DELTA_P3M - DELTA_SPM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "max(DELTA_PM) = 70.734375, min(DELTA_PM) = -1.0\n",
      "max(DELTA_P3M) = 69.35508728027344, min(DELTA_P3M) = -1.0\n",
      "max(diff) = 3.0339012145996094, min(diff) = -12.64963150024414\n"
     ]
    }
   ],
   "source": [
    "print(f\"max(DELTA_PM) = {np.max(DELTA_PM)}, min(DELTA_PM) = {np.min(DELTA_PM)}\")\n",
    "print(f\"max(DELTA_P3M) = {np.max(DELTA_P3M)}, min(DELTA_P3M) = {np.min(DELTA_P3M)}\")\n",
    "print(f\"max(diff) = {np.max(diff_p3m_pm)}, min(diff) = {np.min(diff_p3m_pm)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 160376 characters]",
      "text/plain": [
       "<Figure size 6300x1200 with 14 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# fields = [\"pm\", \"spm\", \"p3m\", \"diff_p3m_pm\"]  # fields to plot\n",
    "fields = [\"lpt\", \"cola\", \"pm\", \"spm\", \"p3m\", \"diff_p3m_pm\", \"diff_p3m_spm\"]  # fields to plot\n",
    "\n",
    "figname = \"_\".join(fields)\n",
    "slices_dict = {\n",
    "    \"lpt\": DELTA_LPT,\n",
    "    \"cola\": DELTA_COLA,\n",
    "    \"pm\": DELTA_PM,\n",
    "    \"spm\": DELTA_SPM,\n",
    "    \"p3m\": DELTA_P3M,\n",
    "    \"diff_p3m_pm\": diff_p3m_pm,\n",
    "    \"diff_p3m_spm\": diff_p3m_spm,\n",
    "}\n",
    "titles_dict = {\n",
    "    \"lpt\": \"LPT\",\n",
    "    \"cola\": f\"COLA $n_\\\\mathrm{{steps}}={nsteps_cola}$\",\n",
    "    \"pm\": f\"PM $n_\\\\mathrm{{steps}}={nsteps_pm}$\",\n",
    "    \"spm\": f\"sPM $n_\\\\mathrm{{steps}}={nsteps_spm}$\",\n",
    "    \"p3m\": f\"P3M $n_\\\\mathrm{{steps}}={nsteps_p3m}$\",\n",
    "    \"diff_p3m_pm\": r\"$\\delta_{\\rm P3M}-\\delta_{\\rm PM}$\",\n",
    "    \"diff_p3m_spm\": r\"$\\delta_{\\rm P3M}-\\delta_{\\rm sPM}$\",\n",
    "}\n",
    "\n",
    "npanels = len(fields)\n",
    "fig, axs = plt.subplots(1, npanels, figsize=(3 * npanels, 4), sharey=True)\n",
    "\n",
    "ims = []\n",
    "for i, key in enumerate(fields):\n",
    "    ax = axs[i]\n",
    "    data = slices_dict[key]\n",
    "    title = titles_dict[key]\n",
    "\n",
    "    if key.startswith(\"diff\"):\n",
    "        norm = TwoSlopeNorm(vmin=-np.log(1 + np.abs(np.min(data))), vcenter=0, vmax=np.log10(1 + np.abs(np.max(data))))\n",
    "        im = ax.imshow(\n",
    "            np.sign(data) * np.log(1 + np.abs(data)), cmap=\"RdBu_r\", norm=norm\n",
    "        )\n",
    "    else:\n",
    "        im = ax.imshow(np.log10(2 + data), cmap=cmap)\n",
    "\n",
    "    ims.append((im, key))\n",
    "    ax.set_title(title, fontsize=fs_titles)\n",
    "    for spine in ax.spines.values():\n",
    "        spine.set_visible(False)\n",
    "\n",
    "axs[0].set_yticks([0, N // 2, N])\n",
    "axs[0].set_yticklabels([f\"{-L/2:.0f}\", \"0\", f\"{L/2:.0f}\"], fontsize=fs)\n",
    "axs[0].set_ylabel(r\"Mpc/$h$\", size=GLOBAL_FS_SMALL)\n",
    "\n",
    "for i, ax in enumerate(axs):\n",
    "    ax.set_xticks([0, N // 2, N])\n",
    "    ax.set_xticklabels([f\"{-L/2:.0f}\", \"0\", f\"{L/2:.0f}\"], fontsize=fs)\n",
    "    ax.set_xlabel(r\"Mpc/$h$\", size=GLOBAL_FS_SMALL)\n",
    "\n",
    "for ax, (im, key) in zip(axs, ims):\n",
    "    divider = make_axes_locatable(ax)\n",
    "    cax = divider.append_axes(\"bottom\", size=\"5%\", pad=0.6)\n",
    "    cb = fig.colorbar(im, cax=cax, orientation=\"horizontal\")\n",
    "    if key.startswith(\"diff\"):\n",
    "        cb.set_label(r\"$\\textrm{sgn}\\left(\\Delta\\delta\\right)\\log_{10}(1 + |\\Delta\\delta|)$\", fontsize=fs)\n",
    "    else:\n",
    "        cb.set_label(r\"$\\log_{10}(2 + \\delta)$\", fontsize=fs)\n",
    "    cb.ax.tick_params(labelsize=fs)\n",
    "    cax.xaxis.set_ticks_position(\"bottom\")\n",
    "    cax.xaxis.set_label_position(\"bottom\")\n",
    "fig.savefig(\n",
    "    simdir + f\"{figname}.png\",\n",
    "    bbox_inches=\"tight\",\n",
    "    dpi=300,\n",
    "    transparent=True,\n",
    ")\n",
    "fig.savefig(\n",
    "    simdir + f\"{figname}.pdf\",\n",
    "    bbox_inches=\"tight\",\n",
    "    dpi=300,\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "full_field_p3m = np.log10(2+read_field(simdir + f\"nsteps{nsteps_p3m}_final_density_p3m.h5\").data)\n",
    "\n",
    "if N <= 128:\n",
    "    fig = plotly_3d(full_field_p3m, size=N, L=L, colormap=thermal_plotly, limits=\"default\")\n",
    "else:\n",
    "    # Downsample the grid for visualisation\n",
    "    downsample_factor = N // 128\n",
    "    downsampled_field = full_field_p3m[\n",
    "        ::downsample_factor, ::downsample_factor, ::downsample_factor\n",
    "    ]\n",
    "    fig = plotly_3d(downsampled_field, size=N, L=L, colormap=thermal_plotly, limits=\"default\")\n",
    "\n",
    "fig.show()\n",
    "clear_large_plot(fig)  # Uncomment to clear the Plotly figure to avoid memory issues"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute and plot the power spectra of the evolved dark matter fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read Fourier grid in data file '/Users/hoellinger/WIP3M/run13/input_ss_k_grid.h5'...\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mL0=16.0, L1=16.0, L2=16.0\u001b[00m\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mN0=32, N1=32, N2=32, N2_HC=17, N_HC=17408, NUM_MODES=40\u001b[00m\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read Fourier grid in data file '/Users/hoellinger/WIP3M/run13/input_ss_k_grid.h5' done.\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/initial_density.h5'...\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/initial_density.h5' done.\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores)...\n",
      "[14:43:09\u001b[00m|[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/lpt_density.h5'...\n",
      "\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores) done.\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/lpt_density.h5' done.\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores)...\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_pm.h5'...\n",
      "]|Getting auto-correlation in Fourier space (using 8 cores) done.\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_pm.h5' done.\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores)...\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps20_final_density_cola.h5'...\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores) done.\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps20_final_density_cola.h5' done.\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores)...\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_spm.h5'...\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores) done.\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_spm.h5' done.\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores)...\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_p3m.h5'...\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores) done.\n",
      "[14:43:09|\u001b[38;5;246mDIAGNOSTIC\u001b[00m]==|\u001b[38;5;246mranges=[np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.float64(-8.0), np.float64(8.0), np.int32(32), np.int32(32), np.int32(32)]\u001b[00m\n",
      "[14:43:09|\u001b[38;5;113mSTATUS    \u001b[00m]|Read field in data file '/Users/hoellinger/WIP3M/run13/nsteps50_final_density_p3m.h5' done.\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores)...\n",
      "[14:43:09\u001b[00m|\u001b[38;5;113mSTATUS    \u001b[00m]|Getting auto-correlation in Fourier space (using 8 cores) done.\n"
     ]
    }
   ],
   "source": [
    "G = read_FourierGrid(simdir + \"input_ss_k_grid.h5\")\n",
    "k = G.k_modes[1:]\n",
    "AliasingCorr = False\n",
    "\n",
    "DELTA = read_field(simdir + \"initial_density.h5\")\n",
    "Pk_INI, Vk_INI = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "Pk_INI, Vk_INI = Pk_INI[1:], Vk_INI[1:]\n",
    "Sk_INI = np.sqrt(Vk_INI)\n",
    "\n",
    "DELTA = read_field(simdir + \"lpt_density.h5\")\n",
    "Pk_LPT, Vk_LPT = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "Pk_LPT, Vk_LPT = Pk_LPT[1:], Vk_LPT[1:]\n",
    "Sk_LPT = np.sqrt(Vk_LPT)\n",
    "\n",
    "DELTA = read_field(simdir + f\"nsteps{nsteps_pm}_final_density_pm.h5\")\n",
    "Pk_PM, Vk_PM = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "Pk_PM, Vk_PM = Pk_PM[1:], Vk_PM[1:]\n",
    "Sk_PM = np.sqrt(Vk_PM)\n",
    "\n",
    "DELTA = read_field(simdir + f\"nsteps{nsteps_cola}_final_density_cola.h5\")\n",
    "Pk_COLA, Vk_COLA = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "Pk_COLA, Vk_COLA = Pk_COLA[1:], Vk_COLA[1:]\n",
    "Sk_COLA = np.sqrt(Vk_COLA)\n",
    "\n",
    "DELTA = read_field(simdir + f\"nsteps{nsteps_spm}_final_density_spm.h5\")\n",
    "Pk_sPM, Vk_sPM = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "Pk_sPM, Vk_sPM = Pk_sPM[1:], Vk_sPM[1:]\n",
    "Sk_sPM = np.sqrt(Vk_sPM)\n",
    "\n",
    "DELTA = read_field(simdir + f\"nsteps{nsteps_p3m}_final_density_p3m.h5\")\n",
    "Pk_P3M, Vk_p3m = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "Pk_P3M, Vk_p3m = Pk_P3M[1:], Vk_p3m[1:]\n",
    "Sk_p3m = np.sqrt(Vk_p3m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "kmin = 1.3454698324203491, kmax = 10.072183609008789\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 174104 characters]",
      "text/plain": [
       "<Figure size 2100x1200 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "Pk_ref = Pk_PM\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 4))\n",
    "\n",
    "ax.set_xscale(\"log\")\n",
    "k = G.k_modes[1:]\n",
    "kmin, kmax = k.min(), k.max()\n",
    "print(f\"kmin = {kmin}, kmax = {kmax}\")\n",
    "log_pad = 0.02\n",
    "log_k_min = np.log10(kmin)\n",
    "log_k_max = np.log10(kmax)\n",
    "log_range = log_k_max - log_k_min\n",
    "xlim_min = 10 ** (log_k_min - log_pad * log_range)\n",
    "xlim_max = 10 ** (log_k_max + log_pad * log_range)\n",
    "\n",
    "plt.xlim(xlim_min, xlim_max)\n",
    "ax.set_ylim([0.2, 1.8])\n",
    "# ax.set_ylim([0.5, 1.5])\n",
    "dark_grey_bnd = 0.05\n",
    "light_grey_bnd = 0.1\n",
    "\n",
    "label_ref = f\"PM with $n_\\\\mathrm{{steps}}={nsteps_pm}$\"\n",
    "\n",
    "line1 = ax.plot([kmin, kmax], [1, 1], color=\"black\", linestyle=\"-\", label=label_ref)\n",
    "\n",
    "# line2 = ax.plot(k, Pk_LPT / Pk_ref, label=\"2LPT\", color=cols[0], linestyle=\"-\")\n",
    "\n",
    "ax.plot(\n",
    "    k,\n",
    "    Pk_COLA / Pk_ref,\n",
    "    label=f\"COLA with $n_\\\\mathrm{{steps}}={nsteps_cola}$\",\n",
    "    linestyle=\"-\",\n",
    "    color=cols[2],\n",
    ")\n",
    "# ax.plot(\n",
    "#     k,\n",
    "#     Pk_PM / Pk_ref,\n",
    "#     label=f\"PM with $n_\\\\mathrm{{steps}}={nsteps_pm}$\",\n",
    "#     linestyle=\"-\",\n",
    "#     color=cols[3],\n",
    "# )\n",
    "ax.plot(\n",
    "    k,\n",
    "    Pk_sPM / Pk_ref,\n",
    "    label=f\"sPM with $n_\\\\mathrm{{steps}}={nsteps_p3m}$\",\n",
    "    linestyle=\"--\",\n",
    "    color=cols[4],\n",
    ")\n",
    "ax.plot(\n",
    "    k,\n",
    "    Pk_P3M / Pk_ref,\n",
    "    label=f\"P3M with $n_\\\\mathrm{{steps}}={nsteps_p3m}$\",\n",
    "    linestyle=\"--\",\n",
    "    color=cols[5],\n",
    ")\n",
    "\n",
    "ax.axhspan(1 - dark_grey_bnd, 1 + dark_grey_bnd, color=\"grey\", alpha=0.2)\n",
    "ax.axhspan(1 - light_grey_bnd, 1 + light_grey_bnd, color=\"grey\", alpha=0.1)\n",
    "\n",
    "for i in range(1, len(k)):\n",
    "    ax.axvline(k[i], color=\"black\", linestyle=\":\", linewidth=1, alpha=0.1)\n",
    "ax.axvline(np.pi * N / L, color=\"grey\", linestyle=\"--\", linewidth=1)  # 1D Nyquist\n",
    "ax.yaxis.set_major_locator(plt.MaxNLocator(6))\n",
    "ax.yaxis.get_major_ticks()[0].label1.set_visible(False)\n",
    "ax.set_xlabel(\"$k$ [$h/\\\\mathrm{Mpc}$]\", fontsize=fs)\n",
    "ax.set_ylabel(\"$P(k)/P_\\\\mathrm{ref}(k)$\", fontsize=fs)\n",
    "ax.tick_params(which=\"both\", direction=\"in\")\n",
    "ax.tick_params(axis=\"both\", which=\"major\", labelsize=fs)\n",
    "ax.tick_params(axis=\"both\", which=\"minor\", labelsize=fs)\n",
    "\n",
    "empty_patch = mpatches.Patch(color=\"none\", label=\"\")\n",
    "handles, labels = plt.gca().get_legend_handles_labels()\n",
    "# handles = [empty_patch, *handles]\n",
    "# labels = [\"\", *labels]\n",
    "plt.legend(\n",
    "    handles,\n",
    "    labels,\n",
    "    loc=\"upper center\",\n",
    "    ncol=2,\n",
    "    bbox_to_anchor=(0.5, -0.2),\n",
    "    fontsize=fs,\n",
    ")\n",
    "fig.savefig(\n",
    "    simdir + \"power_spectrum.png\",\n",
    "    bbox_inches=\"tight\",\n",
    "    dpi=300,\n",
    "    transparent=True,\n",
    ")\n",
    "fig.savefig(\n",
    "    simdir + \"power_spectrum.pdf\",\n",
    "    bbox_inches=\"tight\",\n",
    "    dpi=300,\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Force exerted by particles on other particles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Newton prefactor = 9.95e-03\n",
      "Nyquist: 0.50 Mpc/h\n",
      "Particle length: 0.03 Mpc/h\n",
      "Split scale: 0.31 Mpc/h\n",
      "Short-range reach: 1.41 Mpc/h\n",
      "Figure saved to: /Users/hoellinger/WIP3M/run13/force_diagnostic_comparison.png\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 366068 characters]",
      "text/plain": [
       "<Figure size 1920x1440 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "r1, fmag1, _ = load_force_diagnostic(OutputForceDiagnostic_pm)\n",
    "r4, fmag4, _ = load_force_diagnostic(OutputForceDiagnostic_cola)\n",
    "r2, fmag2, _ = load_force_diagnostic(OutputForceDiagnostic_spm)\n",
    "r3, fmag3, _ = load_force_diagnostic(OutputForceDiagnostic_p3m)\n",
    "\n",
    "rr = np.array([r1, r4, r2, r3], dtype=object)\n",
    "ff = np.array([fmag1, fmag4, fmag2, fmag3], dtype=object)\n",
    "ll = np.array([\"PM\", \"COLA\", \"sPM\", \"P3M\"])\n",
    "\n",
    "ix = [0, 1, 2, 3]\n",
    "\n",
    "Newton_prefactor = (L / Np)**3 / (4*np.pi)\n",
    "print(f\"Newton prefactor = {Newton_prefactor:.2e}\")\n",
    "plot_force_distance_comparison(rr=rr[ix], ff=ff[ix], ll=ll[ix], L=L, Np=Np, Npm=Npm, a=Newton_prefactor, title=\"Particle's contributions to total force\", save_path=simdir + \"force_diagnostic_comparison.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "p3m",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
