hoellinger/wip3m
0
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

validate P3M timestep convergence for baseline linear-in-scale-factor steps

Browse source
This commit is contained in:
Tristan Hoellinger 2025-05-25 18:36:50 +02:00
parent 13e6c3b32d
commit a38f1b2585
7 changed files with 89782 additions and 1889 deletions
Download patch file Download diff file Expand all files Collapse all files

73
src/wip3m/convergence_baseline_ts_parser.py
View file

@ -25,7 +25,12 @@ from pysbmy.fft import FourierGrid
from wip3m import *
from wip3m.tools import get_k_max, generate_sim_params, generate_white_noise_Field, run_simulation
from wip3m.params import params_planck_kmax_missing, cosmo_small_to_full_dict, z2a, BASELINE_SEEDPHASE
from wip3m.params import (
params_planck_kmax_missing,
cosmo_small_to_full_dict,
z2a,
BASELINE_SEEDPHASE,
)
from wip3m.plot_utils import *
from wip3m.logger import getCustomLogger, INDENT, UNINDENT
@ -37,34 +42,48 @@ output_path = OUTPUT_PATH
from argparse import ArgumentParser
parser = ArgumentParser(
description="Run convergence tests towards implementing P3M gravity."
)
parser = ArgumentParser(description="Run convergence tests towards implementing P3M gravity.")
parser.add_argument("--run_id", type=str, help="Simulation run identifier.")
parser.add_argument("--L", type=float, default=32.0, help="Box size in Mpc/h.")
parser.add_argument("--N", type=int, default=128, help="Density grid size per dimension.")
parser.add_argument("--Np", type=int, default=32, help="Number of particles per dimension.")
parser.add_argument("--Npm", type=int, default=64, help="PM mesh resolution per dimension.")
parser.add_argument("--n_Tiles", type=int, default=8, help="Number of tiles per dimension for short-range PP.")
parser.add_argument(
"--n_Tiles", type=int, default=8, help="Number of tiles per dimension for short-range PP."
)
# Timestep settings per method
for scheme in ["pmref", "pm1", "pm2", "cola", "spm", "p3m1", "p3m2", "p3m3"]:
parser.add_argument(f"--nsteps_{scheme}", type=int, default={
"pmref": 200, "pm1": 100, "pm2": 20,
"cola": 10, "spm": 200,
"p3m1": 200, "p3m2": 100, "p3m3": 20
}[scheme], help=f"Number of timesteps for {scheme.upper()}.")
parser.add_argument(f"--tsd_{scheme}", type=int, default=0,
help=f"TimeStepDistribution ID for {scheme.upper()}.")
parser.add_argument(
f"--nsteps_{scheme}",
type=int,
default={
"pmref": 200,
"pm1": 100,
"pm2": 20,
"cola": 10,
"spm": 200,
"p3m1": 200,
"p3m2": 100,
"p3m3": 20,
}[scheme],
help=f"Number of timesteps for {scheme.upper()}.",
)
parser.add_argument(
f"--tsd_{scheme}",
type=int,
default=0,
help=f"TimeStepDistribution ID for {scheme.upper()}.",
)
args = parser.parse_args()
if __name__ == "__main__":
try:
go_beyond_Nyquist_ss = True # for the summary statistics
force = True
force_hard = True
force = False
force_hard = False
run_id = args.run_id
L = args.L
@ -104,14 +123,13 @@ if __name__ == "__main__":
logger.info(f"Number of timesteps for P3M2: {nsteps_p3m2}")
logger.info(f"Number of timesteps for P3M3: {nsteps_p3m3}")
INDENT()
corner = -L / 2.0
RedshiftLPT = 19.0
RedshiftFCs = 0.0
ai = z2a(RedshiftLPT)
af = z2a(RedshiftFCs)
k_max = get_k_max(L, N) # k_max in h/Mpc
print(f"k_max = {k_max}")
cosmo = params_planck_kmax_missing.copy()
cosmo["k_max"] = k_max
@ -120,6 +138,7 @@ if __name__ == "__main__":
logdir = simdir + "logs/"
if force_hard:
import shutil
if Path(simdir).exists():
shutil.rmtree(simdir)
if Path(wd).exists():
@ -143,6 +162,22 @@ if __name__ == "__main__":
"n_Tiles": n_Tiles,
"RedshiftLPT": RedshiftLPT,
"RedshiftFCs": RedshiftFCs,
"nsteps_pmref": nsteps_pmref,
"nsteps_pm1": nsteps_pm1,
"nsteps_pm2": nsteps_pm2,
"nsteps_cola": nsteps_cola,
"nsteps_spm": nsteps_spm,
"nsteps_p3m1": nsteps_p3m1,
"nsteps_p3m2": nsteps_p3m2,
"nsteps_p3m3": nsteps_p3m3,
"tsd_pmref": tsd_pmref,
"tsd_pm1": tsd_pm1,
"tsd_pm2": tsd_pm2,
"tsd_cola": tsd_cola,
"tsd_spm": tsd_spm,
"tsd_p3m1": tsd_p3m1,
"tsd_p3m2": tsd_p3m2,
"tsd_p3m3": tsd_p3m3,
}
with open(wd + "sim_params.txt", "w") as f:
f.write(f"{sim_params}\n")
@ -247,7 +282,9 @@ if __name__ == "__main__":
("p3m3", p3m3_params),
]
for name, params in all_sim_params:
logger.info(f"Generating parameters for {name.upper()} with nsteps = {params['nsteps']}...")
logger.info(
f"Generating parameters for {name.upper()} with nsteps = {params['nsteps']}..."
)
file_ext = f"{name}_nsteps{params['nsteps']}"
generate_sim_params(params, ICs_path, wd, simdir, file_ext, force)
logger.info(f"Generating parameters for {name.upper()} done.")
@ -324,4 +361,4 @@ if __name__ == "__main__":
raise RuntimeError("Failed.") from e
finally:
UNINDENT()
logger.info("Running convergence tests done.")
logger.info("Running convergence tests done.")

1
src/wip3m/plot_params.py
View file

@ -1 +0,0 @@
/Users/hoellinger/Library/CloudStorage/Dropbox/travail/these/science/code/PUBLIC/selfisys_public/src/selfisys/utils/plot_params.py

75
src/wip3m/plot_params.py Normal file
View file

@ -0,0 +1,75 @@
#!/usr/bin/env python3
# ----------------------------------------------------------------------
# Copyright (C) 2025 Tristan Hoellinger
# Distributed under the GNU General Public License v3.0 (GPLv3).
# See the LICENSE file in the root directory for details.
# SPDX-License-Identifier: GPL-3.0-or-later
# ----------------------------------------------------------------------
__author__ = "Tristan Hoellinger"
__version__ = "0.1.0"
__date__ = "2025"
__license__ = "GPLv3"
"""Plotting parameters and custom colormaps for the WIP-P3M package."""
# Global font sizes
GLOBAL_FS = 18
GLOBAL_FS_LARGE = 20
GLOBAL_FS_XLARGE = 22
GLOBAL_FS_SMALL = 16
GLOBAL_FS_TINY = 14
COLOUR_LIST = ["C{}".format(i) for i in range(10)]
def reset_plotting():
import matplotlib as mpl
mpl.rcParams.update(mpl.rcParamsDefault)
def setup_plotting():
"""
Configure Matplotlib plotting settings for consistent appearance.
"""
import matplotlib.pyplot as plt
import importlib.resources
with importlib.resources.open_text("wip3m", "preamble.tex") as f:
preamble = f.read()
# Dictionary with rcParams settings
rcparams = {
"font.family": "serif",
"font.size": GLOBAL_FS, # Base font size
"axes.titlesize": GLOBAL_FS_XLARGE,
"axes.labelsize": GLOBAL_FS_LARGE,
"axes.linewidth": 1.0,
"xtick.labelsize": GLOBAL_FS_SMALL,
"ytick.labelsize": GLOBAL_FS_SMALL,
"xtick.major.width": 1.2,
"ytick.major.width": 1.2,
"xtick.minor.width": 1.0,
"ytick.minor.width": 1.0,
"xtick.direction": "in",
"ytick.direction": "in",
"xtick.major.pad": 5,
"xtick.minor.pad": 5,
"ytick.major.pad": 5,
"ytick.minor.pad": 5,
"legend.fontsize": GLOBAL_FS_SMALL,
"legend.title_fontsize": GLOBAL_FS_LARGE,
"figure.titlesize": GLOBAL_FS_XLARGE,
"figure.dpi": 300,
"grid.color": "gray",
"grid.linestyle": "dotted",
"grid.linewidth": 0.6,
"lines.linewidth": 2,
"lines.markersize": 8,
"text.usetex": True,
"text.latex.preamble": preamble,
}
# Update rcParams
plt.rcParams.update(rcparams)

450
src/wip3m/plots_convergence_baseline_ts_parser.py Normal file
View file

@ -0,0 +1,450 @@
#!/usr/bin/env python3
# ----------------------------------------------------------------------
# Copyright (C) 2025 Tristan Hoellinger
# Distributed under the GNU General Public License v3.0 (GPLv3).
# See the LICENSE file in the root directory for details.
# SPDX-License-Identifier: GPL-3.0-or-later
# ----------------------------------------------------------------------
__author__ = "Tristan Hoellinger"
__version__ = "0.1.0"
__date__ = "2025"
__license__ = "GPLv3"
"""Plot results of time step convergence tests."""
# pyright: reportWildcardImportFromLibrary=false
from pathlib import Path
import numpy as np
from pysbmy.fft import read_FourierGrid
from pysbmy.field import read_field
from pysbmy.correlations import get_autocorrelation
from wip3m import *
from wip3m.tools import none_bool_str
from wip3m.plot_utils import *
from wip3m.logger import getCustomLogger, INDENT, UNINDENT
logger = getCustomLogger(__name__)
workdir = ROOT_PATH + "results/"
output_path = OUTPUT_PATH
from argparse import ArgumentParser
parser = ArgumentParser(description="Plot convergence results for WIP3M.")
parser.add_argument(
"--run_id",
type=str,
help="Run ID to use for plotting.",
)
parser.add_argument(
"--outdir_suffix",
type=str,
default="plots",
help="Suffix to append to the output directory.",
)
parser.add_argument(
"--ref",
type=str,
default="P3M1",
choices=["P3M1", "PMref"],
help="Reference field for power spectrum comparison.",
)
parser.add_argument(
"--lpt",
type=none_bool_str,
default=False,
help="Whether to compute and plot the LPT power spectrum.",
)
parser.add_argument(
"--cola",
type=none_bool_str,
default=True,
help="Whether to compute and plot the COLA power spectrum.",
)
parser.add_argument(
"--plot_fields",
type=none_bool_str,
default=False,
help="Whether to plot the fields or not.",
)
args = parser.parse_args()
run_id = args.run_id
outdir_suffix = args.outdir_suffix
ref = args.ref
lpt = args.lpt
cola = args.cola
plot_fields = args.plot_fields
if not run_id:
raise ValueError("Please provide a run ID using the --run_id argument.")
simdir = output_path + run_id + "/"
outdir = simdir + outdir_suffix + "/"
input_ss_file = simdir + "input_ss_k_grid.h5"
wd = workdir + run_id + "/"
Path(outdir).mkdir(parents=True, exist_ok=True)
if __name__ == "__main__":
logger.info(f"Plotting convergence results for run ID: {run_id}")
logger.info(f"Output directory: {outdir}")
logger.info(f"Reading simulation parameters from {wd}sim_params.txt...")
with open(wd + "sim_params.txt", "r") as f:
sim_params = eval(f.read())
L = sim_params["L"] # Box size in Mpc/h
N = sim_params["N"] # Density grid size
Np = sim_params["Np"] # Number of dark matter particles per spatial dimension
Npm = sim_params["Npm"] # PM grid size
nsteps_pmref = sim_params["nsteps_pmref"]
nsteps_pm1 = sim_params["nsteps_pm1"]
nsteps_pm2 = sim_params["nsteps_pm2"]
nsteps_cola = sim_params["nsteps_cola"]
nsteps_spm = sim_params["nsteps_spm"]
nsteps_p3m1 = sim_params["nsteps_p3m1"]
nsteps_p3m2 = sim_params["nsteps_p3m2"]
nsteps_p3m3 = sim_params["nsteps_p3m3"]
logger.info(f"Reading simulation parameters from {wd}sim_params.txt done.")
if plot_fields:
logger.info(f"Loading fields from {simdir}...")
INDENT()
slice_ijk = (N // 2, slice(None), slice(None))
DELTA_LPT = read_field(simdir + "lpt_density.h5").data[slice_ijk]
DELTA_PMREF = read_field(simdir + f"pmref_nsteps{nsteps_pmref}_final_density_pm.h5").data[
slice_ijk
]
DELTA_PM1 = read_field(simdir + f"pm1_nsteps{nsteps_pm1}_final_density_pm.h5").data[
slice_ijk
]
DELTA_PM2 = read_field(simdir + f"pm2_nsteps{nsteps_pm2}_final_density_pm.h5").data[
slice_ijk
]
DELTA_COLA = read_field(simdir + f"cola_nsteps{nsteps_cola}_final_density_cola.h5").data[
slice_ijk
]
DELTA_SPM = read_field(simdir + f"spm_nsteps{nsteps_spm}_final_density_spm.h5").data[
slice_ijk
]
DELTA_P3M1 = read_field(simdir + f"p3m1_nsteps{nsteps_p3m1}_final_density_p3m.h5").data[
slice_ijk
]
DELTA_P3M2 = read_field(simdir + f"p3m2_nsteps{nsteps_p3m2}_final_density_p3m.h5").data[
slice_ijk
]
DELTA_P3M3 = read_field(simdir + f"p3m3_nsteps{nsteps_p3m3}_final_density_p3m.h5").data[
slice_ijk
]
UNINDENT()
logger.info(f"Loading fields from {simdir} done.")
diff_pm1_pmref = DELTA_PM1 - DELTA_PMREF
diff_pm2_pmref = DELTA_PM2 - DELTA_PMREF
diff_p3m1_pmref = DELTA_P3M1 - DELTA_PMREF
diff_p3m2_pm1 = DELTA_P3M2 - DELTA_PM1
diff_p3m3_pm2 = DELTA_P3M3 - DELTA_PM2
diff_p3m1_spm = DELTA_P3M1 - DELTA_SPM
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
ncols = 6 # Max panels per row
figname = "_".join(fields)
slices_dict = {
"lpt": DELTA_LPT,
"cola": DELTA_COLA,
"pmref": DELTA_PMREF,
"pm1": DELTA_PM1,
"pm2": DELTA_PM2,
"spm": DELTA_SPM,
"p3m1": DELTA_P3M1,
"p3m2": DELTA_P3M2,
"p3m3": DELTA_P3M3,
"diff_pm1_pmref": diff_pm1_pmref,
"diff_pm2_pmref": diff_pm2_pmref,
"diff_p3m1_pmref": diff_p3m1_pmref,
"diff_p3m2_pm1": diff_p3m2_pm1,
"diff_p3m3_pm2": diff_p3m3_pm2,
"diff_p3m1_spm": diff_p3m1_spm,
}
titles_dict = {
"lpt": "LPT",
"pmref": f"PMref $n_\\mathrm{{steps}}={nsteps_pmref}$",
"pm1": f"PM1 $n_\\mathrm{{steps}}={nsteps_pm1}$",
"pm2": f"PM2 $n_\\mathrm{{steps}}={nsteps_pm2}$",
"cola": f"COLA $n_\\mathrm{{steps}}={nsteps_cola}$",
"spm": f"sPM $n_\\mathrm{{steps}}={nsteps_spm}$",
"p3m1": f"P3M1 $n_\\mathrm{{steps}}={nsteps_p3m1}$",
"p3m2": f"P3M2 $n_\\mathrm{{steps}}={nsteps_p3m2}$",
"p3m3": f"P3M3 $n_\\mathrm{{steps}}={nsteps_p3m3}$",
"diff_pm1_pmref": r"$\delta_{\rm PM1}-\delta_{\rm PMref}$",
"diff_pm2_pmref": r"$\delta_{\rm PM2}-\delta_{\rm PMref}$",
"diff_p3m1_pmref": r"$\delta_{\rm P3M1}-\delta_{\rm PMref}$",
"diff_p3m2_pm1": r"$\delta_{\rm P3M2}-\delta_{\rm PM1}$",
"diff_p3m3_pm2": r"$\delta_{\rm P3M3}-\delta_{\rm PM2}$",
"diff_p3m1_spm": r"$\delta_{\rm P3M1}-\delta_{\rm sPM}$",
}
logger.info(f"Plotting fields: {fields}...")
npanels = len(fields)
nrows = np.ceil(npanels / ncols).astype(int)
fig, axs = plt.subplots(
nrows=nrows, ncols=ncols, figsize=(3 * ncols, 4.6 * nrows), sharey=True
)
axs = axs.flatten()
ims = []
for i, key in enumerate(fields):
ax = axs[i]
data = slices_dict[key]
title = titles_dict[key]
if key.startswith("diff"):
norm = TwoSlopeNorm(
vmin=-np.log(1 + np.abs(np.min(data))),
vcenter=0,
vmax=np.log10(1 + np.abs(np.max(data))),
)
im = ax.imshow(np.sign(data) * np.log(1 + np.abs(data)), cmap="RdBu_r", norm=norm)
else:
im = ax.imshow(np.log10(2 + data), cmap=cmap)
ims.append((im, key))
ax.set_title(title, fontsize=fs_titles)
for spine in ax.spines.values():
spine.set_visible(False)
# Hide unused axes
for i in range(npanels, len(axs)):
axs[i].axis("off")
# Shared axes labels
for i, ax in enumerate(axs[:npanels]):
if i % ncols == 0:
ax.set_yticks([0, N // 2, N])
ax.set_yticklabels([f"{-L/2:.0f}", "0", f"{L/2:.0f}"], fontsize=fs)
ax.set_ylabel(r"Mpc/$h$", size=GLOBAL_FS_SMALL)
ax.set_xticks([0, N // 2, N])
ax.set_xticklabels([f"{-L/2:.0f}", "0", f"{L/2:.0f}"], fontsize=fs)
ax.set_xlabel(r"Mpc/$h$", size=GLOBAL_FS_SMALL)
# Colorbars
for ax, (im, key) in zip(axs[:npanels], ims):
divider = make_axes_locatable(ax)
cax = divider.append_axes("bottom", size="5%", pad=0.6)
cb = fig.colorbar(im, cax=cax, orientation="horizontal")
if key.startswith("diff"):
cb.set_label(
r"$\textrm{sgn}\left(\Delta\delta\right)\log_{10}(1 + |\Delta\delta|)$",
fontsize=fs,
)
else:
cb.set_label(r"$\log_{10}(2 + \delta)$", fontsize=fs)
cb.ax.tick_params(labelsize=fs)
cax.xaxis.set_ticks_position("bottom")
cax.xaxis.set_label_position("bottom")
fig.savefig(outdir + f"{figname}.png", bbox_inches="tight", dpi=300, transparent=True)
fig.savefig(outdir + f"{figname}.pdf", bbox_inches="tight", dpi=300)
plt.close(fig)
logger.info(f"Plotting fields done. Saved to {outdir}{figname}.png and {figname}.pdf")
logger.info("Computing and plotting power spectra...")
INDENT()
logger.info("Loading Fourier grid...")
G = read_FourierGrid(input_ss_file)
logger.info("Loading Fourier grid done.")
# Exit here with "0" for debugging purposes:
exit(0)
k = G.k_modes
AliasingCorr = False
logger.info("Computing power spectra...")
DELTA = read_field(simdir + "initial_density.h5")
Pk_INI, _ = get_autocorrelation(DELTA, G, AliasingCorr)
if lpt:
DELTA = read_field(simdir + "lpt_density.h5")
Pk_LPT, _ = get_autocorrelation(DELTA, G, AliasingCorr)
DELTA = read_field(simdir + f"pmref_nsteps{nsteps_pmref}_final_density_pm.h5")
Pk_PMref, _ = get_autocorrelation(DELTA, G, AliasingCorr)
DELTA = read_field(simdir + f"pm1_nsteps{nsteps_pm1}_final_density_pm.h5")
Pk_PM1, _ = get_autocorrelation(DELTA, G, AliasingCorr)
DELTA = read_field(simdir + f"pm2_nsteps{nsteps_pm2}_final_density_pm.h5")
Pk_PM2, _ = get_autocorrelation(DELTA, G, AliasingCorr)
if cola:
DELTA = read_field(simdir + f"cola_nsteps{nsteps_cola}_final_density_cola.h5")
Pk_COLA, _ = get_autocorrelation(DELTA, G, AliasingCorr)
DELTA = read_field(simdir + f"spm_nsteps{nsteps_spm}_final_density_spm.h5")
Pk_sPM, _ = get_autocorrelation(DELTA, G, AliasingCorr)
DELTA = read_field(simdir + f"p3m1_nsteps{nsteps_p3m1}_final_density_p3m.h5")
Pk_P3M1, _ = get_autocorrelation(DELTA, G, AliasingCorr)
DELTA = read_field(simdir + f"p3m2_nsteps{nsteps_p3m2}_final_density_p3m.h5")
Pk_P3M2, _ = get_autocorrelation(DELTA, G, AliasingCorr)
DELTA = read_field(simdir + f"p3m3_nsteps{nsteps_p3m3}_final_density_p3m.h5")
Pk_P3M3, _ = get_autocorrelation(DELTA, G, AliasingCorr)
logger.info("Computing power spectra done.")
if ref == "P3M1":
nsteps_reference = nsteps_p3m1
Pk_ref = Pk_P3M1
label_ref = f"Ref (P3M, $n_\\mathrm{{steps}}={nsteps_reference}$)"
elif ref == "PMref":
nsteps_reference = nsteps_pm1
Pk_ref = Pk_PMref
label_ref = f"Ref (PM, $n_\\mathrm{{steps}}={nsteps_reference}$)"
else:
raise ValueError("Invalid reference field. Choose 'P3M1' or 'PMref'.")
logger.info(f"Plotting power spectra...")
INDENT()
fig, ax = plt.subplots(figsize=(7, 4))
ax.set_xscale("log")
k = G.k_modes
kmin, kmax = k.min(), k.max()
logger.info(f"kmin = {kmin:.2e}, kmax = {kmax:.2e}")
log_pad = 0.02
log_k_min = np.log10(kmin)
log_k_max = np.log10(kmax)
log_range = log_k_max - log_k_min
xlim_min = 10 ** (log_k_min - log_pad * log_range)
xlim_max = 10 ** (log_k_max + log_pad * log_range)
plt.xlim(xlim_min, xlim_max)
# ax.set_ylim([0.2, 1.8])
dark_grey_bnd = 0.01
medium_grey_bnd = 0.05
light_grey_bnd = 0.1
ax.plot([kmin, kmax], [1, 1], color="black", linestyle="-", label=label_ref)
if lpt:
ax.plot(k, Pk_LPT / Pk_ref, label="2LPT", linestyle="--")
fields_to_plot = [
("PMref", Pk_PMref),
("PM1", Pk_PM1),
("PM2", Pk_PM2),
("sPM", Pk_sPM),
("P3M1", Pk_P3M1),
("P3M2", Pk_P3M2),
("P3M3", Pk_P3M3),
]
if cola:
fields_to_plot.append(("COLA", Pk_COLA))
# Remove the field corresponding to the reference:
fields_to_plot = [(field_name, Pk) for field_name, Pk in fields_to_plot if field_name != ref]
for field_name, Pk in fields_to_plot:
label = f"{field_name}, $n_\\mathrm{{steps}}={eval(f'nsteps_{field_name.lower()}')}$"
if field_name in ["PMref", "P3M1"]:
linestyle = "-"
zorder = 1
else:
linestyle = "--"
zorder = 2
ax.plot(k, Pk / Pk_ref, label=label, linestyle=linestyle)
ax.axhspan(1 - dark_grey_bnd, 1 + dark_grey_bnd, color="grey", alpha=0.3)
ax.axhspan(1 - medium_grey_bnd, 1 + medium_grey_bnd, color="grey", alpha=0.2)
ax.axhspan(1 - light_grey_bnd, 1 + light_grey_bnd, color="grey", alpha=0.1)
for i in range(1, len(k)):
ax.axvline(k[i], color="black", linestyle=":", linewidth=1, alpha=0.1, zorder=0)
ax.yaxis.set_major_locator(plt.MaxNLocator(6))
ax.yaxis.get_major_ticks()[0].label1.set_visible(False)
ax.set_xlabel("$k$ [$h/\\mathrm{Mpc}$]", fontsize=fs)
ax.set_ylabel("$P(k)/P_\\mathrm{ref}(k)$", fontsize=fs)
ax.tick_params(which="both", direction="in")
ax.tick_params(axis="both", which="major", labelsize=fs)
ax.tick_params(axis="both", which="minor", labelsize=fs)
# Characteristic vertical reference scales
nyquist = np.pi * N / L
nyquist_PM = np.pi * Npm / L
epsilon = 0.03 * L / Np
particle_length = 2 * epsilon
xs = 1.25 * L / Npm
xr = 4.5 * xs
particle_wavenumber = 2 * np.pi / particle_length # Too large to be shown
xs_inv = 2 * np.pi / xs
xr_inv = 2 * np.pi / xr
line1 = ax.axvline(
x=nyquist, color="black", linestyle="--", lw=2, label="Nyquist (density grid)", zorder=0
)
line1 = ax.axvline(
x=nyquist_PM, color="black", linestyle="-", lw=1, label="Nyquist (PM grid)", zorder=0
)
line3 = ax.axvline(
x=xs_inv, color="gray", linestyle="-.", lw=2, label=r"Split wavenumber $x_s$", zorder=0
)
line4 = ax.axvline(
x=xr_inv, color="gray", linestyle=":", lw=2, label=r"Short-range reach $x_r$", zorder=0
)
logger.info(f"Nyquist (density grid): {nyquist:.2f} h/Mpc")
logger.info(f"Nyquist (PM grid): {nyquist_PM:.2f} h/Mpc")
logger.info(f"Particle wavenumber: {particle_wavenumber:.2f} h/Mpc")
logger.info(f"Split wavenumber: {xs_inv:.2f} h/Mpc")
logger.info(f"Short-range reach: {xr_inv:.2f} h/Mpc")
handles, labels = plt.gca().get_legend_handles_labels()
# empty_patch = mpatches.Patch(color="none", label="")
# handles = [empty_patch, *handles]
# labels = ["", *labels]
plt.legend(
handles,
labels,
loc="upper center",
ncol=2,
bbox_to_anchor=(0.5, -0.2),
fontsize=fs,
frameon=False,
)
fig.savefig(
outdir + "power_spectra.png",
bbox_inches="tight",
dpi=300,
transparent=True,
)
fig.savefig(
outdir + "power_spectra.pdf",
bbox_inches="tight",
dpi=300,
)
plt.close(fig)
UNINDENT()
logger.info(
f"Power spectra plotted and saved to {outdir}power_spectra.png and power_spectra.pdf"
)
UNINDENT()
logger.info("Computing and plotting power spectra done.")
logger.info("All plots completed successfully.")