Add poisson weak scaling test

profiling/clusters.py

@@ -57,6 +57,21 @@ def detect_machine_name() -> str | None:
 
 
 SLURM_PRESETS: dict[str, dict] = {
+    "raven": {
+        "cpus_per_task": 1,
+        # Raven's submission filter requires total memory (not --mem-per-cpu)
+        # for shared jobs. ProfilingCase converts this per-rank value to --mem.
+        "mem_per_rank_gb": 3,
+        "max_mem_per_node_gb": 110,
+        "partition": "general",
+        "qos": "debug",
+        "account": "ipp_cpu",
+        "chdir": "./",
+        "output": "./%x.%j.out",
+        "error": "./%x.%j.err",
+        "mail_type": "none",
+        "time": "00:15:00",
+    },
     "pitagora_dcgp": {
         # "nodes": 1, # Should be set by ProfilingCase.launch()
         # "ntasks_per_node": 1,  # Should be set by ProfilingCase.launch()

profiling/examples/Poisson/cube_weak_scaling/params_poisson.py

@@ -0,0 +1,281 @@
+# -----------------------------
+# Description of the simulation
+# -----------------------------
+# Please fill in a verbal description of the simulation.
+# It will be printed at the beginning of the simulation and can be used to keep track of the different runs.
+
+name = "Poisson weak scaling on 3D cuboid"
+description = """
+Weak scaling test for Poisson equation on a 3D cuboid.
+The manufactured solution is a simple product of sines and cosines.
+Homogeneous Dirichlet boundary conditions are set in direction x.
+"""
+
+import logging
+
+from struphy import set_logging_level
+
+set_logging_level(logging.WARNING)
+
+import argparse
+
+from mpi4py import MPI
+
+# ------------------
+# Import Struphy API
+# ------------------
+from struphy import (
+    BaseUnits,
+    DerhamOptions,
+    EnvironmentOptions,
+    Simulation,
+    Time,
+    domains,
+    grids,
+)
+
+# ---------------------
+# Instance of the model
+# ---------------------
+from struphy.models import Poisson
+
+# Units
+base_units = BaseUnits()
+
+# Model instance
+model = Poisson(base_units=base_units)
+
+# --------------------------
+# Instance of the simulation
+# --------------------------
+
+# Environment options
+# `--id` distinguishes runs that share a rank count but differ in something else; the
+# profiling driver passes its launch counter (see `ProfilingJob.build_commands`).
+# Unknown flags are ignored so the driver can forward other parameters as well.
+parser = argparse.ArgumentParser()
+parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.")
+parser.add_argument(
+    "--cells-per-rank",
+    type=int,
+    default=64,
+    help="Number of cells assigned to each rank in the weak-scaling direction.",
+)
+parser.add_argument(
+    "--verify",
+    action="store_true",
+    help="Save fields, post-process them, and check the manufactured solution.",
+)
+args, _ = parser.parse_known_args()
+
+mpi_size = MPI.COMM_WORLD.Get_size()
+
+# List all variables and decide whether to save their data.
+# Full field output is only needed for the small verification mode; production weak
+# scaling runs keep output compact and rely on profiling_data.h5.
+model.em_fields.phi.save_data = args.verify
+model.em_fields.source.save_data = args.verify
+
+env = EnvironmentOptions(
+    sim_folder=f"sim_{args.id:02d}",
+    restart=False,
+)
+
+# Time stepping
+time_opts = Time()
+
+# Geometry
+Lx = 2.0 * mpi_size
+Ly = 3.0
+Lz = 4.0
+domain = domains.Cuboid(r1=Lx, l2=-Ly/2, r2=Ly/2, r3=Lz)
+
+# Fluid equilibrium (can be used as part of initial conditions)
+equil = None
+
+# Grid
+num_elements = (args.cells_per_rank * mpi_size, args.cells_per_rank, args.cells_per_rank)
+grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, False, False))
+
+# Derham options
+derham_opts = DerhamOptions(degree=(1, 2, 3), bcs=(("dirichlet", "dirichlet"), None, None))
+
+# Simulation object
+sim = Simulation(
+    model=model,
+    name=name,
+    description=description,
+    params_path=__file__,
+    env=env,
+    time_opts=time_opts,
+    domain=domain,
+    equil=equil,
+    grid=grid,
+    derham_opts=derham_opts,
+)
+
+# ------------------
+# Propagator options
+# ------------------
+
+from struphy.linear_algebra.solver import SolverParameters
+solver_params = SolverParameters(tol=1e-8, maxiter=3000, info=True, recycle=True)
+model.propagators.poisson.options = model.propagators.poisson.Options(stab_eps=0.0,
+                                                                      solver="pcg", 
+                                                                      precond=None,
+                                                                      solver_params=solver_params,
+                                                                      )
+
+# ------------------
+# Initial conditions
+# ------------------
+import numpy as np
+from struphy.initial.base import GenericPerturbation
+
+def exact_solution(x, y, z):
+    return np.sin(np.pi / Lx * x) * np.cos(2 * np.pi / Ly * y + 2 * np.pi / Lz * z)
+
+def rhs_fun(x, y, z):
+    return exact_solution(x, y, z) * ((np.pi / Lx) ** 2 + (2 * np.pi / Ly) ** 2 + (2 * np.pi / Lz) ** 2)
+
+rhs_perturbation = GenericPerturbation(rhs_fun, given_in_basis="physical")
+
+model.em_fields.source.add_perturbation(rhs_perturbation)
+
+
+if __name__ == "__main__":
+    sim.run(profiling_activated=True, one_time_step=True)
+
+    if not args.verify:
+        if sim.comm.rank == 0:
+            import os
+
+            results_dir = os.path.join(sim.env.path_out, "results")
+            os.makedirs(results_dir, exist_ok=True)
+
+            np.save(os.path.join(results_dir, "resolution.npy"), sim.grid.num_elements)
+            np.save(os.path.join(results_dir, "spline_degree.npy"), sim.derham_opts.degree)
+            np.save(os.path.join(results_dir, "mpi_size.npy"), mpi_size)
+            np.save(os.path.join(results_dir, "cells_per_rank.npy"), args.cells_per_rank)
+        raise SystemExit(0)
+
+    sim.pproc(parallel_pproc=True)
+
+    def plot_slices(num, exact, name, slice_pt_x=0, slice_pt_y=0, slice_pt_z=0):
+        from matplotlib import pyplot as plt
+
+        fig = plt.figure(figsize=(16, 12))
+
+        plt.subplot(3, 3, 1)
+        plt.pcolor(y[slice_pt_x, :, :], z[slice_pt_x, :, :], num[slice_pt_x, :, :])
+        plt.colorbar()
+        plt.xlabel("y")
+        plt.ylabel("z")
+        plt.title("{} from struphy, slice at x = {:.2f}".format(name, x[slice_pt_x, 0, 0]))
+
+        plt.subplot(3, 3, 4)
+        plt.pcolor(y[slice_pt_x, :, :], z[slice_pt_x, :, :], exact(x[slice_pt_x, :, :], y[slice_pt_x, :, :], z[slice_pt_x, :, :]))
+        plt.colorbar()
+        plt.xlabel("y")
+        plt.ylabel("z")
+        plt.title("{} exact, slice at x = {:.2f}".format(name, x[slice_pt_x, 0, 0]))
+
+        plt.subplot(3, 3, 7)
+        plt.pcolor(y[slice_pt_x, :, :], z[slice_pt_x, :, :], np.abs(num[slice_pt_x, :, :] - exact(x[slice_pt_x, :, :], y[slice_pt_x, :, :], z[slice_pt_x, :, :])))
+        plt.colorbar()
+        plt.xlabel("y")
+        plt.ylabel("z")
+        plt.title("{} error, slice at x = {:.2f}".format(name, x[slice_pt_x, 0, 0]))
+
+        plt.subplot(3, 3, 2)
+        plt.pcolor(x[:, slice_pt_y, :], z[:, slice_pt_y, :], num[:, slice_pt_y, :])
+        plt.colorbar()
+        plt.xlabel("x")
+        plt.ylabel("z")
+        plt.title("{} from struphy, slice at y = {:.2f}".format(name, y[0, slice_pt_y, 0]))
+
+        plt.subplot(3, 3, 5)
+        plt.pcolor(x[:, slice_pt_y, :], z[:, slice_pt_y, :], exact(x[:, slice_pt_y, :], y[:, slice_pt_y, :], z[:, slice_pt_y, :]))
+        plt.colorbar()
+        plt.xlabel("x")
+        plt.ylabel("z")
+        plt.title("{} exact, slice at y = {:.2f}".format(name, y[0, slice_pt_y, 0]))
+
+        plt.subplot(3, 3, 8)
+        plt.pcolor(x[:, slice_pt_y, :], z[:, slice_pt_y, :], np.abs(num[:, slice_pt_y, :] - exact(x[:, slice_pt_y, :], y[:, slice_pt_y, :], z[:, slice_pt_y, :])))
+        plt.colorbar()
+        plt.xlabel("x")
+        plt.ylabel("z")
+        plt.title("{} error, slice at y = {:.2f}".format(name, y[0, slice_pt_y, 0]))
+
+        plt.subplot(3, 3, 3)
+        plt.pcolor(x[:, :, slice_pt_z], y[:, :, slice_pt_z], num[:, :, slice_pt_z])
+        plt.colorbar()
+        plt.xlabel("x")
+        plt.ylabel("y")
+        plt.title("{} from struphy, slice at z = {:.2f}".format(name, z[0, 0, slice_pt_z]))
+
+        plt.subplot(3, 3, 6)
+        plt.pcolor(x[:, :, slice_pt_z], y[:, :, slice_pt_z], exact(x[:, :, slice_pt_z], y[:, :, slice_pt_z], z[:, :, slice_pt_z]))
+        plt.colorbar()
+        plt.xlabel("x")
+        plt.ylabel("y")
+        plt.title("{} exact, slice at z = {:.2f}".format(name, z[0, 0, slice_pt_z]))
+
+        plt.subplot(3, 3, 9)
+        plt.pcolor(x[:, :, slice_pt_z], y[:, :, slice_pt_z], np.abs(num[:, :, slice_pt_z] - exact(x[:, :, slice_pt_z], y[:, :, slice_pt_z], z[:, :, slice_pt_z])))
+        plt.colorbar()
+        plt.xlabel("x")
+        plt.ylabel("y")
+        plt.title("{} error, slice at z = {:.2f}".format(name, z[0, 0, slice_pt_z]))
+
+        return fig
+
+    if sim.comm.rank == 0:
+        sim.load_plotting_data()
+
+        Tstart = sim.t_grid[0]
+        rhs_data = sim.spline_values.em_fields.source_log
+        print(rhs_data)
+        rhs = rhs_data.data[Tstart][0]
+
+        Tend = sim.t_grid[-1]
+        phi_data = sim.spline_values.em_fields.phi_log
+        print(phi_data)
+        phi = phi_data.data[Tend][0]
+        x = sim.grids_phy[0]
+        y = sim.grids_phy[1]
+        z = sim.grids_phy[2]
+
+        slice_pt_x = x.shape[0] // 2
+        slice_pt_y = y.shape[1] // 2
+        slice_pt_z = 0
+
+        fig_rhs = plot_slices(rhs, rhs_fun, "RHS", slice_pt_x=slice_pt_x, slice_pt_y=slice_pt_y, slice_pt_z=slice_pt_z)
+        fig_phi = plot_slices(phi, exact_solution, "Phi", slice_pt_x=slice_pt_x, slice_pt_y=slice_pt_y, slice_pt_z=slice_pt_z)
+
+        rel_err_rhs = np.max(np.abs(rhs - rhs_fun(x, y, z))) / np.max(np.abs(rhs_fun(x, y, z)))
+        rel_err_phi = np.max(np.abs(phi - exact_solution(x, y, z))) / np.max(np.abs(exact_solution(x, y, z)))
+
+        print(f"Max relative error in RHS: {rel_err_rhs:.2e}")
+        print(f"Max relative error in Phi: {rel_err_phi:.2e}")
+
+        assert rel_err_rhs < 5e-3, f"The computed RHS does not match the exact RHS, max rel error = {rel_err_rhs}."
+        assert rel_err_phi < 5e-3, f"The computed solution does not match the exact solution, max rel error = {rel_err_phi}."
+
+        import os
+        # `path_out` is the run's output folder; `sim_folder` alone is a bare name
+        # resolved against the CWD. The profiling packaging picks these files up from
+        # here and uploads them as `results-run<id>`.
+        results_dir = os.path.join(sim.env.path_out, "results")
+        os.makedirs(results_dir, exist_ok=True)
+
+        np.save(os.path.join(results_dir, "rel_err_rhs.npy"), rel_err_rhs)
+        np.save(os.path.join(results_dir, "rel_err_phi.npy"), rel_err_phi)
+        np.save(os.path.join(results_dir, "resolution.npy"), sim.grid.num_elements)
+        np.save(os.path.join(results_dir, "spline_degree.npy"), sim.derham_opts.degree)
+        np.save(os.path.join(results_dir, "mpi_size.npy"), mpi_size)
+        np.save(os.path.join(results_dir, "cells_per_rank.npy"), args.cells_per_rank)
+
+        fig_rhs.savefig(os.path.join(results_dir, "rhs_slices.png"))
+        fig_phi.savefig(os.path.join(results_dir, "phi_slices.png"))

profiling/profiling_job.py

@@ -252,9 +252,20 @@ def launch(
 
             # Pick the cluster preset for this run, either from the caller's override or the default.
             if slurm_presets is not None:
-                cluster_preset = slurm_presets[cluster_name]
+                cluster_preset = dict(slurm_presets[cluster_name])
             else:
-                cluster_preset = SLURM_PRESETS[cluster_name]
+                cluster_preset = dict(SLURM_PRESETS[cluster_name])
+
+            # Some schedulers (notably Raven's submission filter) reject
+            # --mem-per-cpu for shared jobs. Let their preset express the same
+            # scaling policy while emitting the required total --mem value.
+            mem_per_rank_gb = cluster_preset.pop("mem_per_rank_gb", None)
+            max_mem_per_node_gb = cluster_preset.pop("max_mem_per_node_gb", None)
+            if mem_per_rank_gb is not None:
+                memory_gb = mem_per_rank_gb * (num_tasks // num_nodes)
+                if max_mem_per_node_gb is not None:
+                    memory_gb = min(memory_gb, max_mem_per_node_gb)
+                cluster_preset["mem"] = f"{memory_gb}G"
 
             # Build the slurm script
             script = SlurmScript(

profiling/submit_poisson_weak_scaling.py

@@ -0,0 +1,59 @@
+import argparse
+from pathlib import Path
+
+from profiling_job import ProfilingCase
+
+
+def main() -> None:
+
+    parser = argparse.ArgumentParser(
+        description=("Submit Poisson weak scaling profiling jobs and package the results."),
+    )
+    parser.add_argument(
+        "--upload",
+        action="store_true",
+        help="Upload the packaged profiling results to the profiling-data repo.",
+    )
+    parser.add_argument(
+        "--cells-per-rank",
+        type=int,
+        default=64,
+        help="Number of cells assigned to each rank in the weak-scaling direction.",
+    )
+    parser.add_argument(
+        "--ranks",
+        type=int,
+        nargs="+",
+        default=(1, 2, 4),
+        help="MPI rank counts to submit (default: 1 2 4).",
+    )
+    args = parser.parse_args()
+
+    script_dir = Path(__file__).resolve().parent
+    params_dir = script_dir / "examples" / "Poisson" / "cube_weak_scaling"
+
+    profiling_case = ProfilingCase(
+        label="poisson_cube_weak_scaling",
+        name="Poisson on cuboid weak scaling test",
+        description=(
+            "Weak scaling of the Poisson model with a manufactured solution on a "
+            "3D cuboid. The grid grows in x with the MPI rank count."
+        ),
+        physics_problem="Occurs in many plasma applications.",
+        struphy_model_used="Poisson",
+        params_source=params_dir / "params_poisson.py",
+        language="fortran",
+        compiler="GNU",
+        upload=args.upload,
+    )
+
+    param_flags = ["--cells-per-rank", str(args.cells_per_rank)]
+
+    for num_tasks in args.ranks:
+        profiling_case.launch(num_tasks, param_flags=param_flags)
+
+    profiling_case.finalize_run()
+
+
+if __name__ == "__main__":
+    main()