{
 "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": [
    "# Plot time step convergence tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the environment and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pyright: reportWildcardImportFromLibrary=false\n",
    "import numpy as np\n",
    "\n",
    "from pysbmy.fft import read_FourierGrid\n",
    "from pysbmy.field import read_field\n",
    "from pysbmy.correlations import get_autocorrelation\n",
    "\n",
    "from wip3m import *\n",
    "from wip3m.plot_utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workdir = ROOT_PATH + \"results/\"\n",
    "output_path = OUTPUT_PATH\n",
    "\n",
    "# run_id = \"night6\"\n",
    "# nsteps_pmref = 500\n",
    "# nsteps_pm1 = 300\n",
    "# nsteps_pm2 = 100\n",
    "# nsteps_cola = 20\n",
    "# nsteps_spm = nsteps_p3m1 = 1000\n",
    "# nsteps_p3m2 = 800\n",
    "# nsteps_p3m3 = 500\n",
    "\n",
    "run_id = \"night4\"\n",
    "nsteps_pmref = 500\n",
    "nsteps_pm1 = 300\n",
    "nsteps_pm2 = 100\n",
    "nsteps_cola = 20\n",
    "nsteps_spm = nsteps_p3m1 = 500\n",
    "nsteps_p3m2 = 300\n",
    "nsteps_p3m3 = 100\n",
    "\n",
    "# run_id = \"medium_test2\"\n",
    "# nsteps_pmref = 500\n",
    "# nsteps_pm1 = 300\n",
    "# nsteps_pm2 = 200\n",
    "# nsteps_cola = 20\n",
    "# nsteps_spm = nsteps_p3m1 = 500\n",
    "# nsteps_p3m2 = 300\n",
    "# nsteps_p3m3 = 200"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In principle nothing needs to be changed below this cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simdir = output_path + run_id + \"/\"\n",
    "input_ss_file = simdir + \"input_ss_k_grid.h5\"\n",
    "wd = workdir + run_id + \"/\"\n",
    "with open(wd + \"sim_params.txt\", \"r\") as f:\n",
    "    sim_params = eval(f.read())\n",
    "L = sim_params[\"L\"]  # Box size in Mpc/h\n",
    "N = sim_params[\"N\"]  # Density grid size\n",
    "Np = sim_params[\"Np\"]  # Number of dark matter particles per spatial dimension\n",
    "Npm = sim_params[\"Npm\"]  # PM grid size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the evolved dark matter density fields"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice_ijk = (N // 2, slice(None), slice(None))\n",
    "DELTA_LPT = read_field(simdir + \"lpt_density.h5\").data[slice_ijk]\n",
    "DELTA_PMREF = read_field(simdir + f\"pmref_nsteps{nsteps_pmref}_final_density_pm.h5\").data[slice_ijk]\n",
    "DELTA_PM1 = read_field(simdir + f\"pm1_nsteps{nsteps_pm1}_final_density_pm.h5\").data[slice_ijk]\n",
    "DELTA_PM2 = read_field(simdir + f\"pm2_nsteps{nsteps_pm2}_final_density_pm.h5\").data[slice_ijk]\n",
    "DELTA_COLA = read_field(simdir + f\"cola_nsteps{nsteps_cola}_final_density_cola.h5\").data[slice_ijk]\n",
    "DELTA_SPM = read_field(simdir + f\"spm_nsteps{nsteps_spm}_final_density_spm.h5\").data[slice_ijk]\n",
    "DELTA_P3M1 = read_field(simdir + f\"p3m1_nsteps{nsteps_p3m1}_final_density_p3m.h5\").data[slice_ijk]\n",
    "DELTA_P3M2 = read_field(simdir + f\"p3m2_nsteps{nsteps_p3m2}_final_density_p3m.h5\").data[slice_ijk]\n",
    "DELTA_P3M3 = read_field(simdir + f\"p3m3_nsteps{nsteps_p3m3}_final_density_p3m.h5\").data[slice_ijk]\n",
    "diff_pm1_pmref = DELTA_PM1 - DELTA_PMREF\n",
    "diff_pm2_pmref = DELTA_PM2 - DELTA_PMREF\n",
    "diff_p3m1_pmref = DELTA_P3M1 - DELTA_PMREF\n",
    "diff_p3m2_pm1 = DELTA_P3M2 - DELTA_PM1\n",
    "diff_p3m3_pm2 = DELTA_P3M3 - DELTA_PM2\n",
    "diff_p3m1_spm = DELTA_P3M1 - DELTA_SPM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fields = [\"pmref\", \"pm1\", \"pm2\", \"p3m1\", \"p3m2\", \"p3m3\", \"diff_p3m1_spm\", \"diff_pm1_pmref\", \"diff_pm2_pmref\", \"diff_p3m1_pmref\", \"diff_p3m2_pm1\", \"diff_p3m3_pm2\"]  # fields to plot\n",
    "# fields = [\"lpt\", \"pmref\", \"pm1\", \"pm2\", \"cola\", \"spm\", \"p3m1\", \"p3m2\", \"p3m3\", \"diff_p3m1_spm\", \"diff_pm1_pmref\", \"diff_pm2_pmref\", \"diff_p3m1_pmref\", \"diff_p3m2_pm1\", \"diff_p3m3_pm2\"]  # al possible fields\n",
    "\n",
    "figname = \"_\".join(fields)\n",
    "slices_dict = {\n",
    "    \"lpt\": DELTA_LPT,\n",
    "    \"cola\": DELTA_COLA,\n",
    "    \"pmref\": DELTA_PMREF,\n",
    "    \"pm1\": DELTA_PM1,\n",
    "    \"pm2\": DELTA_PM2,\n",
    "    \"spm\": DELTA_SPM,\n",
    "    \"p3m1\": DELTA_P3M1,\n",
    "    \"p3m2\": DELTA_P3M2,\n",
    "    \"p3m3\": DELTA_P3M3,\n",
    "    \"diff_pm1_pmref\": diff_pm1_pmref,\n",
    "    \"diff_pm2_pmref\": diff_pm2_pmref,\n",
    "    \"diff_p3m1_pmref\": diff_p3m1_pmref,\n",
    "    \"diff_p3m2_pm1\": diff_p3m2_pm1,\n",
    "    \"diff_p3m3_pm2\": diff_p3m3_pm2,\n",
    "    \"diff_p3m1_spm\": diff_p3m1_spm,\n",
    "}\n",
    "titles_dict = {\n",
    "    \"lpt\": \"LPT\",\n",
    "    \"pmref\": f\"PMref $n_\\\\mathrm{{steps}}={nsteps_pmref}$\",\n",
    "    \"pm1\": f\"PM1 $n_\\\\mathrm{{steps}}={nsteps_pm1}$\",\n",
    "    \"pm2\": f\"PM2 $n_\\\\mathrm{{steps}}={nsteps_pm2}$\",\n",
    "    \"cola\": f\"COLA $n_\\\\mathrm{{steps}}={nsteps_cola}$\",\n",
    "    \"spm\": f\"sPM $n_\\\\mathrm{{steps}}={nsteps_spm}$\",\n",
    "    \"p3m1\": f\"P3M1 $n_\\\\mathrm{{steps}}={nsteps_p3m1}$\",\n",
    "    \"p3m2\": f\"P3M2 $n_\\\\mathrm{{steps}}={nsteps_p3m2}$\",\n",
    "    \"p3m3\": f\"P3M3 $n_\\\\mathrm{{steps}}={nsteps_p3m3}$\",\n",
    "    \"diff_pm1_pmref\": r\"$\\delta_{\\rm PM1}-\\delta_{\\rm PMref}$\",\n",
    "    \"diff_pm2_pmref\": r\"$\\delta_{\\rm PM2}-\\delta_{\\rm PMref}$\",\n",
    "    \"diff_p3m1_pmref\": r\"$\\delta_{\\rm P3M1}-\\delta_{\\rm PMref}$\",\n",
    "    \"diff_p3m2_pm1\": r\"$\\delta_{\\rm P3M2}-\\delta_{\\rm PM1}$\",\n",
    "    \"diff_p3m3_pm2\": r\"$\\delta_{\\rm P3M3}-\\delta_{\\rm PM2}$\",\n",
    "    \"diff_p3m1_spm\": r\"$\\delta_{\\rm P3M1}-\\delta_{\\rm sPM}$\",\n",
    "}\n",
    "\n",
    "ncols = 6  # Max panels per row\n",
    "npanels = len(fields)\n",
    "nrows = np.ceil(npanels / ncols).astype(int)\n",
    "\n",
    "fig, axs = plt.subplots(\n",
    "    nrows=nrows, ncols=ncols, figsize=(3 * ncols, 4.6 * nrows), sharey=True\n",
    ")\n",
    "\n",
    "axs = axs.flatten()\n",
    "ims = []\n",
    "\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(\n",
    "            vmin=-np.log(1 + np.abs(np.min(data))),\n",
    "            vcenter=0,\n",
    "            vmax=np.log10(1 + np.abs(np.max(data))),\n",
    "        )\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",
    "# Hide unused axes\n",
    "for i in range(npanels, len(axs)):\n",
    "    axs[i].axis(\"off\")\n",
    "\n",
    "# Shared axes labels\n",
    "for i, ax in enumerate(axs[:npanels]):\n",
    "    if i % ncols == 0:\n",
    "        ax.set_yticks([0, N // 2, N])\n",
    "        ax.set_yticklabels([f\"{-L/2:.0f}\", \"0\", f\"{L/2:.0f}\"], fontsize=fs)\n",
    "        ax.set_ylabel(r\"Mpc/$h$\", size=GLOBAL_FS_SMALL)\n",
    "\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",
    "# Colorbars\n",
    "for ax, (im, key) in zip(axs[:npanels], 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(\n",
    "            r\"$\\textrm{sgn}\\left(\\Delta\\delta\\right)\\log_{10}(1 + |\\Delta\\delta|)$\",\n",
    "            fontsize=fs,\n",
    "        )\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",
    "\n",
    "fig.savefig(simdir + f\"{figname}.png\", bbox_inches=\"tight\", dpi=300, transparent=True)\n",
    "fig.savefig(simdir + f\"{figname}.pdf\", bbox_inches=\"tight\", dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# full_field_p3m1 = np.log10(2+read_field(simdir + f\"p3m1_nsteps{nsteps_p3m1}_final_density_p3m.h5\").data)\n",
    "\n",
    "# if N <= 128:\n",
    "#     fig = plotly_3d(full_field_p3m1, 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_p3m1[\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": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = read_FourierGrid(input_ss_file)\n",
    "k = G.k_modes\n",
    "AliasingCorr = False\n",
    "\n",
    "DELTA = read_field(simdir + \"initial_density.h5\")\n",
    "Pk_INI, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + \"lpt_density.h5\")\n",
    "Pk_LPT, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"pmref_nsteps{nsteps_pmref}_final_density_pm.h5\")\n",
    "Pk_PMref, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"pm1_nsteps{nsteps_pm1}_final_density_pm.h5\")\n",
    "Pk_PM1, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"pm2_nsteps{nsteps_pm2}_final_density_pm.h5\")\n",
    "Pk_PM2, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"cola_nsteps{nsteps_cola}_final_density_cola.h5\")\n",
    "Pk_COLA, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"spm_nsteps{nsteps_spm}_final_density_spm.h5\")\n",
    "Pk_sPM, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"p3m1_nsteps{nsteps_p3m1}_final_density_p3m.h5\")\n",
    "Pk_P3M1, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"p3m2_nsteps{nsteps_p3m2}_final_density_p3m.h5\")\n",
    "Pk_P3M2, _ = get_autocorrelation(DELTA, G, AliasingCorr)\n",
    "\n",
    "DELTA = read_field(simdir + f\"p3m3_nsteps{nsteps_p3m3}_final_density_p3m.h5\")\n",
    "Pk_P3M3, _ = get_autocorrelation(DELTA, G, AliasingCorr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Pk_ref = Pk_PMref\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 4))\n",
    "\n",
    "ax.set_xscale(\"log\")\n",
    "k = G.k_modes\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",
    "dark_grey_bnd = 0.05\n",
    "light_grey_bnd = 0.1\n",
    "\n",
    "label_ref = f\"Ref (PM, $n_\\\\mathrm{{steps}}={nsteps_pmref}$)\"\n",
    "\n",
    "ax.plot([kmin, kmax], [1, 1], color=\"black\", linestyle=\"-\", label=label_ref)\n",
    "# ax.plot(k, Pk_LPT / Pk_ref, label=\"2LPT\", linestyle=\"--\")\n",
    "fields_to_plot = [\n",
    "    (\"PM1\", Pk_PM1),\n",
    "    (\"PM2\", Pk_PM2),\n",
    "    (\"COLA\", Pk_COLA),\n",
    "    (\"sPM\", Pk_sPM),\n",
    "    (\"P3M1\", Pk_P3M1),\n",
    "    (\"P3M2\", Pk_P3M2),\n",
    "    (\"P3M3\", Pk_P3M3),\n",
    "]\n",
    "for field_name, Pk in fields_to_plot:\n",
    "    label = f\"{field_name}, $n_\\\\mathrm{{steps}}={eval(f'nsteps_{field_name.lower()}')}$\"\n",
    "    ax.plot(k, Pk / Pk_ref, label=label, linestyle=\"--\")\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.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",
    "# Characteristic vertical reference scales\n",
    "nyquist = np.pi * N / L\n",
    "nyquist_PM = np.pi * Npm / L\n",
    "epsilon = 0.03 * L / Np\n",
    "particle_length = 2*epsilon\n",
    "xs = 1.25 * L / Npm\n",
    "xr = 4.5 * xs\n",
    "particle_wavenumber = 2*np.pi / particle_length  # Too large to be shown\n",
    "xs_inv = 2*np.pi / xs\n",
    "xr_inv = 2*np.pi / xr\n",
    "line1 = ax.axvline(x=nyquist, color=\"black\", linestyle=\"--\", lw=1, label=\"Nyquist (density grid)\")\n",
    "line1 = ax.axvline(x=nyquist_PM, color=\"black\", linestyle=\"-\", lw=1, label=\"Nyquist (PM grid)\")\n",
    "line3 = ax.axvline(x=xs_inv, color=\"gray\", linestyle=\"-.\", lw=2, label=r\"Split wavenumber $x_s$\")\n",
    "line4 = ax.axvline(x=xr_inv, color=\"gray\", linestyle=\":\", lw=2, label=r\"Short-range reach $x_r$\")\n",
    "print(f\"Nyquist (density grid): {nyquist:.2f} h/Mpc\")\n",
    "print(f\"Nyquist (PM grid): {nyquist_PM:.2f} h/Mpc\")\n",
    "print(f\"Particle wavenumber: {particle_wavenumber:.2f} h/Mpc\")\n",
    "print(f\"Split wavenumber: {xs_inv:.2f} h/Mpc\")\n",
    "print(f\"Short-range reach: {xr_inv:.2f} h/Mpc\")\n",
    "\n",
    "handles, labels = plt.gca().get_legend_handles_labels()\n",
    "# empty_patch = mpatches.Patch(color=\"none\", label=\"\")\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",
    "    frameon=False,\n",
    ")\n",
    "fig.savefig(\n",
    "    simdir + \"power_spectra.png\",\n",
    "    bbox_inches=\"tight\",\n",
    "    dpi=300,\n",
    "    transparent=True,\n",
    ")\n",
    "fig.savefig(\n",
    "    simdir + \"power_spectra.pdf\",\n",
    "    bbox_inches=\"tight\",\n",
    "    dpi=300,\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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
}