{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "47c34537",
   "metadata": {},
   "source": [
    "Tristan Hoellinger<br/>\n",
    "Institut d'Astrophysique de Paris</br>\n",
    "tristan.hoellinger@iap.fr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b31e6021",
   "metadata": {},
   "source": [
    "# Fit the P3M limiter from existing simulation data\n",
    "\n",
    "## Set up the environment and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0f8c355d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# pyright: reportWildcardImportFromLibrary=false\n",
    "from wip3m import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "31d95b9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "workdir = ROOT_PATH + \"results/\"\n",
    "output_path = OUTPUT_PATH"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "03aa3f4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Automatic reloading of modules\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "from os.path import isfile, join\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import json\n",
    "\n",
    "from pysbmy.timestepping import P3MTimeStepping\n",
    "\n",
    "from wip3m.plot_utils import *  # type: ignore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "57436422",
   "metadata": {},
   "outputs": [],
   "source": [
    "run_id = \"COLAr7/CV2_cola_Pf_L80_N160_Np160\"\n",
    "# run_id = \"COLAr6/V3cola_Pf_L1024_N512_Np512\"\n",
    "# run_id = \"COLAr6/V3cola_Pf_L1024_N1024_Np512\"\n",
    "# run_id = \"COLAr6/V3cola_Pf_L512_N1024_Np512\"\n",
    "# run_id = \"COLAr6/V3cola_Pf_LNNp512\"\n",
    "# run_id = \"COLAr3/cola_2_07_05_L1024_N1024_Np512\"\n",
    "# run_id = \"COLAr3/cola_L512_N1024_Np512\"\n",
    "wd = workdir + run_id + \"/\"\n",
    "simdir = output_path + run_id + \"/\"\n",
    "outdir = simdir + \"plots/\"\n",
    "\n",
    "OutputTimestepsLog_list = [\n",
    "    simdir + \"colap3m1_timestep_log.txt\",\n",
    "    simdir + \"colap3m2_timestep_log.txt\",\n",
    "]\n",
    "TSpath_list = [\n",
    "    wd + f\"p3m1_rscale0_35_ts_p3m.h5\",\n",
    "    wd + f\"p3m2_rescale0_7_ts_p3m.h5\",\n",
    "]\n",
    "\n",
    "# OutputTimestepsLog_list = [\n",
    "#     simdir + \"colap3m1_timestep_log.txt\",\n",
    "#     simdir + \"colap3m2_timestep_log.txt\",\n",
    "# ]\n",
    "# TSpath_list = [\n",
    "#     wd + f\"p3m1_relax2_0_ts_p3m.h5\",\n",
    "#     wd + f\"p3m2_relax0_7_ts_p3m.h5\",\n",
    "#     # wd + f\"p3m2_relax1_0_ts_p3m.h5\",\n",
    "# ]\n",
    "\n",
    "# run_id = \"notebook8\"\n",
    "# wd = workdir + run_id + \"/\"\n",
    "# simdir = output_path + run_id + \"/\"\n",
    "# OutputTimestepsLog_list = [\n",
    "#     simdir + \"timesteps_log_sim0.txt\",\n",
    "#     simdir + \"timesteps_log_sim1.txt\",\n",
    "#     simdir + \"timesteps_log_sim2.txt\",\n",
    "# ]\n",
    "# SimParams_list = [\n",
    "#     simdir + f\"sim_params_p3m_sim0.json\",\n",
    "#     simdir + f\"sim_params_p3m_sim1.json\",\n",
    "#     simdir + f\"sim_params_p3m_sim2.json\",\n",
    "# ]\n",
    "# TSpath_list = [\n",
    "#     wd + f\"sim0_ts_p3m.h5\",\n",
    "#     wd + f\"sim1_ts_p3m.h5\",\n",
    "#     wd + f\"sim2_ts_p3m.h5\",\n",
    "# ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7126d1c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "N_sim_fit = len(OutputTimestepsLog_list)\n",
    "all_sim_params_p3m = []\n",
    "\n",
    "# Load simulation parameters from JSON files\n",
    "# fac_hubble_list = []\n",
    "# fac_bend_list = []\n",
    "# da_early_list = []\n",
    "# for sim_idx in range(N_sim_fit):\n",
    "#     sim_params_path = SimParams_list[sim_idx]\n",
    "#     if not isfile(sim_params_path):\n",
    "#         raise FileNotFoundError(f\"Simulation parameters file not found: {sim_params_path}\")\n",
    "#     with open(sim_params_path, \"r\") as f:\n",
    "#         params = json.load(f)\n",
    "#     print(f\"Loaded simulation parameters for sim_idx=={sim_idx}: {len(params)} keys\")\n",
    "#     all_sim_params_p3m.append(params)\n",
    "#     fac_hubble = all_sim_params_p3m[sim_idx][\"fac_H_custom\"]\n",
    "#     fac_bend = all_sim_params_p3m[sim_idx][\"fac_bend\"]\n",
    "#     da_early = all_sim_params_p3m[sim_idx][\"da_max_early_custom\"]\n",
    "\n",
    "# Alternatively, let the user specify the parameters directly:\n",
    "fac_hubble_list = [DEFAULT_FAC_H_CUSTOM_COLA]*2\n",
    "fac_bend_list = [DEFAULT_FAC_BEND]*2\n",
    "da_early_list = [DEFAULT_DA_MAX_EARLY_CUSTOM]*2\n",
    "# fac_hubble_list = [0.02, 0.02]\n",
    "# fac_bend_list = [0.3125 / 2, 0.3125]\n",
    "# da_early_list = [0.0013, 0.0013]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "69dcd0c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "aa = []\n",
    "da_p3m_list = []\n",
    "for sim_idx in range(N_sim_fit):\n",
    "    OutputTimestepsLog = OutputTimestepsLog_list[sim_idx][:-4] + \"_custom.txt\"\n",
    "    a, _, _, _, _, da_p3m, da_p3m_fit, _, _, _ = np.loadtxt(\n",
    "        OutputTimestepsLog, delimiter=\",\", unpack=True, skiprows=0\n",
    "    )\n",
    "    aa.append(a)\n",
    "    da_p3m_list.append(da_p3m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "701572ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 189392 characters]",
      "text/plain": [
       "<Figure size 1800x1050 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Simulation 0: Coefficients: [-3.74659557  0.39041298  1.90821078  1.10022755  0.5597775   0.09604354]\n",
      "Coefficients saved to /home/hoellinger/wip3m/results/COLAr7/CV2_cola_Pf_L80_N160_Np160/coeffs_p3m_sim0.json\n",
      "Simulation 1: Coefficients: [-3.74967124  0.32498523  1.68724927  0.84016171  0.45636105  0.08504407]\n",
      "Coefficients saved to /home/hoellinger/wip3m/results/COLAr7/CV2_cola_Pf_L80_N160_Np160/coeffs_p3m_sim1.json\n"
     ]
    }
   ],
   "source": [
    "# For each sim_idx, fit the P3M time step distribution and plot the results\n",
    "amin_fit=1e-1\n",
    "amin_plot = 1e-1\n",
    "# amin_plot = 5e-2\n",
    "deg_fit = 5  # Degree of the polynomial fit for the P3M limiter\n",
    "\n",
    "coeffs_list = []  # List to store the coefficients of each polynomial fit\n",
    "for sim_idx in range(N_sim_fit):\n",
    "    # Fit the P3M time step distribution\n",
    "    a = aa[sim_idx]\n",
    "    da_p3m = da_p3m_list[sim_idx]\n",
    "    coeffs = fit_p3m(a, da_p3m, degree=deg_fit, amin=amin_fit)\n",
    "    coeffs_list.append(coeffs)\n",
    "\n",
    "fig = plt.figure(figsize=(6, 3.5))\n",
    "gs = GridSpec(2, 1, height_ratios=[3, 1], hspace=0.05)\n",
    "ax1 = fig.add_subplot(gs[0])\n",
    "ax2 = fig.add_subplot(gs[1], sharex=ax1)\n",
    "\n",
    "for sim_idx in range(N_sim_fit):\n",
    "    a = aa[sim_idx]\n",
    "    da_p3m = da_p3m_list[sim_idx]\n",
    "    coeffs = coeffs_list[sim_idx]\n",
    "    mask = a > amin_plot  # Filter out values below amin_plot\n",
    "    a = a[mask]\n",
    "    da_p3m = da_p3m[mask]\n",
    "    ratio = da_p3m / a\n",
    "    a_dense = np.logspace(np.log10(a.min()), np.log10(a.max()), 500)\n",
    "    fit_vals = evaluate_polynomial_log(coeffs, a_dense)\n",
    "    fit_interp = evaluate_polynomial_log(coeffs, a)\n",
    "    residuals = np.log(ratio / fit_interp)\n",
    "\n",
    "    # P3M limiter and fit\n",
    "    if sim_idx == 0:\n",
    "        ax1.plot(a, ratio, \".\", alpha=0.7, color=f\"C{sim_idx}\", label=r\"$\\Delta t \\propto x_s/\\sqrt{\\langle v^2 \\rangle}$\")\n",
    "        ax1.plot(a_dense, fit_vals, \"-\", lw=1.8, color=f\"C{sim_idx}\", label=rf\"Fit for $a > {amin_fit}$\")\n",
    "    else:\n",
    "        ax1.plot(a, ratio, \".\", alpha=0.7, color=f\"C{sim_idx}\")\n",
    "        ax1.plot(a_dense, fit_vals, \"-\", lw=1.8, color=f\"C{sim_idx}\")\n",
    "    ax1.set_yscale(\"log\")\n",
    "    ax1.set_xscale(\"log\")\n",
    "    ax1.set_ylabel(r\"$\\Delta a/a$\")\n",
    "    ax1.legend(loc=\"upper right\", frameon=False)\n",
    "    ax1.tick_params(which=\"both\", direction=\"in\", top=True, right=True)\n",
    "    ax1.grid(True, which=\"major\", ls=\":\", lw=0.5)\n",
    "\n",
    "    # Residuals\n",
    "    ax2.plot(a, residuals, \"o\", color=\"grey\", alpha=0.6, markersize=3)\n",
    "\n",
    "ax2.axhline(0.0, color=\"black\", lw=1, ls=\"--\")\n",
    "ax2.set_ylabel(\"res\", labelpad=2)\n",
    "ax2.set_xlabel(r\"Scale factor $a$\")\n",
    "ax2.set_xscale(\"log\")\n",
    "ax2.set_yscale(\"symlog\")\n",
    "\n",
    "ax2.tick_params(which=\"both\", direction=\"in\", top=True, right=True)\n",
    "ax2.grid(True, which=\"major\", ls=\":\", lw=0.5)\n",
    "ax1.tick_params(labelbottom=False)\n",
    "ax2.yaxis.set_major_formatter(plt.FuncFormatter(lambda y, _: f\"{y:.3g}\"))\n",
    "plt.show()\n",
    "\n",
    "for sim_idx, coeffs in enumerate(coeffs_list):\n",
    "    print(f\"Simulation {sim_idx}: Coefficients: {coeffs}\")\n",
    "    coeffs_path = wd + f\"coeffs_p3m_sim{sim_idx}.json\"\n",
    "    with open(coeffs_path, \"w\") as f:\n",
    "        json.dump(coeffs.tolist(), f)\n",
    "    print(f\"Coefficients saved to {coeffs_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c1c096bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14:11:00|\u001b[38;5;113mSTATUS    \u001b[00m]|Read custom timestepping configuration in '/home/hoellinger/wip3m/results/COLAr7/CV2_cola_Pf_L80_N160_Np160/p3m1_rscale0_35_ts_p3m.h5'...\n"
     ]
    },
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] Unable to synchronously open file (unable to open file: name = '/home/hoellinger/wip3m/results/COLAr7/CV2_cola_Pf_L80_N160_Np160/p3m1_rscale0_35_ts_p3m.h5', errno = 2, error message = 'No such file or directory', flags = 0, o_flags = 0)",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mFileNotFoundError\u001b[39m                         Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[8]\u001b[39m\u001b[32m, line 3\u001b[39m\n\u001b[32m      1\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m sim_idx \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(N_sim_fit):\n\u001b[32m      2\u001b[39m     TSpath = TSpath_list[sim_idx]\n\u001b[32m----> \u001b[39m\u001b[32m3\u001b[39m     TS = \u001b[43mP3MTimeStepping\u001b[49m\u001b[43m.\u001b[49m\u001b[43mread\u001b[49m\u001b[43m(\u001b[49m\u001b[43mTSpath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m      4\u001b[39m     aiDrift = TS.aiDrift\n\u001b[32m      5\u001b[39m     nsteps = TS.nsteps\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/simbelmyne2025/DEV_simbelmyne_private/build/__editable__.pysbmy-0.5-py3-none-any/pysbmy/timestepping.py:1064\u001b[39m, in \u001b[36mP3MTimeStepping.read\u001b[39m\u001b[34m(fname)\u001b[39m\n\u001b[32m   1062\u001b[39m PrintMessage(\u001b[32m3\u001b[39m, \u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mRead custom timestepping configuration in \u001b[39m\u001b[33m'\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mfname\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m'\u001b[39m\u001b[33m...\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m   1063\u001b[39m INDENT()\n\u001b[32m-> \u001b[39m\u001b[32m1064\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[43mh5\u001b[49m\u001b[43m.\u001b[49m\u001b[43mFile\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[33;43m'\u001b[39;49m\u001b[33;43mr\u001b[39;49m\u001b[33;43m'\u001b[39;49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m hf:\n\u001b[32m   1065\u001b[39m     ai = hf.attrs[\u001b[33m'\u001b[39m\u001b[33m/info/scalars/ai\u001b[39m\u001b[33m'\u001b[39m]\n\u001b[32m   1066\u001b[39m     af = hf.attrs[\u001b[33m'\u001b[39m\u001b[33m/info/scalars/af\u001b[39m\u001b[33m'\u001b[39m]\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/p3m/lib/python3.13/site-packages/h5py/_hl/files.py:564\u001b[39m, in \u001b[36mFile.__init__\u001b[39m\u001b[34m(self, name, mode, driver, libver, userblock_size, swmr, rdcc_nslots, rdcc_nbytes, rdcc_w0, track_order, fs_strategy, fs_persist, fs_threshold, fs_page_size, page_buf_size, min_meta_keep, min_raw_keep, locking, alignment_threshold, alignment_interval, meta_block_size, **kwds)\u001b[39m\n\u001b[32m    555\u001b[39m     fapl = make_fapl(driver, libver, rdcc_nslots, rdcc_nbytes, rdcc_w0,\n\u001b[32m    556\u001b[39m                      locking, page_buf_size, min_meta_keep, min_raw_keep,\n\u001b[32m    557\u001b[39m                      alignment_threshold=alignment_threshold,\n\u001b[32m    558\u001b[39m                      alignment_interval=alignment_interval,\n\u001b[32m    559\u001b[39m                      meta_block_size=meta_block_size,\n\u001b[32m    560\u001b[39m                      **kwds)\n\u001b[32m    561\u001b[39m     fcpl = make_fcpl(track_order=track_order, fs_strategy=fs_strategy,\n\u001b[32m    562\u001b[39m                      fs_persist=fs_persist, fs_threshold=fs_threshold,\n\u001b[32m    563\u001b[39m                      fs_page_size=fs_page_size)\n\u001b[32m--> \u001b[39m\u001b[32m564\u001b[39m     fid = \u001b[43mmake_fid\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43muserblock_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfapl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfcpl\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mswmr\u001b[49m\u001b[43m=\u001b[49m\u001b[43mswmr\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    566\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(libver, \u001b[38;5;28mtuple\u001b[39m):\n\u001b[32m    567\u001b[39m     \u001b[38;5;28mself\u001b[39m._libver = libver\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/p3m/lib/python3.13/site-packages/h5py/_hl/files.py:238\u001b[39m, in \u001b[36mmake_fid\u001b[39m\u001b[34m(name, mode, userblock_size, fapl, fcpl, swmr)\u001b[39m\n\u001b[32m    236\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m swmr \u001b[38;5;129;01mand\u001b[39;00m swmr_support:\n\u001b[32m    237\u001b[39m         flags |= h5f.ACC_SWMR_READ\n\u001b[32m--> \u001b[39m\u001b[32m238\u001b[39m     fid = \u001b[43mh5f\u001b[49m\u001b[43m.\u001b[49m\u001b[43mopen\u001b[49m\u001b[43m(\u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mflags\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfapl\u001b[49m\u001b[43m=\u001b[49m\u001b[43mfapl\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    239\u001b[39m \u001b[38;5;28;01melif\u001b[39;00m mode == \u001b[33m'\u001b[39m\u001b[33mr+\u001b[39m\u001b[33m'\u001b[39m:\n\u001b[32m    240\u001b[39m     fid = h5f.open(name, h5f.ACC_RDWR, fapl=fapl)\n",
      "\u001b[36mFile \u001b[39m\u001b[32mh5py/_objects.pyx:54\u001b[39m, in \u001b[36mh5py._objects.with_phil.wrapper\u001b[39m\u001b[34m()\u001b[39m\n",
      "\u001b[36mFile \u001b[39m\u001b[32mh5py/_objects.pyx:55\u001b[39m, in \u001b[36mh5py._objects.with_phil.wrapper\u001b[39m\u001b[34m()\u001b[39m\n",
      "\u001b[36mFile \u001b[39m\u001b[32mh5py/h5f.pyx:102\u001b[39m, in \u001b[36mh5py.h5f.open\u001b[39m\u001b[34m()\u001b[39m\n",
      "\u001b[31mFileNotFoundError\u001b[39m: [Errno 2] Unable to synchronously open file (unable to open file: name = '/home/hoellinger/wip3m/results/COLAr7/CV2_cola_Pf_L80_N160_Np160/p3m1_rscale0_35_ts_p3m.h5', errno = 2, error message = 'No such file or directory', flags = 0, o_flags = 0)"
     ]
    }
   ],
   "source": [
    "for sim_idx in range(N_sim_fit):\n",
    "    TSpath = TSpath_list[sim_idx]\n",
    "    TS = P3MTimeStepping.read(TSpath)\n",
    "    aiDrift = TS.aiDrift\n",
    "    nsteps = TS.nsteps\n",
    "    OutputTimestepsLog = OutputTimestepsLog_list[sim_idx]\n",
    "    plot_custom_timestepping_diagnostics(\n",
    "        log_path=OutputTimestepsLog,\n",
    "        aiDrift=aiDrift,\n",
    "        TimeStepDistribution=3,\n",
    "        nsteps=nsteps,\n",
    "        ymin=3e-3,\n",
    "        ymax=0.5,\n",
    "        fac_hubble=fac_hubble_list[sim_idx],\n",
    "        fac_bend=fac_bend_list[sim_idx],\n",
    "        da_max_early=da_early_list[sim_idx],\n",
    "        da_max_late=DEFAULT_DA_MAX_LATE_CUSTOM,\n",
    "        show=False,\n",
    "    )\n",
    "    fit_interp = evaluate_polynomial_log(coeffs_list[sim_idx], aa[sim_idx])\n",
    "    plt.plot(aa[sim_idx], fit_interp, label=r\"$\\Delta t \\propto x_s/\\sqrt{\\langle v^2 \\rangle}$\", color=\"black\", alpha=0.7, zorder=10)\n",
    "    plt.title(f\"Simulation {sim_idx}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f225b2f9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85b01324",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dadb9198",
   "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": 5
}
