{
 "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": null,
   "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 = \"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_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 = [0.02, 0.02]\n",
    "fac_bend_list = [0.04375/2, 0.04375]\n",
    "da_early_list = [0.0013, 0.0013]\n"
   ]
  },
  {
   "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, _, _, _ = 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, 188716 characters]",
      "text/plain": [
       "<Figure size 1800x1050 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Simulation 0: Coefficients: [-4.41928363  0.24380666  1.92034984  1.21626721  0.81694801  0.20590599]\n",
      "Coefficients saved to /home/hoellinger/wip3m/results/COLAr3/cola_2_07_05_L1024_N1024_Np512/coeffs_p3m_sim0.json\n",
      "Simulation 1: Coefficients: [-4.40335594  0.26813626  2.11941734  1.40650529  0.86356166  0.20573432]\n",
      "Coefficients saved to /home/hoellinger/wip3m/results/COLAr3/cola_2_07_05_L1024_N1024_Np512/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=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": null,
   "id": "c1c096bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[05:03:35|\u001b[38;5;113mSTATUS    \u001b[00m]|Read custom timestepping configuration in '/home/hoellinger/wip3m/results/COLAr3/cola_2_07_05_L1024_N1024_Np512/p3m1_relax2_0_ts_p3m.h5'...\n",
      "[05:03:35|\u001b[38;5;113mSTATUS    \u001b[00m]|Read custom timestepping configuration in '/home/hoellinger/wip3m/results/COLAr3/cola_2_07_05_L1024_N1024_Np512/p3m1_relax2_0_ts_p3m.h5' done.\n",
      "[05:03:35|\u001b[1;36mINFO      \u001b[00m]|\u001b[38;5;147m(wip3m.plot_utils)\u001b[00m Plotting timestep limiters from /data70/hoellinger/WIP3M/COLAr3/cola_2_07_05_L1024_N1024_Np512/colap3m1_timestep_log.txt and /data70/hoellinger/WIP3M/COLAr3/cola_2_07_05_L1024_N1024_Np512/colap3m1_timestep_log_custom.txt...\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "not enough values to unpack (expected 10, got 9)",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mValueError\u001b[39m                                Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[8]\u001b[39m\u001b[32m, line 7\u001b[39m\n\u001b[32m      5\u001b[39m nsteps = TS.nsteps\n\u001b[32m      6\u001b[39m OutputTimestepsLog = OutputTimestepsLog_list[sim_idx]\n\u001b[32m----> \u001b[39m\u001b[32m7\u001b[39m \u001b[43mplot_custom_timestepping_diagnostics\u001b[49m\u001b[43m(\u001b[49m\n\u001b[32m      8\u001b[39m \u001b[43m    \u001b[49m\u001b[43mlog_path\u001b[49m\u001b[43m=\u001b[49m\u001b[43mOutputTimestepsLog\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m      9\u001b[39m \u001b[43m    \u001b[49m\u001b[43maiDrift\u001b[49m\u001b[43m=\u001b[49m\u001b[43maiDrift\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     10\u001b[39m \u001b[43m    \u001b[49m\u001b[43mTimeStepDistribution\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m3\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m     11\u001b[39m \u001b[43m    \u001b[49m\u001b[43mnsteps\u001b[49m\u001b[43m=\u001b[49m\u001b[43mnsteps\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     12\u001b[39m \u001b[43m    \u001b[49m\u001b[43mymin\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m3e-3\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m     13\u001b[39m \u001b[43m    \u001b[49m\u001b[43mymax\u001b[49m\u001b[43m=\u001b[49m\u001b[32;43m0.5\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[32m     14\u001b[39m \u001b[43m    \u001b[49m\u001b[43mfac_hubble\u001b[49m\u001b[43m=\u001b[49m\u001b[43mfac_hubble_list\u001b[49m\u001b[43m[\u001b[49m\u001b[43msim_idx\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     15\u001b[39m \u001b[43m    \u001b[49m\u001b[43mfac_bend\u001b[49m\u001b[43m=\u001b[49m\u001b[43mfac_bend_list\u001b[49m\u001b[43m[\u001b[49m\u001b[43msim_idx\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     16\u001b[39m \u001b[43m    \u001b[49m\u001b[43mda_max_early\u001b[49m\u001b[43m=\u001b[49m\u001b[43mda_early_list\u001b[49m\u001b[43m[\u001b[49m\u001b[43msim_idx\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     17\u001b[39m \u001b[43m    \u001b[49m\u001b[43mda_max_late\u001b[49m\u001b[43m=\u001b[49m\u001b[43mDEFAULT_DA_MAX_LATE_CUSTOM\u001b[49m\u001b[43m,\u001b[49m\n\u001b[32m     18\u001b[39m \u001b[43m    \u001b[49m\u001b[43mshow\u001b[49m\u001b[43m=\u001b[49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m,\u001b[49m\n\u001b[32m     19\u001b[39m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     20\u001b[39m fit_interp = evaluate_polynomial_log(coeffs_list[sim_idx], aa[sim_idx])\n\u001b[32m     21\u001b[39m plt.plot(aa[sim_idx], fit_interp, label=\u001b[33mr\u001b[39m\u001b[33m\"\u001b[39m\u001b[33m$\u001b[39m\u001b[33m\\\u001b[39m\u001b[33mDelta t \u001b[39m\u001b[33m\\\u001b[39m\u001b[33mpropto x_s/\u001b[39m\u001b[33m\\\u001b[39m\u001b[33msqrt\u001b[39m\u001b[33m{\u001b[39m\u001b[33m\\\u001b[39m\u001b[33mlangle v^2 \u001b[39m\u001b[33m\\\u001b[39m\u001b[33mrangle}$\u001b[39m\u001b[33m\"\u001b[39m, color=\u001b[33m\"\u001b[39m\u001b[33mblack\u001b[39m\u001b[33m\"\u001b[39m, alpha=\u001b[32m0.7\u001b[39m, zorder=\u001b[32m10\u001b[39m)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/wip3m/src/build/__editable__.wip3m-0.1.0-py3-none-any/wip3m/plot_utils.py:492\u001b[39m, in \u001b[36mplot_custom_timestepping_diagnostics\u001b[39m\u001b[34m(log_path, aiDrift, TimeStepDistribution, nsteps, fac_hubble, fac_bend, fac_p3m_fit, da_max_early, da_max_late, ymin, ymax, save_path, show)\u001b[39m\n\u001b[32m    490\u001b[39m logger.info(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mPlotting timestep limiters from \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mlog_path\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m and \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mlog_path_custom\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m...\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m    491\u001b[39m INDENT()\n\u001b[32m--> \u001b[39m\u001b[32m492\u001b[39m a, da, da_max, da_dyn, da_bend, da_p3m, da_p3m_fit, da_hubble, da_hubble_lenient, da_late = np.loadtxt(log_path_custom, delimiter=\u001b[33m\"\u001b[39m\u001b[33m,\u001b[39m\u001b[33m\"\u001b[39m, unpack=\u001b[38;5;28;01mTrue\u001b[39;00m, skiprows=\u001b[32m0\u001b[39m)\n\u001b[32m    494\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m aiDrift \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m    495\u001b[39m     \u001b[38;5;28;01massert\u001b[39;00m np.allclose(a, aiDrift), \u001b[33m\"\u001b[39m\u001b[33maiDrift does not match the time steps in the log file.\u001b[39m\u001b[33m\"\u001b[39m\n",
      "\u001b[31mValueError\u001b[39m: not enough values to unpack (expected 10, got 9)"
     ]
    }
   ],
   "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
}
