wip3m/notebooks/2_plots_convergence_baseline_ts.ipynb

458 lines
17 KiB
Text

{
"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
}