P3M force validated; starting timestep convergence analysis

.gitignore

@@ -54,6 +54,7 @@ __pycache__*
 ########
 data/
 examples/
+results/
 *.npy
 *.h5
 *.fits
@@ -62,3 +63,8 @@ examples/
 # Notebook #
 ############
 *.ipynb_checkpoints*
+
+##########
+# Others #
+##########
+sync.sh

src/wip3m/__init__.py

@@ -19,3 +19,5 @@ implementing the P3M algorithm in Simbelmynë.
 """
 
 from .params import *
+# pyright: reportUnsupportedDunderAll=false
+__all__ = params.__all__

src/wip3m/convergence_baseline_ts_parser.py

@@ -0,0 +1,327 @@
+#!/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"
+
+"""
+Perform time step convergence tests towards implementing P3M gravity
+"""
+
+# pyright: reportWildcardImportFromLibrary=false
+from os.path import isfile
+from pathlib import Path
+import numpy as np
+
+from pysbmy.power import PowerSpectrum
+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.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="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.")
+
+# 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()}.")
+
+args = parser.parse_args()
+
+if __name__ == "__main__":
+    try:
+        go_beyond_Nyquist_ss = True  # for the summary statistics
+        force = True
+        force_hard = True
+
+        run_id = args.run_id
+        L = args.L
+        N = args.N
+        Np = args.Np
+        Npm = args.Npm
+        n_Tiles = args.n_Tiles
+        nsteps_pmref = args.nsteps_pmref
+        nsteps_pm1 = args.nsteps_pm1
+        nsteps_pm2 = args.nsteps_pm2
+        nsteps_cola = args.nsteps_cola
+        nsteps_spm = args.nsteps_spm
+        nsteps_p3m1 = args.nsteps_p3m1
+        nsteps_p3m2 = args.nsteps_p3m2
+        nsteps_p3m3 = args.nsteps_p3m3
+        tsd_pmref = args.tsd_pmref
+        tsd_pm1 = args.tsd_pm1
+        tsd_pm2 = args.tsd_pm2
+        tsd_cola = args.tsd_cola
+        tsd_spm = args.tsd_spm
+        tsd_p3m1 = args.tsd_p3m1
+        tsd_p3m2 = args.tsd_p3m2
+        tsd_p3m3 = args.tsd_p3m3
+        logger.info("Running convergence tests...")
+        logger.info(f"Run ID: {run_id}")
+        logger.info(f"Box size: {L} Mpc/h")
+        logger.info(f"Density grid size: {N}^3")
+        logger.info(f"Particles per dimension: {Np}^3")
+        logger.info(f"PM grid size: {Npm}^3")
+        logger.info(f"Number of tiles: {n_Tiles}^3")
+        logger.info(f"Number of timesteps for PMREF: {nsteps_pmref}")
+        logger.info(f"Number of timesteps for PM1: {nsteps_pm1}")
+        logger.info(f"Number of timesteps for PM2: {nsteps_pm2}")
+        logger.info(f"Number of timesteps for COLA: {nsteps_cola}")
+        logger.info(f"Number of timesteps for sPM: {nsteps_spm}")
+        logger.info(f"Number of timesteps for P3M1: {nsteps_p3m1}")
+        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
+
+        wd = workdir + run_id + "/"
+        simdir = output_path + run_id + "/"
+        logdir = simdir + "logs/"
+        if force_hard:
+            import shutil
+            if Path(simdir).exists():
+                shutil.rmtree(simdir)
+            if Path(wd).exists():
+                shutil.rmtree(wd)
+        Path(wd).mkdir(parents=True, exist_ok=True)
+        Path(logdir).mkdir(parents=True, exist_ok=True)
+
+        input_white_noise_file = simdir + "input_white_noise.h5"
+        input_seed_phase_file = simdir + "seed"
+        ICs_path = simdir + "initial_density.h5"
+        simpath = simdir
+
+        # Path to the input matter power spectrum (generated later)
+        input_power_file = simdir + "input_power.h5"
+
+        sim_params = {
+            "L": L,
+            "N": N,
+            "Np": Np,
+            "Npm": Npm,
+            "n_Tiles": n_Tiles,
+            "RedshiftLPT": RedshiftLPT,
+            "RedshiftFCs": RedshiftFCs,
+        }
+        with open(wd + "sim_params.txt", "w") as f:
+            f.write(f"{sim_params}\n")
+
+        logger.info("Setting up simulation parameters...")
+
+        common_params = {
+            "Np": Np,
+            "N": N,
+            "L": L,
+            "corner0": corner,
+            "corner1": corner,
+            "corner2": corner,
+            "h": cosmo["h"],
+            "Omega_m": cosmo["Omega_m"],
+            "Omega_b": cosmo["Omega_b"],
+            "n_s": cosmo["n_s"],
+            "sigma8": cosmo["sigma8"],
+        }
+
+        lpt_params = common_params.copy()
+        lpt_params["method"] = "lpt"
+        lpt_params["InputPowerSpectrum"] = input_power_file
+        lpt_params["ICsMode"] = 1
+        lpt_params["InputWhiteNoise"] = input_white_noise_file
+
+        common_params_num = common_params.copy()
+        common_params_num["ai"] = ai
+        common_params_num["af"] = af
+        common_params_num["RedshiftLPT"] = RedshiftLPT
+        common_params_num["RedshiftFCs"] = RedshiftFCs
+        common_params_num["Npm"] = Npm
+        common_params_num["RunForceDiagnostic"] = False
+        common_params_num["nBinsForceDiagnostic"] = 20
+        common_params_num["nPairsForceDiagnostic"] = 3
+        common_params_num["maxTrialsForceDiagnostic"] = int(1e8)
+        common_params_num["OutputForceDiagnostic"] = simdir + "force_diagnostic.txt"
+
+        pmref_params = common_params_num.copy()
+        pmref_params["method"] = "pm"
+        pmref_params["TimeStepDistribution"] = tsd_pmref
+        pmref_params["nsteps"] = nsteps_pmref
+
+        pm1_params = common_params_num.copy()
+        pm1_params["method"] = "pm"
+        pm1_params["TimeStepDistribution"] = tsd_pm1
+        pm1_params["nsteps"] = nsteps_pm1
+
+        pm2_params = common_params_num.copy()
+        pm2_params["method"] = "pm"
+        pm2_params["TimeStepDistribution"] = tsd_pm2
+        pm2_params["nsteps"] = nsteps_pm2
+
+        cola_params = common_params_num.copy()
+        cola_params["method"] = "cola"
+        cola_params["TimeStepDistribution"] = tsd_cola
+        cola_params["nsteps"] = nsteps_cola
+
+        spm_params = common_params_num.copy()
+        spm_params["method"] = "spm"
+        spm_params["EvolutionMode"] = 5
+        spm_params["TimeStepDistribution"] = tsd_spm
+        spm_params["nsteps"] = nsteps_spm
+        spm_params["n_Tiles"] = n_Tiles
+
+        p3m1_params = common_params_num.copy()
+        p3m1_params["method"] = "p3m"
+        p3m1_params["EvolutionMode"] = 4
+        p3m1_params["TimeStepDistribution"] = tsd_p3m1
+        p3m1_params["nsteps"] = nsteps_p3m1
+        p3m1_params["n_Tiles"] = n_Tiles
+
+        p3m2_params = common_params_num.copy()
+        p3m2_params["method"] = "p3m"
+        p3m2_params["EvolutionMode"] = 4
+        p3m2_params["TimeStepDistribution"] = tsd_p3m2
+        p3m2_params["nsteps"] = nsteps_p3m2
+        p3m2_params["n_Tiles"] = n_Tiles
+
+        p3m3_params = common_params_num.copy()
+        p3m3_params["method"] = "p3m"
+        p3m3_params["EvolutionMode"] = 4
+        p3m3_params["TimeStepDistribution"] = tsd_p3m3
+        p3m3_params["nsteps"] = nsteps_p3m3
+        p3m3_params["n_Tiles"] = n_Tiles
+
+        logger.info("Setting up simulation parameters done.")
+
+        logger.info("Generating simulation parameters...")
+        INDENT()
+        reset_plotting()  # Default style for Simbelmynë
+        generate_sim_params(lpt_params, ICs_path, wd, simdir, None, force)
+
+        all_sim_params = [
+            ("pmref", pmref_params),
+            ("pm1", pm1_params),
+            ("pm2", pm2_params),
+            ("cola", cola_params),
+            ("spm", spm_params),
+            ("p3m1", p3m1_params),
+            ("p3m2", p3m2_params),
+            ("p3m3", p3m3_params),
+        ]
+        for name, params in all_sim_params:
+            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.")
+
+        UNINDENT()
+        logger.info("Generating simulation parameters done.")
+
+        setup_plotting()  # Reset plotting style for this project
+
+        logger.info("Generating white noise field...")
+        generate_white_noise_Field(
+            L=L,
+            size=N,
+            corner=corner,
+            seedphase=BASELINE_SEEDPHASE,
+            fname_whitenoise=input_white_noise_file,
+            seedname_whitenoise=input_seed_phase_file,
+            force_phase=force,
+        )
+        logger.info("Generating white noise field done.")
+
+        # If cosmo["WhichSpectrum"] == "class", then classy is required.
+        if not isfile(input_power_file) or force:
+            logger.info("Generating input power spectrum...")
+            Pk = PowerSpectrum(L, L, L, N, N, N, cosmo_small_to_full_dict(cosmo))
+            Pk.write(input_power_file)
+            logger.info("Generating input power spectrum done.")
+
+        logger.info("Generating Fourier grid...")
+        # k grid used to compute the final overdensity power spectrum
+        Pinit = 100
+        trim_threshold = 100  # Merge bins until this minimum number of modes per bin is reached
+        log_kmin = np.log10(2 * np.pi / (np.sqrt(3) * L))  # Minimum non-zero k in h/Mpc
+        if go_beyond_Nyquist_ss:
+            k_max_ss = get_k_max(L, N)
+        else:
+            k_max_ss = get_k_max(L, N) / np.sqrt(3)  # 1D Nyquist frequency
+        Pbins_left_bnds = np.logspace(log_kmin, np.log10(k_max_ss), Pinit + 1, dtype=np.float32)
+        Pbins_left_bnds = Pbins_left_bnds[:-1]
+        input_ss_file = simdir + "input_ss_k_grid.h5"
+        Gk = FourierGrid(
+            L,
+            L,
+            L,
+            N,
+            N,
+            N,
+            k_modes=Pbins_left_bnds,
+            kmax=k_max_ss,
+            trim_bins=True,
+            trim_threshold=trim_threshold,
+        )
+        Gk.write(input_ss_file)
+        logger.info("Generating Fourier grid done.")
+
+        logger.info("Running simulations...")
+        INDENT()
+        logger.info("Running LPT simulation...")
+        run_simulation("lpt", lpt_params, wd, logdir)
+        logger.info("Running LPT simulation done.")
+
+        for name, parameters in all_sim_params:
+            logger.info(f"Running {name.upper()} simulation...")
+            run_simulation(name, parameters, wd, logdir)
+            logger.info(f"Running {name.upper()} simulation done.")
+        UNINDENT()
+        logger.info("All simulations done.")
+
+    except OSError as e:
+        logger.error("Directory or file access error: %s", str(e))
+        raise
+    except Exception as e:
+        logger.critical("An unexpected error occurred: %s", str(e))
+        raise RuntimeError("Failed.") from e
+    finally:
+        UNINDENT()
+        logger.info("Running convergence tests done.")

src/wip3m/params.py

@@ -15,15 +15,16 @@ __license__ = "GPLv3"
 
 import os
 from pathlib import Path
+from typing import cast
 import numpy as np
 
 WHICH_SPECTRUM = "eh"  # available options are "eh" and "class"
 
-# Load global paths from environment variables
-ROOT_PATH = os.getenv("WIP3M_ROOT_PATH")
+ROOT_PATH = cast(str, os.getenv("WIP3M_ROOT_PATH"))
 if ROOT_PATH is None:
     raise EnvironmentError("Please set the 'WIP3M_ROOT_PATH' environment variable.")
-OUTPUT_PATH = os.getenv("WIP3M_OUTPUT_PATH")
+
+OUTPUT_PATH = cast(str, os.getenv("WIP3M_OUTPUT_PATH"))
 if OUTPUT_PATH is None:
     raise EnvironmentError("Please set the 'WIP3M_OUTPUT_PATH' environment variable.")
 
@@ -107,3 +108,6 @@ def cosmo_small_to_full_dict(cosmo_min):
         "WhichSpectrum": cosmo_min["WhichSpectrum"],
     }
     return cosmo_full
+
+# pyright: reportUnsupportedDunderAll=false
+__all__ = [k for k in globals() if not k.startswith("_")]

src/wip3m/plot_params.py

@@ -1,80 +0,0 @@
-#!/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.
-
-This module provides custom Matplotlib settings, formatter classes, and
-colormaps used for visualising results in the SelfiSys project.
-"""
-
-
-# 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)

src/wip3m/plot_params.py

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

src/wip3m/plot_utils.py

@@ -19,8 +19,9 @@ import matplotlib.patches as mpatches
 from matplotlib.colors import TwoSlopeNorm
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 import cmocean.cm as cm
+import plotly.graph_objects as go
 
-from wip3m.plot_params import *
+from wip3m.plot_params import *  # type: ignore
 
 # Configure global plotting settings
 setup_plotting()
@@ -28,9 +29,26 @@ setup_plotting()
 fs = GLOBAL_FS_SMALL
 fs_titles = GLOBAL_FS_LARGE
 cols = COLOUR_LIST
-cmap = cm.thermal
+cmap = cm.thermal  # type: ignore
 
-def plotly_3d(field, size=128, L=None, colormap="RdYlBu", limits="max"):
+def matplotlib_to_plotly(cmap, n=255):
+    """Convert a matplotlib colormap to a Plotly colorscale."""
+    colorscale = []
+    for i in range(n):
+        norm = i / (n - 1)
+        r, g, b, _ = cmap(norm)
+        colorscale.append([norm, f'rgb({int(r*255)}, {int(g*255)}, {int(b*255)})'])
+    return colorscale
+
+thermal_plotly = matplotlib_to_plotly(cmap)
+
+def plotly_3d(
+    field: np.ndarray,
+    size: int = 128,
+    L: float = None,  # type: ignore
+    colormap: str | list = "RdYlBu",
+    limits: str = "max",
+) -> go.Figure:
     """
     Create an interactive 3D plot of volume slices using Plotly.
 
@@ -53,9 +71,6 @@ def plotly_3d(field, size=128, L=None, colormap="RdYlBu", limits="max"):
     go.Figure
         Plotly figure object.
     """
-    import numpy as np
-    import plotly.graph_objects as go
-
     volume = field.T
     rows, cols = volume[0].shape
 
@@ -172,17 +187,6 @@ def plotly_3d(field, size=128, L=None, colormap="RdYlBu", limits="max"):
     fig.update_layout(**layout_config)
     return fig
 
-def matplotlib_to_plotly(cmap, n=255):
-    """Convert a matplotlib colormap to a Plotly colorscale."""
-    colorscale = []
-    for i in range(n):
-        norm = i / (n - 1)
-        r, g, b, _ = cmap(norm)
-        colorscale.append([norm, f'rgb({int(r*255)}, {int(g*255)}, {int(b*255)})'])
-    return colorscale
-
-thermal_plotly = matplotlib_to_plotly(cm.thermal)
-
 def clear_large_plot(fig):
     """
     Clear a figure to free up memory.
@@ -334,7 +338,7 @@ def plot_force_distance_comparison(rr, ff, ll, L, Np, Npm, ss=None, a=None, titl
     epsilon = 0.03 * L / Np
     xs = 1.25 * L / Npm
     xr = 4.5 * xs
-    line1 = ax.axvline(x=nyquist, color="black", linestyle="-", lw=1, label="Nyquist")
+    line1 = ax.axvline(x=nyquist, color="black", linestyle="-", lw=1, label="Nyquist (PM grid)")
     line2 = ax.axvline(x=2*epsilon, color="gray", linestyle="--", lw=2, label=r"Particle length $2\epsilon$")
     line3 = ax.axvline(x=xs, color="gray", linestyle="-.", lw=2, label=r"Split scale $x_s$")
     line4 = ax.axvline(x=xr, color="gray", linestyle=":", lw=2, label=r"Short-range reach $x_r$")

src/wip3m/tools.py

@@ -422,3 +422,19 @@ def generate_white_noise_Field(
             raise RuntimeError("generate_white_noise_Field failed.") from e
         finally:
             gc.collect()
+
+def run_simulation(name, params, wd, logdir):
+    from io import BytesIO
+    from wip3m.low_level import stderr_redirector
+    from pysbmy import pySbmy
+
+    file_ext = f"{name}_nsteps{params['nsteps']}" if params.get("nsteps") is not None else None
+    method = params["method"]
+    path = wd + file_ext + "_" if file_ext else wd
+    sbmy_path = joinstrs([path, "example_", method, ".sbmy"])
+    log_path = joinstrs([logdir, file_ext, method, ".txt"])
+    
+    f = BytesIO()
+    with stderr_redirector(f):
+        pySbmy(sbmy_path, log_path)
+    f.close()

submit/convergence_baseline_ts.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# ----------------------------------------------------------------------
+# 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
+# Date: 2025
+# License: GPLv3
+
+eval "$(conda shell.bash hook)"
+conda activate p3m
+
+
+export OMP_NUM_THREADS=8
+python $WIP3M_ROOT_PATH"src/wip3m/convergence_baseline_ts_parser.py" \
+    --run_id ts1/ \
+    --L 64 \
+    --N 64 \
+    --Np 64 \
+    --Npm 3600 \
+    --n_Tiles 32 \
+    --nsteps_pmref 200 \
+    --nsteps_pm1 100 \
+    --nsteps_pm2 20 \
+    --nsteps_cola 10 \
+    --nsteps_spm 200 \
+    --nsteps_p3m1 200 \
+    --nsteps_p3m2 100 \
+    --nsteps_p3m3 20 \
+    --tsd_pmref 0 \
+    --tsd_pm1 0 \
+    --tsd_pm2 0 \
+    --tsd_cola 0 \
+    --tsd_spm 0 \
+    --tsd_p3m1 0 \
+    --tsd_p3m2 0 \
+    --tsd_p3m3 0
+
+exit 0