From a24546d1496f31ba7c4e9176f8ecdb779daf8d69 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 6 Aug 2026 16:21:48 +0200 Subject: [PATCH 01/32] Added petsc and petsc4py to dependencies --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 97c62e82b..41fb7e5e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,6 +57,10 @@ file = "LICENSE" mpi = [ "mpi4py<=4.1.1", ] +petsc = [ + "petsc", + "petsc4py", +] phys = [ "gvec>=1.1.0, <=1.4.1", "desc-opt<=0.17.1", @@ -106,6 +110,7 @@ all = [ "struphy[mpi]", "struphy[doc]", "struphy[likwid]", + "struphy[petsc]", ] [project.urls] From 2606cf7e75efd7e0e977a1fd0502f3764053e2aa Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 6 Aug 2026 16:23:10 +0200 Subject: [PATCH 02/32] Adde PETSc solver --- src/struphy/linear_algebra/petsc_solver.py | 148 ++++++++++++++++++ .../linear_algebra/petsc_solver_example.py | 69 ++++++++ .../linear_algebra/tests/test_petsc_solver.py | 66 ++++++++ 3 files changed, 283 insertions(+) create mode 100644 src/struphy/linear_algebra/petsc_solver.py create mode 100644 src/struphy/linear_algebra/petsc_solver_example.py create mode 100644 src/struphy/linear_algebra/tests/test_petsc_solver.py diff --git a/src/struphy/linear_algebra/petsc_solver.py b/src/struphy/linear_algebra/petsc_solver.py new file mode 100644 index 000000000..98b4637ee --- /dev/null +++ b/src/struphy/linear_algebra/petsc_solver.py @@ -0,0 +1,148 @@ +import logging + +from feectools.linalg.basic import InverseLinearOperator, Vector +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix +from feectools.linalg.topetsc import mat_topetsc, vec_topetsc +from feectools.linalg.utilities import petsc_to_psydac + +logger = logging.getLogger("struphy") + + +class PETScSolver(InverseLinearOperator): + """(Approximate) inverse of an assembled operator, computed via a PETSc ``KSP`` Krylov solver. + + ``A`` is converted to a ``PETSc.Mat`` via :func:`feectools.linalg.topetsc.mat_topetsc` + and the right-hand side is converted to a ``PETSc.Vec`` via + :func:`feectools.linalg.topetsc.vec_topetsc`; the solve itself is delegated to + ``petsc4py.PETSc.KSP``. Requires the optional ``petsc4py`` dependency + (``pip install struphy[petsc]``). + + Parameters + ---------- + A : feectools.linalg.stencil.StencilMatrix | feectools.linalg.block.BlockLinearOperator + Left-hand-side matrix of the linear system. Only assembled operators can be + converted to a ``PETSc.Mat``, see :func:`feectools.linalg.topetsc.mat_topetsc`. + + x0 : feectools.linalg.basic.Vector, default=None + Kept for interface compatibility with the other + :class:`~feectools.linalg.basic.InverseLinearOperator` subclasses; unused by PETSc's KSP. + + tol : float, default=1e-6 + Relative tolerance, passed to ``KSP.setTolerances(rtol=tol)``. + + maxiter : int, default=1000 + Maximum number of KSP iterations. + + verbose : bool, default=False + If True, log convergence information after each solve. + + recycle : bool, default=False + Kept for interface compatibility; unused by PETSc's KSP. + + ksp_type : str, default="cg" + PETSc Krylov solver type, see ``petsc4py.PETSc.KSP.Type``. + + pc_type : str, default="none" + PETSc preconditioner type, see ``petsc4py.PETSc.PC.Type``. + """ + + def __init__( + self, + A, + *, + x0=None, + tol=1e-6, + maxiter=1000, + verbose=False, + recycle=False, + ksp_type="cg", + pc_type="none", + ): + assert isinstance(A, (StencilMatrix, BlockLinearOperator)), ( + f"PETScSolver only supports assembled operators (StencilMatrix or BlockLinearOperator), got {type(A)}." + ) + + self._options = { + "x0": x0, + "tol": tol, + "maxiter": maxiter, + "verbose": verbose, + "recycle": recycle, + "ksp_type": ksp_type, + "pc_type": pc_type, + } + + super().__init__(A, **self._options) + + self._info = None + self._ksp = None + # operator for which self._ksp's PETSc.Mat was last built, used to avoid + # re-assembling the matrix on every solve() call when `linop` is unchanged + self._ksp_linop = None + + def _get_ksp(self): + from petsc4py import PETSc + + A = self._A + if self._ksp is None or self._ksp_linop is not A: + gmat = mat_topetsc(A) + + if self._ksp is None: + self._ksp = PETSc.KSP().create(comm=gmat.getComm()) + + self._ksp.setType(self._options["ksp_type"]) + self._ksp.getPC().setType(self._options["pc_type"]) + self._ksp.setTolerances(rtol=self._options["tol"], max_it=self._options["maxiter"]) + self._ksp.setOperators(gmat) + self._ksp.setFromOptions() + + self._ksp_linop = A + + return self._ksp + + def solve(self, b, out=None): + """Solve ``A x = b`` using a PETSc KSP Krylov solver. + + Parameters + ---------- + b : feectools.linalg.basic.Vector + Right-hand-side vector of the linear system. + + out : feectools.linalg.basic.Vector | None + The output vector, or None (optional). + + Returns + ------- + x : feectools.linalg.basic.Vector + Numerical solution of the linear system. Convergence info is available + via :meth:`get_info`. + """ + assert isinstance(b, Vector) + assert b.space is self._domain + + ksp = self._get_ksp() + + gvec_b = vec_topetsc(b) + gvec_x = gvec_b.duplicate() + + ksp.solve(gvec_b, gvec_x) + + out = petsc_to_psydac(gvec_x, self._codomain, out=out) + + self._info = { + "niter": ksp.getIterationNumber(), + "success": ksp.getConvergedReason() > 0, + "res_norm": ksp.getResidualNorm(), + } + + if self._options["verbose"]: + logger.info(f"PETSc KSP solver info: {self._info}") + + gvec_b.destroy() + gvec_x.destroy() + + return out + + def dot(self, b, out=None): + return self.solve(b, out=out) diff --git a/src/struphy/linear_algebra/petsc_solver_example.py b/src/struphy/linear_algebra/petsc_solver_example.py new file mode 100644 index 000000000..2b1ce6426 --- /dev/null +++ b/src/struphy/linear_algebra/petsc_solver_example.py @@ -0,0 +1,69 @@ +"""Example: solve a struphy mass-matrix system with :class:`~struphy.linear_algebra.petsc_solver.PETScSolver`. + +Builds the ``H1`` mass matrix ``M0`` of a small 3D Derham complex on a cuboid domain, +manufactures a right-hand side from a known exact solution, and solves ``M0 x = b`` +with a PETSc KSP (CG + Jacobi preconditioner), comparing against feectools' native +preconditioned CG solver. + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_solver_example +""" + +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.solvers import inverse + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.fields_background.equils import HomogenSlab +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.petsc_solver import PETScSolver +from struphy.topology.grids import TensorProductGrid + + +def main(): + comm = MPI.COMM_WORLD + + # domain, equilibrium and Derham complex + domain = Cuboid() + equil = HomogenSlab(n0=2.0) + equil.domain = domain + + grid = TensorProductGrid(num_elements=[8, 8, 8]) + derham_opts = DerhamOptions(degree=[2, 2, 2]) + derham = Derham(grid, derham_opts, comm=comm, domain=domain) + + # weighted mass operators -- M0.matrix is the assembled StencilMatrix on the H1 space + mass_ops = WeightedMassOperators(derham, domain, eq_mhd=equil) + M0 = mass_ops.M0.matrix + + # manufacture a right-hand side from a known exact solution + xe = M0.domain.zeros() + xe[:] = xp.random.random(xe[:].shape) + xe.update_ghost_regions() + b = M0.dot(xe) + + # reference solve with feectools' preconditioned CG + pc = M0.diagonal(inverse=True) + cg_solver = inverse(M0, "pcg", pc=pc, tol=1e-12, maxiter=2000, verbose=False, recycle=False) + x_cg = cg_solver.solve(b) + + # solve the same system with PETSc's CG + Jacobi preconditioner + petsc_solver = PETScSolver(M0, tol=1e-12, maxiter=2000, ksp_type="cg", pc_type="jacobi") + x_petsc = petsc_solver.solve(b) + + error_vs_exact = xp.linalg.norm((x_petsc - xe).toarray()) + error_vs_cg = xp.linalg.norm((x_petsc - x_cg).toarray()) + + if comm.Get_rank() == 0: + print(f"PETSc KSP info: {petsc_solver.get_info()}") + print(f"||x_petsc - x_exact|| = {error_vs_exact:.3e}") + print(f"||x_petsc - x_cg|| = {error_vs_cg:.3e}") + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/tests/test_petsc_solver.py b/src/struphy/linear_algebra/tests/test_petsc_solver.py new file mode 100644 index 000000000..176b28fd4 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_solver.py @@ -0,0 +1,66 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.solvers import inverse +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace + +from struphy.linear_algebra.petsc_solver import PETScSolver + + +def _define_tridiagonal_spd_system(n, p): + """Banded, symmetric positive-definite StencilMatrix with 2p+1 diagonals, and a random exact solution.""" + domain_decomposition = DomainDecomposition([n - p], [False], comm=MPI.COMM_WORLD) + cart = CartDecomposition(domain_decomposition, [n], [xp.array([0])], [xp.array([n - 1])], [p], [1]) + V = StencilVectorSpace(cart) + s = V.starts[0] + e = V.ends[0] + + A = StencilMatrix(V, V) + A[:, -p:0] = -1.0 + A[:, 0:1] = 2 * p + A[:, 1 : p + 1] = -1.0 + A.remove_spurious_entries() + + xe = StencilVector(V) + xe[s : e + 1] = xp.random.random(e + 1 - s) + + return V, A, xe + + +@pytest.mark.parametrize("n", [8, 15]) +@pytest.mark.parametrize("p", [1, 2]) +def test_petsc_solver_matches_cg(n, p): + """PETScSolver must solve Ax=b to the same accuracy as feectools' native CG solver.""" + xp.random.seed(n * p) + + _, A, xe = _define_tridiagonal_spd_system(n, p) + + b = A @ xe + + ref_solver = inverse(A, "cg", tol=1e-13, maxiter=2000, verbose=False, recycle=False) + x_ref = ref_solver.solve(b) + + petsc_solver = PETScSolver(A, tol=1e-13, maxiter=2000, ksp_type="cg", pc_type="none") + x_petsc = petsc_solver.solve(b) + + info = petsc_solver.get_info() + assert info["success"] + + error_vs_exact = xp.linalg.norm((x_petsc - xe).toarray()) + assert error_vs_exact < 1e-8 + + error_vs_ref = xp.linalg.norm((x_petsc - x_ref).toarray()) + assert error_vs_ref < 1e-6 + + # re-solving with an unchanged operator (KSP/Mat cache reused) must still be correct + b2 = A @ x_petsc + x_petsc2 = petsc_solver.solve(b2) + assert xp.linalg.norm((x_petsc2 - x_petsc).toarray()) < 1e-8 + + +if __name__ == "__main__": + test_petsc_solver_matches_cg(15, 2) From 2377ecf791af5ec2af489ca316a667995ac3c327 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 7 Aug 2026 07:50:12 +0200 Subject: [PATCH 03/32] Added a wrapper of inverse() with if solver == petsc logic --- src/struphy/feec/mass.py | 3 +- src/struphy/io/options.py | 2 +- src/struphy/linear_algebra/solver.py | 46 +++++++++++++++++++ .../tests/test_petsc_l2_projector.py | 46 +++++++++++++++++++ 4 files changed, 94 insertions(+), 3 deletions(-) create mode 100644 src/struphy/linear_algebra/tests/test_petsc_l2_projector.py diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 908ab0762..cde1e1605 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -11,7 +11,6 @@ from feectools.fem.vector import VectorFemSpace from feectools.linalg.basic import IdentityOperator, InverseLinearOperator, LinearOperator, Vector from feectools.linalg.block import BlockLinearOperator, BlockVector -from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilDiagonalMatrix, StencilMatrix, StencilVector from struphy import equils @@ -22,7 +21,7 @@ from struphy.fields_background.base import MHDequilibrium from struphy.geometry.base import Domain from struphy.io.options import LiteralOptions -from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import SolverParameters, inverse from struphy.polar.basic import PolarVector from struphy.polar.linear_operators import PolarExtractionOperator from struphy.utils.docstring_converter import auto_convert_docstring, info diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index a66d67c43..ab1dc2cfa 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -67,7 +67,7 @@ class LiteralOptions: GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None] # solvers - OptsSymmSolver = Literal["pcg", "cg"] + OptsSymmSolver = Literal["pcg", "cg", "petsc"] OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] OptsSaddlePointSolver = Literal["uzawa"] diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index f01ff07d9..aa2e1bef0 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -5,6 +5,52 @@ logger = logging.getLogger("struphy") +# kwargs accepted by struphy.linear_algebra.petsc_solver.PETScSolver.__init__ +_PETSC_SOLVER_KWARGS = ("x0", "tol", "maxiter", "verbose", "recycle", "ksp_type", "pc_type") + + +def inverse(A, solver: str, **kwargs): + """Create an (approximate) inverse of ``A``. + + Thin wrapper around :func:`feectools.linalg.solvers.inverse` that additionally + supports ``solver="petsc"``, dispatching to + :class:`~struphy.linear_algebra.petsc_solver.PETScSolver`. For all other solver + names this simply delegates to the feectools implementation. + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Left-hand-side matrix of the linear system. For ``solver="petsc"``, ``A`` + must be (or expose via ``A.matrix``) an assembled + ``StencilMatrix``/``BlockLinearOperator``, see + :class:`~struphy.linear_algebra.petsc_solver.PETScSolver`. + + solver : str + Preferred iterative solver, one of feectools' options ('cg', 'pcg', + 'bicg', 'bicgstab', 'pbicgstab', 'minres', 'lsmr', 'gmres') or 'petsc'. + + Returns + ------- + obj : feectools.linalg.basic.InverseLinearOperator + A linear operator acting as the (approximate) inverse of A. + """ + if solver == "petsc": + from struphy.linear_algebra.petsc_solver import PETScSolver + + if kwargs.get("pc") is not None: + logger.debug("PETScSolver ignores the feectools 'pc' preconditioner; use 'pc_type' instead.") + + matrix = getattr(A, "matrix", A) + petsc_kwargs = {k: v for k, v in kwargs.items() if k in _PETSC_SOLVER_KWARGS} + petsc_kwargs.setdefault("ksp_type", "cg") + petsc_kwargs.setdefault("pc_type", "jacobi") + + return PETScSolver(matrix, **petsc_kwargs) + + from feectools.linalg.solvers import inverse as feectools_inverse + + return feectools_inverse(A, solver, **kwargs) + @dataclass class SolverParameters: diff --git a/src/struphy/linear_algebra/tests/test_petsc_l2_projector.py b/src/struphy/linear_algebra/tests/test_petsc_l2_projector.py new file mode 100644 index 000000000..f0d492d2b --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_l2_projector.py @@ -0,0 +1,46 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.mass import L2Projector, WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.fields_background.equils import HomogenSlab +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.topology.grids import TensorProductGrid + + +@pytest.mark.parametrize("space_id", ["H1", "L2"]) +def test_l2_projector_petsc_matches_pcg(space_id): + """L2Projector(solver_name="petsc") must match L2Projector(solver_name="pcg") for a real mass matrix.""" + comm = MPI.COMM_WORLD + + domain = Cuboid() + equil = HomogenSlab(n0=2.0) + equil.domain = domain + + grid = TensorProductGrid(num_elements=[8, 8, 8]) + derham_opts = DerhamOptions(degree=[2, 2, 2]) + derham = Derham(grid, derham_opts, comm=comm, domain=domain) + + mass_ops = WeightedMassOperators(derham, domain, eq_mhd=equil) + + def rhs(e1, e2, e3): + return xp.sin(2 * xp.pi * e1) * xp.cos(2 * xp.pi * e2) * xp.cos(2 * xp.pi * e3) + + proj_pcg = L2Projector(space_id, mass_ops, solver_name="pcg") + proj_petsc = L2Projector(space_id, mass_ops, solver_name="petsc") + + b = proj_pcg.get_dofs(rhs, apply_bc=True) + + x_pcg = proj_pcg.solve(b) + x_petsc = proj_petsc.solve(b) + + assert xp.linalg.norm((x_petsc - x_pcg).toarray()) < 1e-6 + + +if __name__ == "__main__": + test_l2_projector_petsc_matches_pcg("H1") From ae4f9a582496fa54c91eb5e67dbdde189e1fba1d Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 7 Aug 2026 09:15:15 +0200 Subject: [PATCH 04/32] Improved topetsc in feectools --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 74c88399a..11ebcd418 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 74c88399a7adb72a07c5bdc30d6e352f10e851e4 +Subproject commit 11ebcd418dee98f6521d08a543040751e0bf1f86 From 94e3457f536501c5bd266bbb421a5053ef7fb4d1 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 7 Aug 2026 12:10:36 +0200 Subject: [PATCH 05/32] Update feectools --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 11ebcd418..7c371c39e 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 11ebcd418dee98f6521d08a543040751e0bf1f86 +Subproject commit 7c371c39e4e7d81c6745ff2c07ba0bc1ba85c7ff From 049e508782e8318bf4b9a12be0a263e6adb77e1f Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 7 Aug 2026 12:14:25 +0200 Subject: [PATCH 06/32] Added src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py --- .../petsc_solver_ill_conditioned_example.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) create mode 100644 src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py diff --git a/src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py b/src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py new file mode 100644 index 000000000..6aa8a4c37 --- /dev/null +++ b/src/struphy/linear_algebra/petsc_solver_ill_conditioned_example.py @@ -0,0 +1,93 @@ +"""Example: PETSc beats feectools' native solver on an ill-conditioned SPD system. + +Mass-matrix solves (see ``petsc_solver_example.py``) are *not* where PETSc helps: +they are already well-conditioned and feectools' diagonal-preconditioned CG converges +in a couple of iterations, so per-solve overhead dominates and makes PETSc slower there. + +Where PETSc *does* win is on badly-conditioned elliptic systems, where its algebraic +multigrid preconditioner (``pc_type="gamg"``) keeps the iteration count roughly constant +while plain (or diagonally-preconditioned) CG needs an iteration count that grows like +``sqrt(condition number)``. + +This example builds the standard 1D discrete Laplacian (tridiagonal, -1/2/-1), whose +condition number scales like ``O(n^2)`` in the number of unknowns ``n``, and compares: + +- feectools' plain, unpreconditioned CG +- :class:`~struphy.linear_algebra.petsc_solver.PETScSolver` with ``ksp_type="cg", pc_type="gamg"`` + +For this to show a genuine win (not just fewer iterations but less wall time), the +per-solve vector conversion (:func:`feectools.linalg.topetsc.vec_topetsc` / +:func:`~feectools.linalg.utilities.petsc_to_psydac`) needs to be vectorized rather than +looping in pure Python over every DOF -- that conversion cost otherwise swamps any +iteration-count savings. There's a rough sweet spot: below a few thousand unknowns +GAMG's one-time setup cost dominates and plain CG wins; above it, PETSc's roughly +constant iteration count pulls ahead, increasingly so as ``n`` grows (and, as a bonus, +plain unpreconditioned CG's accuracy degrades from floating-point error accumulation +once it needs tens of thousands of iterations, while PETSc stays accurate). + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_solver_ill_conditioned_example +""" + +import time + +import cunumpy as xp +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.solvers import inverse +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace + +from struphy.linear_algebra.petsc_solver import PETScSolver + + +def build_1d_laplacian(n, comm): + """Standard 1D discrete Laplacian (tridiagonal, -1, 2, -1); condition number ~ O(n^2).""" + p = 1 + dd = DomainDecomposition([n - p], [False], comm=comm) + cart = CartDecomposition(dd, [n], [xp.array([0])], [xp.array([n - 1])], [p], [1]) + V = StencilVectorSpace(cart) + s = V.starts[0] + e = V.ends[0] + + A = StencilMatrix(V, V) + A[:, -1:0] = -1.0 + A[:, 0:1] = 2.0 + A[:, 1:2] = -1.0 + A.remove_spurious_entries() + + xe = StencilVector(V) + xe[s : e + 1] = xp.random.random(e + 1 - s) + + return V, A, xe + + +def main(n=50_000, tol=1e-8, maxiter=200_000): + comm = MPI.COMM_WORLD + + _, A, xe = build_1d_laplacian(n, comm) + b = A @ xe + + t0 = time.perf_counter() + cg_solver = inverse(A, "cg", tol=tol, maxiter=maxiter, verbose=False, recycle=False) + x_cg = cg_solver.solve(b) + t_cg = time.perf_counter() - t0 + + t0 = time.perf_counter() + petsc_solver = PETScSolver(A, tol=tol, maxiter=maxiter, ksp_type="cg", pc_type="gamg") + x_petsc = petsc_solver.solve(b) + t_petsc = time.perf_counter() - t0 + + if comm.Get_rank() == 0: + print(f"n={n} dofs (1D Laplacian, condition number ~ O(n^2))") + print(f" cg (unpreconditioned) : {t_cg * 1e3:9.2f} ms, niter={cg_solver.get_info()['niter']}") + print(f" petsc (cg + gamg) : {t_petsc * 1e3:9.2f} ms, niter={petsc_solver.get_info()['niter']}") + print(f" speedup: {t_cg / t_petsc:.2f}x") + print(f" ||x_cg - x_exact|| = {xp.linalg.norm((x_cg - xe).toarray()):.3e}") + print(f" ||x_petsc - x_exact|| = {xp.linalg.norm((x_petsc - xe).toarray()):.3e}") + + +if __name__ == "__main__": + main() From a49eba533085733a063ca2949f740539a08697cd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 7 Aug 2026 19:11:51 +0200 Subject: [PATCH 07/32] PoissonSolve/ImplicitDiffusion (bad performance) --- src/struphy/linear_algebra/petsc_solver.py | 219 ++++++++++++++++-- src/struphy/linear_algebra/solver.py | 11 +- .../test_petsc_directional_derivative.py | 68 ++++++ .../tests/test_petsc_poisson_solve.py | 85 +++++++ src/struphy/propagators/implicit_diffusion.py | 6 +- 5 files changed, 363 insertions(+), 26 deletions(-) create mode 100644 src/struphy/linear_algebra/tests/test_petsc_directional_derivative.py create mode 100644 src/struphy/linear_algebra/tests/test_petsc_poisson_solve.py diff --git a/src/struphy/linear_algebra/petsc_solver.py b/src/struphy/linear_algebra/petsc_solver.py index 98b4637ee..091eec211 100644 --- a/src/struphy/linear_algebra/petsc_solver.py +++ b/src/struphy/linear_algebra/petsc_solver.py @@ -1,28 +1,214 @@ import logging -from feectools.linalg.basic import InverseLinearOperator, Vector -from feectools.linalg.block import BlockLinearOperator +import cunumpy as xp +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ( + ComposedLinearOperator, + IdentityOperator, + InverseLinearOperator, + LinearOperator, + ScaledLinearOperator, + SumLinearOperator, + Vector, +) +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace from feectools.linalg.stencil import StencilMatrix -from feectools.linalg.topetsc import mat_topetsc, vec_topetsc +from feectools.linalg.topetsc import get_npts_local, mat_topetsc, vec_topetsc from feectools.linalg.utilities import petsc_to_psydac logger = logging.getLogger("struphy") -class PETScSolver(InverseLinearOperator): - """(Approximate) inverse of an assembled operator, computed via a PETSc ``KSP`` Krylov solver. +def _directional_derivative_to_stencil_matrix(op): + """ Build a :class:`~feectools.linalg.stencil.StencilMatrix` equivalent to a + (matrix-free) :class:`~feectools.feec.derivatives.DirectionalDerivativeOperator`, so it can + be handed to :func:`feectools.linalg.topetsc.mat_topetsc`. + + ``DirectionalDerivativeOperator.dot`` computes, along its differentiation axis + ``diffdir`` (identity along every other axis): + + - ``out = v[..., k+1, ...] - v[..., k, ...]`` if not negative, not transposed + - ``out = v[..., k, ...] - v[..., k+1, ...]`` if negative, not transposed + - ``out = v[..., k-1, ...] - v[..., k, ...]`` if not negative, transposed + - ``out = v[..., k, ...] - v[..., k-1, ...]`` if negative, transposed + + i.e. a plain two-point (identity, shift-by-one) stencil. + + Note + ---- + Only verified for a *periodic* differentiation axis. For a non-periodic axis under a + parallel (MPI-comm-attached) space -- which is how struphy always builds its Derham + complex, even with a single rank -- this construction (and feectools' own + ``DirectionalDerivativeOperator.tokronstencil().tostencil()``) was found to disagree with + the operator's actual ``.dot()`` at the two boundary planes along that axis. The root + cause was not identified; rather than risk silently wrong results, this case raises + ``NotImplementedError``. + """ + assert isinstance(op, DirectionalDerivativeOperator) + + V = op.domain + W = op.codomain + ndim = V.ndim + diffdir = op._diffdir + negative = op._negative + transposed = op._transposed + + if not V.periods[diffdir]: + raise NotImplementedError( + "PETScSolver cannot (yet) assemble a DirectionalDerivativeOperator along a " + f"non-periodic axis (diffdir={diffdir}, periods={V.periods}) of a parallel " + "(MPI-comm-attached) space: this was found to disagree with the operator's actual " + "action at the domain boundary, for a reason not yet root-caused. Only fully " + "periodic operators (e.g. derham.grad on a fully periodic domain) are supported." + ) + + M = StencilMatrix(V, W) + + def off(o): + return slice(o, o + 1) + + rows = tuple(slice(None) for _ in range(ndim)) + identity_key = tuple(off(0) for _ in range(ndim)) + + shift = -1 if transposed else 1 + shifted_key = tuple(off(shift) if d == diffdir else off(0) for d in range(ndim)) + + if negative: + M[rows + identity_key] = 1.0 + M[rows + shifted_key] = -1.0 + else: + M[rows + identity_key] = -1.0 + M[rows + shifted_key] = 1.0 + + M.remove_spurious_entries() + return M + + +def _assemble_leaf_operator(A): + """ Return an operator equivalent to `A` that is directly convertible via + :func:`feectools.linalg.topetsc.mat_topetsc` (i.e. a ``StencilMatrix`` or a + ``BlockLinearOperator`` whose blocks are all ``StencilMatrix``), replacing any + ``DirectionalDerivativeOperator`` (block or bare) by its assembled equivalent. + """ + if isinstance(A, DirectionalDerivativeOperator): + return _directional_derivative_to_stencil_matrix(A) + + if isinstance(A, BlockLinearOperator): + out = BlockLinearOperator(A.domain, A.codomain) + for i, j in A.nonzero_block_indices: + block = A[i, j] + out[i, j] = _directional_derivative_to_stencil_matrix(block) if isinstance( + block, DirectionalDerivativeOperator + ) else block + return out + + return A + + +def _comm_of(space): + """ MPI communicator of a StencilVectorSpace/BlockVectorSpace, matching mat_topetsc's convention. """ + if isinstance(space, BlockVectorSpace): + return space.spaces[0].cart.global_comm + return space.cart.global_comm + + +def _identity_petsc_mat(space): + """ Build a PETSc.Mat representing the identity operator on `space`. """ + from petsc4py import PETSc + + comm = _comm_of(space) + localsize = int(xp.sum(xp.prod(get_npts_local(space), axis=1))) + globalsize = space.dimension - ``A`` is converted to a ``PETSc.Mat`` via :func:`feectools.linalg.topetsc.mat_topetsc` - and the right-hand side is converted to a ``PETSc.Vec`` via - :func:`feectools.linalg.topetsc.vec_topetsc`; the solve itself is delegated to - ``petsc4py.PETSc.KSP``. Requires the optional ``petsc4py`` dependency + gmat = PETSc.Mat().create(comm=comm) + gmat.setSizes(size=((localsize, globalsize), (localsize, globalsize))) + gmat.setType("mpiaij" if comm else "seqaij") + gmat.setUp() + + ones = space.zeros() + ones._data[:] = 1.0 + gmat.setDiagonal(vec_topetsc(ones)) + gmat.assemble() + + return gmat + + +def _assemble_petsc_matrix(A): + """ Recursively assemble a ``PETSc.Mat`` for a (possibly composite) feectools + ``LinearOperator``, by converting every assembled leaf via + :func:`feectools.linalg.topetsc.mat_topetsc` and combining the pieces with PETSc's own + matrix algebra (``matMult`` for composition, ``axpy`` for sums, ``scale`` for scalar + multiples). This lets algebraic preconditioners (jacobi, gamg, ...) work on operators such + as ``grad.T @ M @ grad`` that are not themselves a ``StencilMatrix``/``BlockLinearOperator``. + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Operator to assemble. Supported: ``StencilMatrix``, ``BlockLinearOperator``, any operator + exposing an assembled ``.matrix`` (e.g. ``WeightedMassOperator``), ``IdentityOperator``, + ``ScaledLinearOperator``, ``SumLinearOperator`` and ``ComposedLinearOperator`` built out of + the above (as produced e.g. by ``derham.grad.T @ mass_ops.M1 @ derham.grad``). + + Returns + ------- + gmat : PETSc.Mat + """ + if isinstance(A, (StencilMatrix, BlockLinearOperator, DirectionalDerivativeOperator)): + return mat_topetsc(_assemble_leaf_operator(A)) + + matrix = getattr(A, "matrix", None) + if isinstance(matrix, (StencilMatrix, BlockLinearOperator, DirectionalDerivativeOperator)): + return mat_topetsc(_assemble_leaf_operator(matrix)) + + if isinstance(A, IdentityOperator): + return _identity_petsc_mat(A.domain) + + if isinstance(A, ScaledLinearOperator): + gmat = _assemble_petsc_matrix(A.operator) + gmat.scale(A.scalar) + return gmat + + if isinstance(A, ComposedLinearOperator): + from petsc4py import PETSc + + mats = [_assemble_petsc_matrix(m) for m in A.multiplicants] + gmat = mats[0] + for m in mats[1:]: + gmat = gmat.matMult(m) + return gmat + + if isinstance(A, SumLinearOperator): + from petsc4py import PETSc + + mats = [_assemble_petsc_matrix(a) for a in A.addends] + gmat = mats[0].copy() + for m in mats[1:]: + gmat.axpy(1.0, m, structure=PETSc.Mat.Structure.DIFFERENT_NONZERO_PATTERN) + return gmat + + raise NotImplementedError( + f"PETScSolver cannot assemble a PETSc matrix for operator of type {type(A)}. " + "Supported: StencilMatrix, BlockLinearOperator, operators exposing an assembled " + "'.matrix', IdentityOperator, and Scaled/Sum/Composed combinations thereof." + ) + + +class PETScSolver(InverseLinearOperator): + """(Approximate) inverse of a feectools ``LinearOperator``, computed via a PETSc ``KSP`` + Krylov solver. + + ``A`` is assembled into a ``PETSc.Mat`` (see :func:`_assemble_petsc_matrix` -- this also + handles composite operators such as ``grad.T @ M @ grad``, not just plain + ``StencilMatrix``/``BlockLinearOperator``) and the right-hand side is converted to a + ``PETSc.Vec`` via :func:`feectools.linalg.topetsc.vec_topetsc`; the solve itself is delegated + to ``petsc4py.PETSc.KSP``. Requires the optional ``petsc4py`` dependency (``pip install struphy[petsc]``). Parameters ---------- - A : feectools.linalg.stencil.StencilMatrix | feectools.linalg.block.BlockLinearOperator - Left-hand-side matrix of the linear system. Only assembled operators can be - converted to a ``PETSc.Mat``, see :func:`feectools.linalg.topetsc.mat_topetsc`. + A : feectools.linalg.basic.LinearOperator + Left-hand-side matrix of the linear system, see :func:`_assemble_petsc_matrix` for the + supported operator types. x0 : feectools.linalg.basic.Vector, default=None Kept for interface compatibility with the other @@ -44,7 +230,8 @@ class PETScSolver(InverseLinearOperator): PETSc Krylov solver type, see ``petsc4py.PETSc.KSP.Type``. pc_type : str, default="none" - PETSc preconditioner type, see ``petsc4py.PETSc.PC.Type``. + PETSc preconditioner type, see ``petsc4py.PETSc.PC.Type``. E.g. ``"gamg"`` (algebraic + multigrid) for large, ill-conditioned elliptic systems. """ def __init__( @@ -59,9 +246,7 @@ def __init__( ksp_type="cg", pc_type="none", ): - assert isinstance(A, (StencilMatrix, BlockLinearOperator)), ( - f"PETScSolver only supports assembled operators (StencilMatrix or BlockLinearOperator), got {type(A)}." - ) + assert isinstance(A, LinearOperator), f"PETScSolver requires a LinearOperator, got {type(A)}." self._options = { "x0": x0, @@ -86,7 +271,7 @@ def _get_ksp(self): A = self._A if self._ksp is None or self._ksp_linop is not A: - gmat = mat_topetsc(A) + gmat = _assemble_petsc_matrix(A) if self._ksp is None: self._ksp = PETSc.KSP().create(comm=gmat.getComm()) diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index aa2e1bef0..2ba503562 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -20,10 +20,10 @@ def inverse(A, solver: str, **kwargs): Parameters ---------- A : feectools.linalg.basic.LinearOperator - Left-hand-side matrix of the linear system. For ``solver="petsc"``, ``A`` - must be (or expose via ``A.matrix``) an assembled - ``StencilMatrix``/``BlockLinearOperator``, see - :class:`~struphy.linear_algebra.petsc_solver.PETScSolver`. + Left-hand-side matrix of the linear system. For ``solver="petsc"``, see + :func:`struphy.linear_algebra.petsc_solver._assemble_petsc_matrix` for the + supported operator types -- this includes plain assembled matrices as well as + composite operators such as ``grad.T @ M @ grad``. solver : str Preferred iterative solver, one of feectools' options ('cg', 'pcg', @@ -40,12 +40,11 @@ def inverse(A, solver: str, **kwargs): if kwargs.get("pc") is not None: logger.debug("PETScSolver ignores the feectools 'pc' preconditioner; use 'pc_type' instead.") - matrix = getattr(A, "matrix", A) petsc_kwargs = {k: v for k, v in kwargs.items() if k in _PETSC_SOLVER_KWARGS} petsc_kwargs.setdefault("ksp_type", "cg") petsc_kwargs.setdefault("pc_type", "jacobi") - return PETScSolver(matrix, **petsc_kwargs) + return PETScSolver(A, **petsc_kwargs) from feectools.linalg.solvers import inverse as feectools_inverse diff --git a/src/struphy/linear_algebra/tests/test_petsc_directional_derivative.py b/src/struphy/linear_algebra/tests/test_petsc_directional_derivative.py new file mode 100644 index 000000000..0e81bb08d --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_directional_derivative.py @@ -0,0 +1,68 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.psydac_derham import Derham +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.petsc_solver import _directional_derivative_to_stencil_matrix +from struphy.topology.grids import TensorProductGrid + + +def _random_fill(v, seed): + from feectools.linalg.block import BlockVector + + xp.random.seed(seed) + if isinstance(v, BlockVector): + for b in v.blocks: + b._data[:] = xp.random.random(b._data.shape) + else: + v._data[:] = xp.random.random(v._data.shape) + v.update_ghost_regions() + return v + + +def test_directional_derivative_matches_grad_on_periodic_domain(): + """The StencilMatrix built by _directional_derivative_to_stencil_matrix must reproduce + every block of derham.grad and derham.grad.T exactly, on a fully periodic domain. + """ + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + grid = TensorProductGrid(num_elements=[6, 6, 6]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + + for op, name in [(derham.grad, "grad"), (derham.grad.T, "grad.T")]: + for i, j in op.nonzero_block_indices: + block = op[i, j] + M = _directional_derivative_to_stencil_matrix(block) + + v = _random_fill(block.domain.zeros(), seed=100 + i + 10 * j + rank) + err = xp.linalg.norm((block.dot(v) - M.dot(v)).toarray()) + assert err < 1e-12, f"{name}[{i},{j}] mismatch: err={err:.3e}" + + +def test_directional_derivative_raises_on_nonperiodic_axis(): + """A non-periodic differentiation axis must raise NotImplementedError rather than + silently produce wrong results (see the docstring of + _directional_derivative_to_stencil_matrix for the unresolved root cause). + """ + comm = MPI.COMM_WORLD + + grid = TensorProductGrid(num_elements=[6, 6, 6]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(("free", "free"), None, None)) + derham = Derham(grid, derham_opts, comm=comm) + + op = derham.grad.T[0, 0] + assert op.domain.periods[op._diffdir] is False + + with pytest.raises(NotImplementedError): + _directional_derivative_to_stencil_matrix(op) + + +if __name__ == "__main__": + test_directional_derivative_matches_grad_on_periodic_domain() + test_directional_derivative_raises_on_nonperiodic_axis() diff --git a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve.py b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve.py new file mode 100644 index 000000000..2e48d6338 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve.py @@ -0,0 +1,85 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.propagators.poisson_solve import PoissonSolve +from struphy.topology.grids import TensorProductGrid + + +def test_poisson_solve_petsc_matches_pcg(): + """PoissonSolve(solver="petsc") must match PoissonSolve(solver="pcg") on a fully periodic domain. + + PETScSolver can only assemble derham.grad on a *periodic* differentiation axis (see + struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix); this is why + the domain here is fully periodic rather than using Dirichlet/Neumann boundaries. + """ + comm = MPI.COMM_WORLD + + domain = Cuboid() + + grid = TensorProductGrid(num_elements=[10, 10, 10]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + + mass_ops = WeightedMassOperators(derham, domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + def sol_xyz(x, y, z): + return xp.sin(2 * xp.pi * x) * xp.cos(2 * xp.pi * y) + + def rho_xyz(x, y, z): + return sol_xyz(x, y, z) * ((2 * xp.pi) ** 2 + (2 * xp.pi) ** 2) + + def rho_pulled(e1, e2, e3): + return domain.pull(rho_xyz, e1, e2, e3, kind="0", squeeze_out=False) + + def run(solver_name): + solver_params = SolverParameters(tol=1e-11, maxiter=3000, info=False, recycle=False) + + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + + prop = PoissonSolve(rho=rho_pulled) + prop.variables.phi = phi + prop.options = prop.Options( + stab_eps=1e-12, + solver=solver_name, + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) + prop.allocate() + prop(1.0) + return phi + + phi_pcg = run("pcg") + phi_petsc = run("petsc") + + e1 = xp.linspace(0.0, 1.0, 20) + e2 = xp.linspace(0.0, 1.0, 20) + e3 = xp.array([0.5]) + + val_pcg = domain.push(phi_pcg.spline, e1, e2, e3, kind="0") + val_petsc = domain.push(phi_petsc.spline, e1, e2, e3, kind="0") + + x, y, z = domain(e1, e2, e3) + analytic = sol_xyz(x, y, z) + + assert xp.max(xp.abs(val_petsc - analytic)) < 1e-2 + assert xp.max(xp.abs(val_petsc - val_pcg)) < 1e-6 + + +if __name__ == "__main__": + test_poisson_solve_petsc_matches_pcg() diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index fb09cd349..607d55f0d 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -1,16 +1,16 @@ import logging +from collections.abc import Callable from dataclasses import dataclass -from typing import Callable, Literal +from typing import Literal import cunumpy as xp from feectools.linalg.basic import IdentityOperator -from feectools.linalg.solvers import inverse from feectools.linalg.stencil import StencilVector from line_profiler import profile from struphy.feec.mass import L2Projector, WeightedMassOperator from struphy.io.options import LiteralOptions, OptionsBase -from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import SolverParameters, inverse from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.pic.accumulation.filter import FilterParameters from struphy.pic.accumulation.particles_to_grid import AccumulatorVector, ParticlesToGrid From 757113c6ac0ec4cd668ab3ec106444e0df7795a8 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 7 Aug 2026 19:38:01 +0200 Subject: [PATCH 08/32] cache petsc solver in implicit diffusion --- src/struphy/propagators/implicit_diffusion.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 607d55f0d..797042737 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -370,6 +370,11 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: self._rhs2 = phi.space.zeros() self._tmp_src = phi.space.zeros() + # cache for the lhs operator (see __call__): avoids rebuilding (and re-assembling, for + # e.g. solver="petsc") a fresh operator every call when sig_1 (hence dt) is unchanged + self._lhs_op = None + self._lhs_op_sig_1 = None + @property def sources(self) -> list[StencilVector | FEECVariable | AccumulatorVector]: """ @@ -457,8 +462,13 @@ def __call__(self, dt): proj = L2Projector("H1", self.mass_ops) self.diagnostic.spline.vector = proj.solve(rhs) - # compute lhs - self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + # compute lhs (reuse the cached operator when sig_1 is unchanged, e.g. constant dt -- + # this lets InverseLinearOperator subclasses that cache on `linop` identity, such as + # PETScSolver, avoid re-assembling the operator on every call) + if self._lhs_op is None or sig_1 != self._lhs_op_sig_1: + self._lhs_op = sig_1 * self._stab_mat + self._diffusion_op + self._lhs_op_sig_1 = sig_1 + self._solver.linop = self._lhs_op # solve out = self._solver.solve(rhs, out=self._tmp) From 7bf21796d0fee4fdab6764ef42c3c5f224f50f70 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 08:25:13 +0200 Subject: [PATCH 09/32] Added struphy.log* to gitignore --- .gitignore | 1 + src/struphy/linear_algebra/petsc_solver.py | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index 8e8873977..a155451da 100644 --- a/.gitignore +++ b/.gitignore @@ -112,3 +112,4 @@ pyvenv.cfg *profile_output*.txt *kernels.txt struphy.log +struphy.log* diff --git a/src/struphy/linear_algebra/petsc_solver.py b/src/struphy/linear_algebra/petsc_solver.py index 091eec211..559c2b783 100644 --- a/src/struphy/linear_algebra/petsc_solver.py +++ b/src/struphy/linear_algebra/petsc_solver.py @@ -215,7 +215,11 @@ class PETScSolver(InverseLinearOperator): :class:`~feectools.linalg.basic.InverseLinearOperator` subclasses; unused by PETSc's KSP. tol : float, default=1e-6 - Relative tolerance, passed to ``KSP.setTolerances(rtol=tol)``. + Relative tolerance, passed to ``KSP.setTolerances(rtol=tol)``. Note this differs from + feectools' own solvers, whose ``tol`` is an *absolute* tolerance on the residual norm -- + for a poorly-scaled system (e.g. a right-hand side far from order 1) the two are not + directly comparable; see git history for a reverted attempt to unify them via + ``atol``, which caused severe slowdowns/inaccuracy for such systems. maxiter : int, default=1000 Maximum number of KSP iterations. From 499c9b8e4c2626f2d0f80f418ea58ad7f20467ee Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 08:35:33 +0200 Subject: [PATCH 10/32] Added profiling/examples/VlasovAmpereOneSpecies/ --- .../bump_on/params_bump_on_pcg.py | 151 ++++++++++++++++++ .../bump_on/params_bump_on_petsc.py | 151 ++++++++++++++++++ .../params_strong_Landau_damping_pcg.py | 145 +++++++++++++++++ .../params_strong_Landau_damping_petsc.py | 145 +++++++++++++++++ .../two_stream/params_two_stream_pcg.py | 150 +++++++++++++++++ .../two_stream/params_two_stream_petsc.py | 150 +++++++++++++++++ .../params_weak_Landau_damping_pcg.py | 146 +++++++++++++++++ .../params_weak_Landau_damping_petsc.py | 146 +++++++++++++++++ 8 files changed, 1184 insertions(+) create mode 100644 profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py create mode 100644 profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py create mode 100644 profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py create mode 100644 profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py create mode 100644 profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py create mode 100644 profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py create mode 100644 profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py create mode 100644 profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py new file mode 100644 index 000000000..98cbb9fec --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py @@ -0,0 +1,151 @@ +# ----------------------------- +# 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. + +description = """ +Nonlinear bump-on-tail instability: A kinetic plasma instability test case for the Vlasov-Ampère model. +This test features a "bump" (localized excess) in the high-velocity tail of the electron velocity distribution. +The bump-on-tail configuration is unstable to the generation of Langmuir waves, leading to energy transfer +from the hot electron population to the growing wave field. This nonlinear process exhibits complex dynamics +including mode coupling and particle trapping in the wave potential. +This benchmark validates the particle-in-cell treatment of velocity-space instabilities and wave-particle interactions. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_pcg") + +# Time stepping +time_opts = Time(dt = 0.1, Tend = 60.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 62.83) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot_1 = BinningPlot(slice="e1_v1", n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) #for initial velocity distribution +binplot_2 = BinningPlot(slice = "v1", n_bins = 128, ranges = (-10.0,10.0)) # for progression of velocity and space distribution +saving_params = SavingParameters(binning_plots=(binplot_1, binplot_2)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4, + ) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +maxwellian_1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(1/10, None), u1 = (-4.5, None), vth1 = (0.5, None)) +background = maxwellian_1 + maxwellian_2 +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.05,), ls = (1,)) +init1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) +init2 = maxwellians.Maxwellian3D(n = (1/10, perturbation), u1 = (-4.5, None), vth1 = (0.5, None)) +init = init1 + init2 +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py new file mode 100644 index 000000000..8a086c8a0 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py @@ -0,0 +1,151 @@ +# ----------------------------- +# 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. + +description = """ +Nonlinear bump-on-tail instability: A kinetic plasma instability test case for the Vlasov-Ampère model. +This test features a "bump" (localized excess) in the high-velocity tail of the electron velocity distribution. +The bump-on-tail configuration is unstable to the generation of Langmuir waves, leading to energy transfer +from the hot electron population to the growing wave field. This nonlinear process exhibits complex dynamics +including mode coupling and particle trapping in the wave potential. +This benchmark validates the particle-in-cell treatment of velocity-space instabilities and wave-particle interactions. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_petsc") + +# Time stepping +time_opts = Time(dt = 0.1, Tend = 60.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 62.83) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot_1 = BinningPlot(slice="e1_v1", n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) #for initial velocity distribution +binplot_2 = BinningPlot(slice = "v1", n_bins = 128, ranges = (-10.0,10.0)) # for progression of velocity and space distribution +saving_params = SavingParameters(binning_plots=(binplot_1, binplot_2)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4, + ) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +maxwellian_1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(1/10, None), u1 = (-4.5, None), vth1 = (0.5, None)) +background = maxwellian_1 + maxwellian_2 +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.05,), ls = (1,)) +init1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) +init2 = maxwellians.Maxwellian3D(n = (1/10, perturbation), u1 = (-4.5, None), vth1 = (0.5, None)) +init = init1 + init2 +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py new file mode 100644 index 000000000..461ece59f --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py @@ -0,0 +1,145 @@ +# ----------------------------- +# 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. + +description = """ +Strong (nonlinear) Landau damping: A nonlinear test case for the VlasovAmpereOneSpecies model. +This test involves a large amplitude electrostatic perturbation in a uniform, collisionless plasma. +Unlike weak Landau damping, the nonlinear regime exhibits trapping of particles in the potential wells +of the self-consistent electric field, leading to vortex formation and complex phase space structures. +This benchmark tests the ability of the particle-in-cell method to capture nonlinear kinetic effects +and validates the long-term stability and accuracy of the Vlasov-Ampère discretization. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_pcg") + +# Time stepping +time_opts = Time(dt = 0.05, Tend = 75.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 12.56) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions() + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4,) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.5,), ls = (1,)) +init = maxwellians.Maxwellian3D(n = (1.0, perturbation)) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py new file mode 100644 index 000000000..65bf94554 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py @@ -0,0 +1,145 @@ +# ----------------------------- +# 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. + +description = """ +Strong (nonlinear) Landau damping: A nonlinear test case for the VlasovAmpereOneSpecies model. +This test involves a large amplitude electrostatic perturbation in a uniform, collisionless plasma. +Unlike weak Landau damping, the nonlinear regime exhibits trapping of particles in the potential wells +of the self-consistent electric field, leading to vortex formation and complex phase space structures. +This benchmark tests the ability of the particle-in-cell method to capture nonlinear kinetic effects +and validates the long-term stability and accuracy of the Vlasov-Ampère discretization. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_petsc") + +# Time stepping +time_opts = Time(dt = 0.05, Tend = 75.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 12.56) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions() + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4,) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.5,), ls = (1,)) +init = maxwellians.Maxwellian3D(n = (1.0, perturbation)) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py new file mode 100644 index 000000000..2b159a304 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py @@ -0,0 +1,150 @@ +# ----------------------------- +# 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. + +description = """ +Nonlinear two-stream instability: A fundamental kinetic test case for the Vlasov-Ampère model. +This test involves two counter-streaming particle populations with a small perturbation that triggers +the two-stream instability. The instability leads to the formation of electron acoustic waves and +subsequent nonlinear effects including particle trapping and energy exchange between modes. +This benchmark validates the numerical treatment of beam-plasma interactions and tests the accuracy +of the particle-in-cell method in capturing mode coupling and energy transfer phenomena. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_pcg") + +# Time stepping +time_opts = Time(dt = 0.1, Tend = 50.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 31.42) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4, + ) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +maxwellian_1 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (3.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (-3.0, None)) +background = maxwellian_1 + maxwellian_2 +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) +init1 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (3.0, None)) +init2 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (-3.0, None)) +init = init1 + init2 +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py new file mode 100644 index 000000000..7eb1e5248 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py @@ -0,0 +1,150 @@ +# ----------------------------- +# 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. + +description = """ +Nonlinear two-stream instability: A fundamental kinetic test case for the Vlasov-Ampère model. +This test involves two counter-streaming particle populations with a small perturbation that triggers +the two-stream instability. The instability leads to the formation of electron acoustic waves and +subsequent nonlinear effects including particle trapping and energy exchange between modes. +This benchmark validates the numerical treatment of beam-plasma interactions and tests the accuracy +of the particle-in-cell method in capturing mode coupling and energy transfer phenomena. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_petsc") + +# Time stepping +time_opts = Time(dt = 0.1, Tend = 50.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 31.42) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4, + ) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +maxwellian_1 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (3.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (-3.0, None)) +background = maxwellian_1 + maxwellian_2 +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) +init1 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (3.0, None)) +init2 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (-3.0, None)) +init = init1 + init2 +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py new file mode 100644 index 000000000..ffafdae60 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py @@ -0,0 +1,146 @@ +# ----------------------------- +# 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. + +description = """ +Weak Landau damping: A linear test case for the VlasovAmpereOneSpecies model. +This test involves a small amplitude electrostatic perturbation in a uniform, collisionless plasma. +The perturbation is damped due to phase mixing effects (Landau damping) as particles interact with +the self-consistent electric field. This benchmark validates the numerical discretization of the +Vlasov-Ampère system and the accuracy of particle-in-cell methods. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_pcg") + +# Time stepping +time_opts = Time(dt = 0.05, Tend = 20.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 12.56) # r1 -> pi * 4 -> k = 0.5 + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42) +weights_params = WeightsParameters(control_variate= True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4, + ) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) +init = maxwellians.Maxwellian3D(n = (1.0,perturbation)) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py new file mode 100644 index 000000000..3ee39461f --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py @@ -0,0 +1,146 @@ +# ----------------------------- +# 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. + +description = """ +Weak Landau damping: A linear test case for the VlasovAmpereOneSpecies model. +This test involves a small amplitude electrostatic perturbation in a uniform, collisionless plasma. +The perturbation is damped due to phase mixing effects (Landau damping) as particles interact with +the self-consistent electric field. This benchmark validates the numerical discretization of the +Vlasov-Ampère system and the accuracy of particle-in-cell methods. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + FieldsBackground, + Simulation, + Time, + domains, + equils, + grids, + perturbations, +) + +# For particles: +from struphy import ( + BinningPlot, + BoundaryParameters, + KernelDensityPlot, + LoadingParameters, + WeightsParameters, + SortingParameters, + SavingParameters, + maxwellians, +) + +# --------------------- +# Instance of the model +# --------------------- + +from struphy.models import VlasovAmpereOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions(sim_folder="sim_data_petsc") + +# Time stepping +time_opts = Time(dt = 0.05, Tend = 20.0, split_algo = "LieTrotter") + +# Geometry +domain = domains.Cuboid(r1 = 12.56) # r1 -> pi * 4 -> k = 0.5 + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42) +weights_params = WeightsParameters(control_variate= True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) + +binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) +saving_params = SavingParameters(binning_plots=(binplot,)) + +model.kinetic_ions.set_markers(loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize = 0.4, + ) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_eta.options = model.propagators.push_eta.Options() +if model.with_B0: + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +# For kinetic species the background is mandatory. +# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. +# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). + +# Background for kinetic species +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +# Perturbations for (some) kinetic species +perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) +init = maxwellians.Maxwellian3D(n = (1.0,perturbation)) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + sim.run() \ No newline at end of file From 71eb643d5f9c3a54f3e3c271a6179a35f8d453c2 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 08:50:18 +0200 Subject: [PATCH 11/32] Updated examples --- .../bump_on/params_bump_on_pcg.py | 27 +++++------ .../bump_on/params_bump_on_petsc.py | 27 +++++------ .../params_strong_Landau_damping_pcg.py | 47 ++++++++++++------- .../params_strong_Landau_damping_petsc.py | 47 ++++++++++++------- .../two_stream/params_two_stream_pcg.py | 27 +++++------ .../two_stream/params_two_stream_petsc.py | 27 +++++------ .../params_weak_Landau_damping_pcg.py | 27 +++++------ .../params_weak_Landau_damping_petsc.py | 27 +++++------ 8 files changed, 134 insertions(+), 122 deletions(-) diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py index 98cbb9fec..22feb5aa2 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py +++ b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -17,35 +18,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -64,7 +61,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_pcg") +env = EnvironmentOptions(sim_folder="sim_data_pcg", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) # Time stepping time_opts = Time(dt = 0.1, Tend = 60.0, split_algo = "LieTrotter") @@ -100,7 +97,7 @@ loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot_1 = BinningPlot(slice="e1_v1", n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) #for initial velocity distribution binplot_2 = BinningPlot(slice = "v1", n_bins = 128, ranges = (-10.0,10.0)) # for progression of velocity and space distribution diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py index 8a086c8a0..d98daea4e 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py +++ b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -17,35 +18,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -64,7 +61,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_petsc") +env = EnvironmentOptions(sim_folder="sim_data_petsc", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) # Time stepping time_opts = Time(dt = 0.1, Tend = 60.0, split_algo = "LieTrotter") @@ -100,7 +97,7 @@ loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot_1 = BinningPlot(slice="e1_v1", n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) #for initial velocity distribution binplot_2 = BinningPlot(slice = "v1", n_bins = 128, ranges = (-10.0,10.0)) # for progression of velocity and space distribution diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py index 461ece59f..37c554e48 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py +++ b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -17,35 +18,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -63,8 +60,22 @@ # Instance of the simulation # -------------------------- +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +args, _ = parser.parse_known_args() + # Environment options -env = EnvironmentOptions(sim_folder="sim_data_pcg") +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, +) # Time stepping time_opts = Time(dt = 0.05, Tend = 75.0, split_algo = "LieTrotter") @@ -100,7 +111,7 @@ loading_params = LoadingParameters(ppc=20, seed=42) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) saving_params = SavingParameters(binning_plots=(binplot,)) @@ -142,4 +153,8 @@ model.kinetic_ions.var.add_initial_condition(init) if __name__ == "__main__": - sim.run() \ No newline at end of file + # one_time_step=True isolates the initial Poisson solve (the only part of this model that + # solver= affects -- the field then evolves via VlasovAmpereCoupling, unrelated to + # PETScSolver) from the ~1500 identical transport steps a full Tend=75.0 run would otherwise + # dilute the comparison with. + sim.run(one_time_step=True) \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py index 65bf94554..d55845621 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py +++ b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -17,35 +18,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -63,8 +60,22 @@ # Instance of the simulation # -------------------------- +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +args, _ = parser.parse_known_args() + # Environment options -env = EnvironmentOptions(sim_folder="sim_data_petsc") +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, +) # Time stepping time_opts = Time(dt = 0.05, Tend = 75.0, split_algo = "LieTrotter") @@ -100,7 +111,7 @@ loading_params = LoadingParameters(ppc=20, seed=42) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) saving_params = SavingParameters(binning_plots=(binplot,)) @@ -142,4 +153,8 @@ model.kinetic_ions.var.add_initial_condition(init) if __name__ == "__main__": - sim.run() \ No newline at end of file + # one_time_step=True isolates the initial Poisson solve (the only part of this model that + # solver= affects -- the field then evolves via VlasovAmpereCoupling, unrelated to + # PETScSolver) from the ~1500 identical transport steps a full Tend=75.0 run would otherwise + # dilute the comparison with. + sim.run(one_time_step=True) \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py index 2b159a304..b96c5cea5 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py +++ b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -17,35 +18,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -64,7 +61,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_pcg") +env = EnvironmentOptions(sim_folder="sim_data_pcg", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) # Time stepping time_opts = Time(dt = 0.1, Tend = 50.0, split_algo = "LieTrotter") @@ -100,7 +97,7 @@ loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) saving_params = SavingParameters(binning_plots=(binplot,)) diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py index 7eb1e5248..83e9a0174 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py +++ b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -17,35 +18,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -64,7 +61,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_petsc") +env = EnvironmentOptions(sim_folder="sim_data_petsc", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) # Time stepping time_opts = Time(dt = 0.1, Tend = 50.0, split_algo = "LieTrotter") @@ -100,7 +97,7 @@ loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) saving_params = SavingParameters(binning_plots=(binplot,)) diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py index ffafdae60..568a8e499 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py +++ b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -16,35 +17,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -63,7 +60,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_pcg") +env = EnvironmentOptions(sim_folder="sim_data_pcg", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) # Time stepping time_opts = Time(dt = 0.05, Tend = 20.0, split_algo = "LieTrotter") @@ -99,7 +96,7 @@ loading_params = LoadingParameters(ppc=20, seed=42) weights_params = WeightsParameters(control_variate= True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) saving_params = SavingParameters(binning_plots=(binplot,)) diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py index 3ee39461f..658327c2a 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py +++ b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py @@ -1,3 +1,4 @@ +import os # ----------------------------- # Description of the simulation # ----------------------------- @@ -16,35 +17,31 @@ # Import Struphy API # ------------------ +# For particles: from struphy import ( BaseUnits, + BinningPlot, + BoundaryParameters, DerhamOptions, EnvironmentOptions, FieldsBackground, + KernelDensityPlot, + LoadingParameters, + SavingParameters, Simulation, + SortingParameters, Time, + WeightsParameters, domains, equils, grids, - perturbations, -) - -# For particles: -from struphy import ( - BinningPlot, - BoundaryParameters, - KernelDensityPlot, - LoadingParameters, - WeightsParameters, - SortingParameters, - SavingParameters, maxwellians, + perturbations, ) # --------------------- # Instance of the model # --------------------- - from struphy.models import VlasovAmpereOneSpecies # Units @@ -63,7 +60,7 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_petsc") +env = EnvironmentOptions(sim_folder="sim_data_petsc", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) # Time stepping time_opts = Time(dt = 0.05, Tend = 20.0, split_algo = "LieTrotter") @@ -99,7 +96,7 @@ loading_params = LoadingParameters(ppc=20, seed=42) weights_params = WeightsParameters(control_variate= True) boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(16, 1, 1), do_sort=True) +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) saving_params = SavingParameters(binning_plots=(binplot,)) From 94b36f290b71d89a835b01deb49016fc56846a12 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 09:13:44 +0200 Subject: [PATCH 12/32] Updated feectools --- feectools | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/feectools b/feectools index 7c371c39e..c10542411 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 7c371c39e4e7d81c6745ff2c07ba0bc1ba85c7ff +Subproject commit c10542411b89884d9a2167d0f71f49d00f3006e1 diff --git a/pyproject.toml b/pyproject.toml index 41fb7e5e0..8fc92f0c0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "numpy<=2.5.0", "cunumpy<=0.1.1", "pyccel>=2.2.0, <=2.2.3", - "feectools<=0.1.3", + "feectools<=0.1.11", "scipy<=1.18.0", "h5py<=3.16.0", "matplotlib<=3.11.0", From 555faf8656c1454b9b87dae7d1a7009fa6269017 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 11:40:40 +0200 Subject: [PATCH 13/32] Added toydrift_petsc example to profiling --- .../periodic_slab/params_periodic_slab_pcg.py | 163 ++++++++++++++++++ .../params_periodic_slab_petsc.py | 163 ++++++++++++++++++ profiling/submit_toydrift_petsc.py | 95 ++++++++++ 3 files changed, 421 insertions(+) create mode 100644 profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py create mode 100644 profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py create mode 100644 profiling/submit_toydrift_petsc.py diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py new file mode 100644 index 000000000..47ff72dc1 --- /dev/null +++ b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py @@ -0,0 +1,163 @@ +import os + +# ----------------------------- +# Description of the simulation +# ----------------------------- + +description = """ +Periodic-slab variant of the ToyDrift model (the model behind +examples/ToyGyrokinetic/diocotron_instability, which uses a physically non-periodic +HollowCylinder domain -- radial confinement is inherent to the diocotron instability). This case +swaps in a periodic Cuboid domain instead: unlike PoissonAdiabaticGyrokinetic (used by +DriftKineticElectrostaticAdiabatic), ToyDrift's field solve is a plain PoissonSolve with no +geometry-coupled averaging, so it works correctly on a periodic domain out of the box. + +Unlike VlasovAmpereOneSpecies (which only solves Poisson once, as an initial condition), this +model's gc_poisson runs as a *regular per-step propagator* -- exactly the repeated-solve pattern +where PETSc's algebraic multigrid preconditioner shows a genuine win (see +struphy.linear_algebra.petsc_examples_benchmark's module docstring for the general story). +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import ToyDrift + +# Units +base_units = BaseUnits(kBT=1.0) + +# Model instance +model = ToyDrift(base_units=base_units) + +# List all variables and decide whether to save their data +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +args, _ = parser.parse_known_args() + +# Environment options +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, +) + +# Time stepping +time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid() + +# Fluid equilibrium: straight B field, homogeneous density (default n0=1.0) +equil = equils.HomogenSlab(B0z=1.0, n0=1.0) + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble +# DirectionalDerivativeOperator along a non-periodic axis, see +# struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix) +derham_opts = DerhamOptions(degree=(2, 2, 2), bcs=(None, None, None)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) + +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + bufsize=0.4, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.gc_poisson.options.solver = "pcg" +model.propagators.gc_poisson.options.solver_params = SolverParameters(tol=1e-10, maxiter=20_000) +model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( + algo="explicit", + evaluate_e_field=True, +) + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). + +# Background for kinetic species +background = maxwellians.GyroMaxwellian2D( + n=(1.0, None), + vth_para=(1.0, None), + vth_perp=(1.0, None), + equil=equil, +) +model.kinetic_ions.var.add_background(background) + +# Perturbation, matching the Landau-damping style used elsewhere in this benchmark suite +perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) +init = maxwellians.GyroMaxwellian2D( + n=(1.0, perturbation), + vth_para=(1.0, None), + vth_perp=(1.0, None), + equil=equil, +) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + # one_time_step=True isolates the (still per-step, unlike VlasovAmpereOneSpecies) + # gc_poisson solve for a single-step timing snapshot. + sim.run(one_time_step=True) diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py new file mode 100644 index 000000000..ef4de1970 --- /dev/null +++ b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py @@ -0,0 +1,163 @@ +import os + +# ----------------------------- +# Description of the simulation +# ----------------------------- + +description = """ +Periodic-slab variant of the ToyDrift model (the model behind +examples/ToyGyrokinetic/diocotron_instability, which uses a physically non-periodic +HollowCylinder domain -- radial confinement is inherent to the diocotron instability). This case +swaps in a periodic Cuboid domain instead: unlike PoissonAdiabaticGyrokinetic (used by +DriftKineticElectrostaticAdiabatic), ToyDrift's field solve is a plain PoissonSolve with no +geometry-coupled averaging, so it works correctly on a periodic domain out of the box. + +Unlike VlasovAmpereOneSpecies (which only solves Poisson once, as an initial condition), this +model's gc_poisson runs as a *regular per-step propagator* -- exactly the repeated-solve pattern +where PETSc's algebraic multigrid preconditioner shows a genuine win (see +struphy.linear_algebra.petsc_examples_benchmark's module docstring for the general story). +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import ToyDrift + +# Units +base_units = BaseUnits(kBT=1.0) + +# Model instance +model = ToyDrift(base_units=base_units) + +# List all variables and decide whether to save their data +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +args, _ = parser.parse_known_args() + +# Environment options +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, +) + +# Time stepping +time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid() + +# Fluid equilibrium: straight B field, homogeneous density (default n0=1.0) +equil = equils.HomogenSlab(B0z=1.0, n0=1.0) + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble +# DirectionalDerivativeOperator along a non-periodic axis, see +# struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix) +derham_opts = DerhamOptions(degree=(2, 2, 2), bcs=(None, None, None)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=20, seed=42) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) + +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + bufsize=0.4, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.gc_poisson.options.solver = "petsc" +model.propagators.gc_poisson.options.solver_params = SolverParameters(tol=1e-10, maxiter=20_000) +model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( + algo="explicit", + evaluate_e_field=True, +) + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). + +# Background for kinetic species +background = maxwellians.GyroMaxwellian2D( + n=(1.0, None), + vth_para=(1.0, None), + vth_perp=(1.0, None), + equil=equil, +) +model.kinetic_ions.var.add_background(background) + +# Perturbation, matching the Landau-damping style used elsewhere in this benchmark suite +perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) +init = maxwellians.GyroMaxwellian2D( + n=(1.0, perturbation), + vth_para=(1.0, None), + vth_perp=(1.0, None), + equil=equil, +) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + # one_time_step=True isolates the (still per-step, unlike VlasovAmpereOneSpecies) + # gc_poisson solve for a single-step timing snapshot. + sim.run(one_time_step=True) diff --git a/profiling/submit_toydrift_petsc.py b/profiling/submit_toydrift_petsc.py new file mode 100644 index 000000000..6b953183c --- /dev/null +++ b/profiling/submit_toydrift_petsc.py @@ -0,0 +1,95 @@ +"""ToyDrift periodic-slab: PETSc vs. pcg for a real per-step Poisson solve. + +This file defines two profiling cases (the `ProfilingCase`s) built from the same periodic-slab +ToyDrift setup (see `profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_{pcg,petsc}.py` +-- there is no periodic ToyDrift example under `examples/`: the real one, +`examples/ToyGyrokinetic/diocotron_instability`, uses a physically non-periodic HollowCylinder +domain, since radial confinement is inherent to the diocotron instability; this case swaps in a +periodic Cuboid domain instead, which works because ToyDrift's field solve is a plain +`PoissonSolve` with no geometry-coupled averaging -- unlike `PoissonAdiabaticGyrokinetic`, used by +`DriftKineticElectrostaticAdiabatic`, which was tried first and diverges outright on a periodic +domain regardless of options), differing only in the solver used for the field solve +(`model.propagators.gc_poisson.options.solver`, `"pcg"` vs. `"petsc"`). + +Unlike VlasovAmpereOneSpecies (whose Poisson solve only runs once, as an initial condition), +ToyDrift's `gc_poisson` runs as a *regular per-step propagator* -- exactly the repeated-solve +pattern where PETSc's algebraic multigrid preconditioner shows a genuine win, without needing the +initial-Poisson-only benchmark's workaround of re-invoking the propagator by hand after +`sim.run()`. Each generated script runs the simulation itself by invoking the corresponding +`params_periodic_slab_{pcg,petsc}.py` directly (its `__main__` block calls +`sim.run(one_time_step=True)`, a single-step timing snapshot). + +See `submit_strong_landau_damping_petsc.py` for the same comparison on VlasovAmpereOneSpecies, and +its module docstring / `struphy.linear_algebra.petsc_examples_benchmark`'s for the general PETSc +vs. pcg story (including the known MPI correctness caveat for this near-singular stab_eps regime). +""" + +import argparse +from pathlib import Path + +from profiling_job import ProfilingCase + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "ToyDrift" / "periodic_slab" + + profiling_cases = { + "pcg": ProfilingCase( + label="toydrift_periodic_slab_pcg", + name="ToyDrift periodic slab, per-step Poisson solve with pcg", + description=( + "Periodic-slab ToyDrift setup, solving the per-step guiding-center Poisson " + "problem with feectools' native preconditioned CG." + ), + physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", + struphy_model_used="ToyDrift", + params_source=params_dir / "params_periodic_slab_pcg.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ), + "petsc": ProfilingCase( + label="toydrift_periodic_slab_petsc", + name="ToyDrift periodic slab, per-step Poisson solve with PETSc", + description=( + "Same setup as 'toydrift_periodic_slab_pcg', but solving the per-step " + "guiding-center Poisson problem with PETScSolver (KSP=cg, PC=gamg) instead." + ), + physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", + struphy_model_used="ToyDrift", + params_source=params_dir / "params_periodic_slab_petsc.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ), + } + + for profiling_case in profiling_cases.values(): + profiling_case.use_slurm = False + + # Launch one run per rank count, same counts for both solvers so they are directly + # comparable. + for num_tasks in (1, 2, 4): + profiling_case.launch(num_tasks) + + # Package and push each run as its own job finishes. + for profiling_case in profiling_cases.values(): + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() From 6c6abf07bda747b92313b8ac19b64abc32cecf22 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 16:13:11 +0200 Subject: [PATCH 14/32] Updated params for the examples --- .../ToyDrift/periodic_slab/params_periodic_slab_pcg.py | 6 +++--- .../ToyDrift/periodic_slab/params_periodic_slab_petsc.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py index 47ff72dc1..92d42db01 100644 --- a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py +++ b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py @@ -86,12 +86,12 @@ equil = equils.HomogenSlab(B0z=1.0, n0=1.0) # Grid -grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) +grid = grids.TensorProductGrid(num_elements=(24, 24, 24)) # Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble # DirectionalDerivativeOperator along a non-periodic axis, see # struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix) -derham_opts = DerhamOptions(degree=(2, 2, 2), bcs=(None, None, None)) +derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None)) # Simulation object sim = Simulation( @@ -109,7 +109,7 @@ # Particle parameters # ------------------- -loading_params = LoadingParameters(ppc=20, seed=42) +loading_params = LoadingParameters(ppc=5, seed=42) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py index ef4de1970..17b615798 100644 --- a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py +++ b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py @@ -86,12 +86,12 @@ equil = equils.HomogenSlab(B0z=1.0, n0=1.0) # Grid -grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) +grid = grids.TensorProductGrid(num_elements=(24, 24, 24)) # Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble # DirectionalDerivativeOperator along a non-periodic axis, see # struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix) -derham_opts = DerhamOptions(degree=(2, 2, 2), bcs=(None, None, None)) +derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None)) # Simulation object sim = Simulation( @@ -109,7 +109,7 @@ # Particle parameters # ------------------- -loading_params = LoadingParameters(ppc=20, seed=42) +loading_params = LoadingParameters(ppc=5, seed=42) weights_params = WeightsParameters(control_variate=True) boundary_params = BoundaryParameters() sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) From ad32feb7e0200d2042e1735d52af0ec40068253f Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 16:55:37 +0200 Subject: [PATCH 15/32] Updated profiling examples --- ...ic_slab_pcg.py => params_periodic_slab.py} | 5 +- .../params_periodic_slab_petsc.py | 163 ------------------ ...arams_bump_on_pcg.py => params_bump_on.py} | 71 +++++--- .../bump_on/params_bump_on_petsc.py | 148 ---------------- ...pcg.py => params_strong_Landau_damping.py} | 36 ++-- ...two_stream_pcg.py => params_two_stream.py} | 65 ++++--- .../two_stream/params_two_stream_petsc.py | 147 ---------------- .../params_weak_Landau_damping.py} | 56 +++--- .../params_weak_Landau_damping_pcg.py | 143 --------------- .../params_weak_Landau_damping_petsc.py | 143 --------------- 10 files changed, 149 insertions(+), 828 deletions(-) rename profiling/examples/ToyDrift/periodic_slab/{params_periodic_slab_pcg.py => params_periodic_slab.py} (96%) delete mode 100644 profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py rename profiling/examples/VlasovAmpereOneSpecies/bump_on/{params_bump_on_pcg.py => params_bump_on.py} (62%) delete mode 100644 profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py rename profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/{params_strong_Landau_damping_pcg.py => params_strong_Landau_damping.py} (85%) rename profiling/examples/VlasovAmpereOneSpecies/two_stream/{params_two_stream_pcg.py => params_two_stream.py} (64%) delete mode 100644 profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py rename profiling/examples/VlasovAmpereOneSpecies/{strong_Landau_damping/params_strong_Landau_damping_petsc.py => weak_Landau_damping/params_weak_Landau_damping.py} (72%) delete mode 100644 profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py delete mode 100644 profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py similarity index 96% rename from profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py rename to profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py index 92d42db01..64da7c98f 100644 --- a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_pcg.py +++ b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py @@ -66,6 +66,9 @@ parser = argparse.ArgumentParser() parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." +) args, _ = parser.parse_known_args() # Environment options @@ -126,7 +129,7 @@ # Propagator options # ------------------ -model.propagators.gc_poisson.options.solver = "pcg" +model.propagators.gc_poisson.options.solver = args.solver model.propagators.gc_poisson.options.solver_params = SolverParameters(tol=1e-10, maxiter=20_000) model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( algo="explicit", diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py deleted file mode 100644 index 17b615798..000000000 --- a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_petsc.py +++ /dev/null @@ -1,163 +0,0 @@ -import os - -# ----------------------------- -# Description of the simulation -# ----------------------------- - -description = """ -Periodic-slab variant of the ToyDrift model (the model behind -examples/ToyGyrokinetic/diocotron_instability, which uses a physically non-periodic -HollowCylinder domain -- radial confinement is inherent to the diocotron instability). This case -swaps in a periodic Cuboid domain instead: unlike PoissonAdiabaticGyrokinetic (used by -DriftKineticElectrostaticAdiabatic), ToyDrift's field solve is a plain PoissonSolve with no -geometry-coupled averaging, so it works correctly on a periodic domain out of the box. - -Unlike VlasovAmpereOneSpecies (which only solves Poisson once, as an initial condition), this -model's gc_poisson runs as a *regular per-step propagator* -- exactly the repeated-solve pattern -where PETSc's algebraic multigrid preconditioner shows a genuine win (see -struphy.linear_algebra.petsc_examples_benchmark's module docstring for the general story). -""" - -# ------------------ -# Import Struphy API -# ------------------ - -from struphy import ( - BaseUnits, - BoundaryParameters, - DerhamOptions, - EnvironmentOptions, - LoadingParameters, - Simulation, - SortingParameters, - Time, - WeightsParameters, - domains, - equils, - grids, - maxwellians, - perturbations, -) -from struphy.linear_algebra.solver import SolverParameters - -# --------------------- -# Instance of the model -# --------------------- -from struphy.models import ToyDrift - -# Units -base_units = BaseUnits(kBT=1.0) - -# Model instance -model = ToyDrift(base_units=base_units) - -# List all variables and decide whether to save their data -model.em_fields.phi.save_data = True -model.kinetic_ions.var.save_data = False - -# -------------------------- -# Instance of the simulation -# -------------------------- - -# `--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. -import argparse - -parser = argparse.ArgumentParser() -parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") -args, _ = parser.parse_known_args() - -# Environment options -env = EnvironmentOptions( - sim_folder=f"sim_{args.id:02d}", - out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), - profiling_activated=True, - profiling_trace=True, -) - -# Time stepping -time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") - -# Geometry -domain = domains.Cuboid() - -# Fluid equilibrium: straight B field, homogeneous density (default n0=1.0) -equil = equils.HomogenSlab(B0z=1.0, n0=1.0) - -# Grid -grid = grids.TensorProductGrid(num_elements=(24, 24, 24)) - -# Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble -# DirectionalDerivativeOperator along a non-periodic axis, see -# struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix) -derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None)) - -# Simulation object -sim = Simulation( - model=model, - params_path=__file__, - env=env, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, -) - -# ------------------- -# Particle parameters -# ------------------- - -loading_params = LoadingParameters(ppc=5, seed=42) -weights_params = WeightsParameters(control_variate=True) -boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) - -model.kinetic_ions.set_markers( - loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - bufsize=0.4, -) - -# ------------------ -# Propagator options -# ------------------ - -model.propagators.gc_poisson.options.solver = "petsc" -model.propagators.gc_poisson.options.solver_params = SolverParameters(tol=1e-10, maxiter=20_000) -model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( - algo="explicit", - evaluate_e_field=True, -) - -# ------------------ -# Initial conditions -# ------------------ -# Initial conditions are the sum of the background(s) and the perturbation(s). - -# Background for kinetic species -background = maxwellians.GyroMaxwellian2D( - n=(1.0, None), - vth_para=(1.0, None), - vth_perp=(1.0, None), - equil=equil, -) -model.kinetic_ions.var.add_background(background) - -# Perturbation, matching the Landau-damping style used elsewhere in this benchmark suite -perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) -init = maxwellians.GyroMaxwellian2D( - n=(1.0, perturbation), - vth_para=(1.0, None), - vth_perp=(1.0, None), - equil=equil, -) -model.kinetic_ions.var.add_initial_condition(init) - -if __name__ == "__main__": - # one_time_step=True isolates the (still per-step, unlike VlasovAmpereOneSpecies) - # gc_poisson solve for a single-step timing snapshot. - sim.run(one_time_step=True) diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py similarity index 62% rename from profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py rename to profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py index 22feb5aa2..36f08eb86 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_pcg.py +++ b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py @@ -1,8 +1,9 @@ import os + # ----------------------------- # Description of the simulation # ----------------------------- -# Please fill in a verbal 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. description = """ @@ -49,7 +50,7 @@ base_units = BaseUnits() # Model instance -model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) # List all variables and decide whether to save their data model.em_fields.e_field.save_data = True @@ -61,13 +62,30 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_pcg", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." +) +args, _ = parser.parse_known_args() + +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, +) # Time stepping -time_opts = Time(dt = 0.1, Tend = 60.0, split_algo = "LieTrotter") +time_opts = Time(dt=0.1, Tend=60.0, split_algo="LieTrotter") # Geometry -domain = domains.Cuboid(r1 = 62.83) +domain = domains.Cuboid(r1=62.83) # Fluid equilibrium (can be used as part of initial conditions) equil = None @@ -99,27 +117,32 @@ boundary_params = BoundaryParameters() sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) -binplot_1 = BinningPlot(slice="e1_v1", n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) #for initial velocity distribution -binplot_2 = BinningPlot(slice = "v1", n_bins = 128, ranges = (-10.0,10.0)) # for progression of velocity and space distribution +binplot_1 = BinningPlot( + slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-10.0, 10.0)) +) # for initial velocity distribution +binplot_2 = BinningPlot( + slice="v1", n_bins=128, ranges=(-10.0, 10.0) +) # for progression of velocity and space distribution saving_params = SavingParameters(binning_plots=(binplot_1, binplot_2)) -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4, - ) +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=0.4, +) # ------------------ # Propagator options # ------------------ -model.propagators.push_eta.options = model.propagators.push_eta.Options() +model.propagators.push_eta.options = model.propagators.push_eta.Options() if model.with_B0: model.propagators.push_vxb.options = model.propagators.push_vxb.Options() model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver=args.solver) # ------------------ # Initial conditions @@ -132,17 +155,21 @@ # For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). # Background for kinetic species -maxwellian_1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) -maxwellian_2 = maxwellians.Maxwellian3D(n=(1/10, None), u1 = (-4.5, None), vth1 = (0.5, None)) +maxwellian_1 = maxwellians.Maxwellian3D(n=(9 / 10, None), u1=(3.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(1 / 10, None), u1=(-4.5, None), vth1=(0.5, None)) background = maxwellian_1 + maxwellian_2 model.kinetic_ions.var.add_background(background) # Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.05,), ls = (1,)) -init1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) -init2 = maxwellians.Maxwellian3D(n = (1/10, perturbation), u1 = (-4.5, None), vth1 = (0.5, None)) +perturbation = perturbations.ModesCos(amps=(0.05,), ls=(1,)) +init1 = maxwellians.Maxwellian3D(n=(9 / 10, None), u1=(3.0, None)) +init2 = maxwellians.Maxwellian3D(n=(1 / 10, perturbation), u1=(-4.5, None), vth1=(0.5, None)) init = init1 + init2 model.kinetic_ions.var.add_initial_condition(init) if __name__ == "__main__": - sim.run() \ No newline at end of file + # one_time_step=True isolates the initial Poisson solve (the only part of this model that + # solver= affects -- the field then evolves via VlasovAmpereCoupling, unrelated to + # PETScSolver) from the many identical transport steps a full run would otherwise dilute + # the comparison with. + sim.run(one_time_step=True) diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py deleted file mode 100644 index d98daea4e..000000000 --- a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on_petsc.py +++ /dev/null @@ -1,148 +0,0 @@ -import os -# ----------------------------- -# 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. - -description = """ -Nonlinear bump-on-tail instability: A kinetic plasma instability test case for the Vlasov-Ampère model. -This test features a "bump" (localized excess) in the high-velocity tail of the electron velocity distribution. -The bump-on-tail configuration is unstable to the generation of Langmuir waves, leading to energy transfer -from the hot electron population to the growing wave field. This nonlinear process exhibits complex dynamics -including mode coupling and particle trapping in the wave potential. -This benchmark validates the particle-in-cell treatment of velocity-space instabilities and wave-particle interactions. -""" - -# ------------------ -# Import Struphy API -# ------------------ - -# For particles: -from struphy import ( - BaseUnits, - BinningPlot, - BoundaryParameters, - DerhamOptions, - EnvironmentOptions, - FieldsBackground, - KernelDensityPlot, - LoadingParameters, - SavingParameters, - Simulation, - SortingParameters, - Time, - WeightsParameters, - domains, - equils, - grids, - maxwellians, - perturbations, -) - -# --------------------- -# Instance of the model -# --------------------- -from struphy.models import VlasovAmpereOneSpecies - -# Units -base_units = BaseUnits() - -# Model instance -model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) - -# List all variables and decide whether to save their data -model.em_fields.e_field.save_data = True -model.em_fields.phi.save_data = True -model.kinetic_ions.var.save_data = True - -# -------------------------- -# Instance of the simulation -# -------------------------- - -# Environment options -env = EnvironmentOptions(sim_folder="sim_data_petsc", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) - -# Time stepping -time_opts = Time(dt = 0.1, Tend = 60.0, split_algo = "LieTrotter") - -# Geometry -domain = domains.Cuboid(r1 = 62.83) - -# Fluid equilibrium (can be used as part of initial conditions) -equil = None - -# Grid -grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) - -# Derham options -derham_opts = DerhamOptions(degree=(3, 1, 1)) - -# Simulation object -sim = Simulation( - model=model, - params_path=__file__, - env=env, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, -) - -# ------------------- -# Particle parameters -# ------------------- - -loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) -weights_params = WeightsParameters(control_variate=True) -boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) - -binplot_1 = BinningPlot(slice="e1_v1", n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) #for initial velocity distribution -binplot_2 = BinningPlot(slice = "v1", n_bins = 128, ranges = (-10.0,10.0)) # for progression of velocity and space distribution -saving_params = SavingParameters(binning_plots=(binplot_1, binplot_2)) - -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4, - ) - -# ------------------ -# Propagator options -# ------------------ - -model.propagators.push_eta.options = model.propagators.push_eta.Options() -if model.with_B0: - model.propagators.push_vxb.options = model.propagators.push_vxb.Options() -model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") - -# ------------------ -# Initial conditions -# ------------------ -# Initial conditions are the sum of the background(s) and the perturbation(s). -# If backgrounds or perturbations are not specified, they are assumed to be zero. - -# For kinetic species the background is mandatory. -# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. -# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). - -# Background for kinetic species -maxwellian_1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) -maxwellian_2 = maxwellians.Maxwellian3D(n=(1/10, None), u1 = (-4.5, None), vth1 = (0.5, None)) -background = maxwellian_1 + maxwellian_2 -model.kinetic_ions.var.add_background(background) - -# Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.05,), ls = (1,)) -init1 = maxwellians.Maxwellian3D(n=(9/10, None), u1 = (3.0, None)) -init2 = maxwellians.Maxwellian3D(n = (1/10, perturbation), u1 = (-4.5, None), vth1 = (0.5, None)) -init = init1 + init2 -model.kinetic_ions.var.add_initial_condition(init) - -if __name__ == "__main__": - sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py similarity index 85% rename from profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py rename to profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py index 37c554e48..77aa639ce 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_pcg.py +++ b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py @@ -1,8 +1,9 @@ import os + # ----------------------------- # Description of the simulation # ----------------------------- -# Please fill in a verbal 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. description = """ @@ -67,6 +68,9 @@ parser = argparse.ArgumentParser() parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." +) args, _ = parser.parse_known_args() # Environment options @@ -78,10 +82,10 @@ ) # Time stepping -time_opts = Time(dt = 0.05, Tend = 75.0, split_algo = "LieTrotter") +time_opts = Time(dt=0.05, Tend=75.0, split_algo="LieTrotter") # Geometry -domain = domains.Cuboid(r1 = 12.56) +domain = domains.Cuboid(r1=12.56) # Fluid equilibrium (can be used as part of initial conditions) equil = None @@ -113,25 +117,27 @@ boundary_params = BoundaryParameters() sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) -binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) +binplot = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-5.0, 5.0))) saving_params = SavingParameters(binning_plots=(binplot,)) -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4,) +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=0.4, +) # ------------------ # Propagator options # ------------------ -model.propagators.push_eta.options = model.propagators.push_eta.Options() +model.propagators.push_eta.options = model.propagators.push_eta.Options() if model.with_B0: model.propagators.push_vxb.options = model.propagators.push_vxb.Options() model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver=args.solver) # ------------------ # Initial conditions @@ -148,8 +154,8 @@ model.kinetic_ions.var.add_background(background) # Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.5,), ls = (1,)) -init = maxwellians.Maxwellian3D(n = (1.0, perturbation)) +perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) +init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) model.kinetic_ions.var.add_initial_condition(init) if __name__ == "__main__": @@ -157,4 +163,4 @@ # solver= affects -- the field then evolves via VlasovAmpereCoupling, unrelated to # PETScSolver) from the ~1500 identical transport steps a full Tend=75.0 run would otherwise # dilute the comparison with. - sim.run(one_time_step=True) \ No newline at end of file + sim.run(one_time_step=True) diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py similarity index 64% rename from profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py rename to profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py index b96c5cea5..0d8565982 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_pcg.py +++ b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py @@ -1,8 +1,9 @@ import os + # ----------------------------- # Description of the simulation # ----------------------------- -# Please fill in a verbal 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. description = """ @@ -49,7 +50,7 @@ base_units = BaseUnits() # Model instance -model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) # List all variables and decide whether to save their data model.em_fields.e_field.save_data = True @@ -61,13 +62,30 @@ # -------------------------- # Environment options -env = EnvironmentOptions(sim_folder="sim_data_pcg", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." +) +args, _ = parser.parse_known_args() + +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, +) # Time stepping -time_opts = Time(dt = 0.1, Tend = 50.0, split_algo = "LieTrotter") +time_opts = Time(dt=0.1, Tend=50.0, split_algo="LieTrotter") # Geometry -domain = domains.Cuboid(r1 = 31.42) +domain = domains.Cuboid(r1=31.42) # Fluid equilibrium (can be used as part of initial conditions) equil = None @@ -99,26 +117,27 @@ boundary_params = BoundaryParameters() sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) -binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) +binplot = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-10.0, 10.0))) saving_params = SavingParameters(binning_plots=(binplot,)) -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4, - ) +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=0.4, +) # ------------------ # Propagator options # ------------------ -model.propagators.push_eta.options = model.propagators.push_eta.Options() +model.propagators.push_eta.options = model.propagators.push_eta.Options() if model.with_B0: model.propagators.push_vxb.options = model.propagators.push_vxb.Options() model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver=args.solver) # ------------------ # Initial conditions @@ -131,17 +150,21 @@ # For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). # Background for kinetic species -maxwellian_1 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (3.0, None)) -maxwellian_2 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (-3.0, None)) +maxwellian_1 = maxwellians.Maxwellian3D(n=(0.5, None), u1=(3.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(0.5, None), u1=(-3.0, None)) background = maxwellian_1 + maxwellian_2 model.kinetic_ions.var.add_background(background) # Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) -init1 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (3.0, None)) -init2 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (-3.0, None)) +perturbation = perturbations.ModesCos(amps=(0.001,), ls=(1,)) +init1 = maxwellians.Maxwellian3D(n=(0.5, perturbation), u1=(3.0, None)) +init2 = maxwellians.Maxwellian3D(n=(0.5, perturbation), u1=(-3.0, None)) init = init1 + init2 model.kinetic_ions.var.add_initial_condition(init) if __name__ == "__main__": - sim.run() \ No newline at end of file + # one_time_step=True isolates the initial Poisson solve (the only part of this model that + # solver= affects -- the field then evolves via VlasovAmpereCoupling, unrelated to + # PETScSolver) from the many identical transport steps a full run would otherwise dilute + # the comparison with. + sim.run(one_time_step=True) diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py deleted file mode 100644 index 83e9a0174..000000000 --- a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream_petsc.py +++ /dev/null @@ -1,147 +0,0 @@ -import os -# ----------------------------- -# 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. - -description = """ -Nonlinear two-stream instability: A fundamental kinetic test case for the Vlasov-Ampère model. -This test involves two counter-streaming particle populations with a small perturbation that triggers -the two-stream instability. The instability leads to the formation of electron acoustic waves and -subsequent nonlinear effects including particle trapping and energy exchange between modes. -This benchmark validates the numerical treatment of beam-plasma interactions and tests the accuracy -of the particle-in-cell method in capturing mode coupling and energy transfer phenomena. -""" - -# ------------------ -# Import Struphy API -# ------------------ - -# For particles: -from struphy import ( - BaseUnits, - BinningPlot, - BoundaryParameters, - DerhamOptions, - EnvironmentOptions, - FieldsBackground, - KernelDensityPlot, - LoadingParameters, - SavingParameters, - Simulation, - SortingParameters, - Time, - WeightsParameters, - domains, - equils, - grids, - maxwellians, - perturbations, -) - -# --------------------- -# Instance of the model -# --------------------- -from struphy.models import VlasovAmpereOneSpecies - -# Units -base_units = BaseUnits() - -# Model instance -model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) - -# List all variables and decide whether to save their data -model.em_fields.e_field.save_data = True -model.em_fields.phi.save_data = True -model.kinetic_ions.var.save_data = True - -# -------------------------- -# Instance of the simulation -# -------------------------- - -# Environment options -env = EnvironmentOptions(sim_folder="sim_data_petsc", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) - -# Time stepping -time_opts = Time(dt = 0.1, Tend = 50.0, split_algo = "LieTrotter") - -# Geometry -domain = domains.Cuboid(r1 = 31.42) - -# Fluid equilibrium (can be used as part of initial conditions) -equil = None - -# Grid -grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) - -# Derham options -derham_opts = DerhamOptions(degree=(3, 1, 1)) - -# Simulation object -sim = Simulation( - model=model, - params_path=__file__, - env=env, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, -) - -# ------------------- -# Particle parameters -# ------------------- - -loading_params = LoadingParameters(ppc=20, seed=42, moments=(0.0, 0.0, 0.0, 3.0, 1.0, 1.0)) -weights_params = WeightsParameters(control_variate=True) -boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) - -binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-10.0,10.0))) -saving_params = SavingParameters(binning_plots=(binplot,)) - -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4, - ) - -# ------------------ -# Propagator options -# ------------------ - -model.propagators.push_eta.options = model.propagators.push_eta.Options() -if model.with_B0: - model.propagators.push_vxb.options = model.propagators.push_vxb.Options() -model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") - -# ------------------ -# Initial conditions -# ------------------ -# Initial conditions are the sum of the background(s) and the perturbation(s). -# If backgrounds or perturbations are not specified, they are assumed to be zero. - -# For kinetic species the background is mandatory. -# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. -# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). - -# Background for kinetic species -maxwellian_1 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (3.0, None)) -maxwellian_2 = maxwellians.Maxwellian3D(n=(0.5, None), u1 = (-3.0, None)) -background = maxwellian_1 + maxwellian_2 -model.kinetic_ions.var.add_background(background) - -# Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) -init1 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (3.0, None)) -init2 = maxwellians.Maxwellian3D(n = (0.5, perturbation), u1 = (-3.0, None)) -init = init1 + init2 -model.kinetic_ions.var.add_initial_condition(init) - -if __name__ == "__main__": - sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py similarity index 72% rename from profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py rename to profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py index d55845621..88e60a37b 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping_petsc.py +++ b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py @@ -1,17 +1,17 @@ import os + # ----------------------------- # Description of the simulation # ----------------------------- -# Please fill in a verbal 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. description = """ -Strong (nonlinear) Landau damping: A nonlinear test case for the VlasovAmpereOneSpecies model. -This test involves a large amplitude electrostatic perturbation in a uniform, collisionless plasma. -Unlike weak Landau damping, the nonlinear regime exhibits trapping of particles in the potential wells -of the self-consistent electric field, leading to vortex formation and complex phase space structures. -This benchmark tests the ability of the particle-in-cell method to capture nonlinear kinetic effects -and validates the long-term stability and accuracy of the Vlasov-Ampère discretization. +Weak Landau damping: A linear test case for the VlasovAmpereOneSpecies model. +This test involves a small amplitude electrostatic perturbation in a uniform, collisionless plasma. +The perturbation is damped due to phase mixing effects (Landau damping) as particles interact with +the self-consistent electric field. This benchmark validates the numerical discretization of the +Vlasov-Ampère system and the accuracy of particle-in-cell methods. """ # ------------------ @@ -60,6 +60,7 @@ # 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. @@ -67,9 +68,11 @@ parser = argparse.ArgumentParser() parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." +) args, _ = parser.parse_known_args() -# Environment options env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), @@ -78,10 +81,10 @@ ) # Time stepping -time_opts = Time(dt = 0.05, Tend = 75.0, split_algo = "LieTrotter") +time_opts = Time(dt=0.05, Tend=20.0, split_algo="LieTrotter") # Geometry -domain = domains.Cuboid(r1 = 12.56) +domain = domains.Cuboid(r1=12.56) # r1 -> pi * 4 -> k = 0.5 # Fluid equilibrium (can be used as part of initial conditions) equil = None @@ -90,7 +93,7 @@ grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) # Derham options -derham_opts = DerhamOptions() +derham_opts = DerhamOptions(degree=(3, 1, 1)) # Simulation object sim = Simulation( @@ -113,25 +116,28 @@ boundary_params = BoundaryParameters() sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) -binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) +binplot = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-5.0, 5.0))) saving_params = SavingParameters(binning_plots=(binplot,)) -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4,) +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=0.4, +) # ------------------ # Propagator options # ------------------ -model.propagators.push_eta.options = model.propagators.push_eta.Options() +model.propagators.push_eta.options = model.propagators.push_eta.Options() if model.with_B0: model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver=args.solver) # ------------------ # Initial conditions @@ -148,13 +154,13 @@ model.kinetic_ions.var.add_background(background) # Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.5,), ls = (1,)) -init = maxwellians.Maxwellian3D(n = (1.0, perturbation)) +perturbation = perturbations.ModesCos(amps=(0.001,), ls=(1,)) +init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) model.kinetic_ions.var.add_initial_condition(init) if __name__ == "__main__": # one_time_step=True isolates the initial Poisson solve (the only part of this model that # solver= affects -- the field then evolves via VlasovAmpereCoupling, unrelated to - # PETScSolver) from the ~1500 identical transport steps a full Tend=75.0 run would otherwise - # dilute the comparison with. - sim.run(one_time_step=True) \ No newline at end of file + # PETScSolver) from the many identical transport steps a full run would otherwise dilute + # the comparison with. + sim.run(one_time_step=True) diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py deleted file mode 100644 index 568a8e499..000000000 --- a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_pcg.py +++ /dev/null @@ -1,143 +0,0 @@ -import os -# ----------------------------- -# 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. - -description = """ -Weak Landau damping: A linear test case for the VlasovAmpereOneSpecies model. -This test involves a small amplitude electrostatic perturbation in a uniform, collisionless plasma. -The perturbation is damped due to phase mixing effects (Landau damping) as particles interact with -the self-consistent electric field. This benchmark validates the numerical discretization of the -Vlasov-Ampère system and the accuracy of particle-in-cell methods. -""" - -# ------------------ -# Import Struphy API -# ------------------ - -# For particles: -from struphy import ( - BaseUnits, - BinningPlot, - BoundaryParameters, - DerhamOptions, - EnvironmentOptions, - FieldsBackground, - KernelDensityPlot, - LoadingParameters, - SavingParameters, - Simulation, - SortingParameters, - Time, - WeightsParameters, - domains, - equils, - grids, - maxwellians, - perturbations, -) - -# --------------------- -# Instance of the model -# --------------------- -from struphy.models import VlasovAmpereOneSpecies - -# Units -base_units = BaseUnits() - -# Model instance -model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) - -# List all variables and decide whether to save their data -model.em_fields.e_field.save_data = True -model.em_fields.phi.save_data = True -model.kinetic_ions.var.save_data = True - -# -------------------------- -# Instance of the simulation -# -------------------------- - -# Environment options -env = EnvironmentOptions(sim_folder="sim_data_pcg", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) - -# Time stepping -time_opts = Time(dt = 0.05, Tend = 20.0, split_algo = "LieTrotter") - -# Geometry -domain = domains.Cuboid(r1 = 12.56) # r1 -> pi * 4 -> k = 0.5 - -# Fluid equilibrium (can be used as part of initial conditions) -equil = None - -# Grid -grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) - -# Derham options -derham_opts = DerhamOptions(degree=(3, 1, 1)) - -# Simulation object -sim = Simulation( - model=model, - params_path=__file__, - env=env, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, -) - -# ------------------- -# Particle parameters -# ------------------- - -loading_params = LoadingParameters(ppc=20, seed=42) -weights_params = WeightsParameters(control_variate= True) -boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) - -binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) -saving_params = SavingParameters(binning_plots=(binplot,)) - -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4, - ) - -# ------------------ -# Propagator options -# ------------------ - -model.propagators.push_eta.options = model.propagators.push_eta.Options() -if model.with_B0: - model.propagators.push_vxb.options = model.propagators.push_vxb.Options() - -model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="pcg") - -# ------------------ -# Initial conditions -# ------------------ -# Initial conditions are the sum of the background(s) and the perturbation(s). -# If backgrounds or perturbations are not specified, they are assumed to be zero. - -# For kinetic species the background is mandatory. -# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. -# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). - -# Background for kinetic species -background = maxwellians.Maxwellian3D(n=(1.0, None)) -model.kinetic_ions.var.add_background(background) - -# Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) -init = maxwellians.Maxwellian3D(n = (1.0,perturbation)) -model.kinetic_ions.var.add_initial_condition(init) - -if __name__ == "__main__": - sim.run() \ No newline at end of file diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py deleted file mode 100644 index 658327c2a..000000000 --- a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping_petsc.py +++ /dev/null @@ -1,143 +0,0 @@ -import os -# ----------------------------- -# 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. - -description = """ -Weak Landau damping: A linear test case for the VlasovAmpereOneSpecies model. -This test involves a small amplitude electrostatic perturbation in a uniform, collisionless plasma. -The perturbation is damped due to phase mixing effects (Landau damping) as particles interact with -the self-consistent electric field. This benchmark validates the numerical discretization of the -Vlasov-Ampère system and the accuracy of particle-in-cell methods. -""" - -# ------------------ -# Import Struphy API -# ------------------ - -# For particles: -from struphy import ( - BaseUnits, - BinningPlot, - BoundaryParameters, - DerhamOptions, - EnvironmentOptions, - FieldsBackground, - KernelDensityPlot, - LoadingParameters, - SavingParameters, - Simulation, - SortingParameters, - Time, - WeightsParameters, - domains, - equils, - grids, - maxwellians, - perturbations, -) - -# --------------------- -# Instance of the model -# --------------------- -from struphy.models import VlasovAmpereOneSpecies - -# Units -base_units = BaseUnits() - -# Model instance -model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0 = False) - -# List all variables and decide whether to save their data -model.em_fields.e_field.save_data = True -model.em_fields.phi.save_data = True -model.kinetic_ions.var.save_data = True - -# -------------------------- -# Instance of the simulation -# -------------------------- - -# Environment options -env = EnvironmentOptions(sim_folder="sim_data_petsc", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd())) - -# Time stepping -time_opts = Time(dt = 0.05, Tend = 20.0, split_algo = "LieTrotter") - -# Geometry -domain = domains.Cuboid(r1 = 12.56) # r1 -> pi * 4 -> k = 0.5 - -# Fluid equilibrium (can be used as part of initial conditions) -equil = None - -# Grid -grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) - -# Derham options -derham_opts = DerhamOptions(degree=(3, 1, 1)) - -# Simulation object -sim = Simulation( - model=model, - params_path=__file__, - env=env, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, -) - -# ------------------- -# Particle parameters -# ------------------- - -loading_params = LoadingParameters(ppc=20, seed=42) -weights_params = WeightsParameters(control_variate= True) -boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) - -binplot = BinningPlot(slice='e1_v1', n_bins= (128, 128), ranges= ((0.,1.), (-5.,5.))) -saving_params = SavingParameters(binning_plots=(binplot,)) - -model.kinetic_ions.set_markers(loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - saving_params=saving_params, - bufsize = 0.4, - ) - -# ------------------ -# Propagator options -# ------------------ - -model.propagators.push_eta.options = model.propagators.push_eta.Options() -if model.with_B0: - model.propagators.push_vxb.options = model.propagators.push_vxb.Options() - -model.propagators.coupling_va.options = model.propagators.coupling_va.Options() -model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver="petsc") - -# ------------------ -# Initial conditions -# ------------------ -# Initial conditions are the sum of the background(s) and the perturbation(s). -# If backgrounds or perturbations are not specified, they are assumed to be zero. - -# For kinetic species the background is mandatory. -# For kinetic species, if add_initial_condition() is not called, the background is taken as the kinetic initial condition. -# For kinetic species the perturbations are added to the moments of the distribution function (defined as tuples). - -# Background for kinetic species -background = maxwellians.Maxwellian3D(n=(1.0, None)) -model.kinetic_ions.var.add_background(background) - -# Perturbations for (some) kinetic species -perturbation = perturbations.ModesCos(amps = (0.001,), ls = (1,)) -init = maxwellians.Maxwellian3D(n = (1.0,perturbation)) -model.kinetic_ions.var.add_initial_condition(init) - -if __name__ == "__main__": - sim.run() \ No newline at end of file From 1b936ad82a5ab504af09f9e0b35104d810672d06 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 16:55:59 +0200 Subject: [PATCH 16/32] petsc and pcg in one case --- profiling/submit_toydrift_petsc.py | 87 ++++++++++++------------------ 1 file changed, 35 insertions(+), 52 deletions(-) diff --git a/profiling/submit_toydrift_petsc.py b/profiling/submit_toydrift_petsc.py index 6b953183c..5aede1a8c 100644 --- a/profiling/submit_toydrift_petsc.py +++ b/profiling/submit_toydrift_petsc.py @@ -1,23 +1,23 @@ """ToyDrift periodic-slab: PETSc vs. pcg for a real per-step Poisson solve. -This file defines two profiling cases (the `ProfilingCase`s) built from the same periodic-slab -ToyDrift setup (see `profiling/examples/ToyDrift/periodic_slab/params_periodic_slab_{pcg,petsc}.py` --- there is no periodic ToyDrift example under `examples/`: the real one, -`examples/ToyGyrokinetic/diocotron_instability`, uses a physically non-periodic HollowCylinder -domain, since radial confinement is inherent to the diocotron instability; this case swaps in a -periodic Cuboid domain instead, which works because ToyDrift's field solve is a plain -`PoissonSolve` with no geometry-coupled averaging -- unlike `PoissonAdiabaticGyrokinetic`, used by -`DriftKineticElectrostaticAdiabatic`, which was tried first and diverges outright on a periodic -domain regardless of options), differing only in the solver used for the field solve -(`model.propagators.gc_poisson.options.solver`, `"pcg"` vs. `"petsc"`). +This file defines a single profiling case built from the periodic-slab ToyDrift setup (see +`profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py` -- there is no periodic +ToyDrift example under `examples/`: the real one, `examples/ToyGyrokinetic/diocotron_instability`, +uses a physically non-periodic HollowCylinder domain, since radial confinement is inherent to the +diocotron instability; this case swaps in a periodic Cuboid domain instead, which works because +ToyDrift's field solve is a plain `PoissonSolve` with no geometry-coupled averaging -- unlike +`PoissonAdiabaticGyrokinetic`, used by `DriftKineticElectrostaticAdiabatic`, which was tried first +and diverges outright on a periodic domain regardless of options), launched once per (rank count, +solver) combination via `--solver pcg`/`--solver petsc` (`param_flags`), which sets +`model.propagators.gc_poisson.options.solver`. Unlike VlasovAmpereOneSpecies (whose Poisson solve only runs once, as an initial condition), ToyDrift's `gc_poisson` runs as a *regular per-step propagator* -- exactly the repeated-solve pattern where PETSc's algebraic multigrid preconditioner shows a genuine win, without needing the initial-Poisson-only benchmark's workaround of re-invoking the propagator by hand after -`sim.run()`. Each generated script runs the simulation itself by invoking the corresponding -`params_periodic_slab_{pcg,petsc}.py` directly (its `__main__` block calls -`sim.run(one_time_step=True)`, a single-step timing snapshot). +`sim.run()`. Each generated script runs the simulation itself by invoking +`params_periodic_slab.py` directly (its `__main__` block calls `sim.run(one_time_step=True)`, a +single-step timing snapshot). See `submit_strong_landau_damping_petsc.py` for the same comparison on VlasovAmpereOneSpecies, and its module docstring / `struphy.linear_algebra.petsc_examples_benchmark`'s for the general PETSc @@ -47,48 +47,31 @@ def main() -> None: script_dir = Path(__file__).resolve().parent params_dir = script_dir / "examples" / "ToyDrift" / "periodic_slab" - profiling_cases = { - "pcg": ProfilingCase( - label="toydrift_periodic_slab_pcg", - name="ToyDrift periodic slab, per-step Poisson solve with pcg", - description=( - "Periodic-slab ToyDrift setup, solving the per-step guiding-center Poisson " - "problem with feectools' native preconditioned CG." - ), - physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", - struphy_model_used="ToyDrift", - params_source=params_dir / "params_periodic_slab_pcg.py", - language="fortran", - compiler="GNU", - upload=args.upload, + profiling_case = ProfilingCase( + label="toydrift_periodic_slab_petsc", + name="ToyDrift periodic slab, per-step Poisson solve: pcg vs. PETSc", + description=( + "Periodic-slab ToyDrift setup, solving the per-step guiding-center Poisson problem " + "with either feectools' native preconditioned CG (--solver pcg) or PETScSolver " + "(KSP=cg, PC=gamg, --solver petsc)." ), - "petsc": ProfilingCase( - label="toydrift_periodic_slab_petsc", - name="ToyDrift periodic slab, per-step Poisson solve with PETSc", - description=( - "Same setup as 'toydrift_periodic_slab_pcg', but solving the per-step " - "guiding-center Poisson problem with PETScSolver (KSP=cg, PC=gamg) instead." - ), - physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", - struphy_model_used="ToyDrift", - params_source=params_dir / "params_periodic_slab_petsc.py", - language="fortran", - compiler="GNU", - upload=args.upload, - ), - } - - for profiling_case in profiling_cases.values(): - profiling_case.use_slurm = False + physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", + struphy_model_used="ToyDrift", + params_source=params_dir / "params_periodic_slab.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + profiling_case.use_slurm = False - # Launch one run per rank count, same counts for both solvers so they are directly - # comparable. - for num_tasks in (1, 2, 4): - profiling_case.launch(num_tasks) + # Launch one run per (rank count, solver) combination, same rank counts for both solvers so + # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` + # output folder), so the two solvers never collide even at the same rank count. + for num_tasks in (1, 2, 4): + for solver in ("pcg", "petsc"): + profiling_case.launch(num_tasks, param_flags=["--solver", solver]) - # Package and push each run as its own job finishes. - for profiling_case in profiling_cases.values(): - profiling_case.finalize_run() + profiling_case.finalize_run() if __name__ == "__main__": From e678db83e602ee5231209d4bb6ff56ab76f5a582 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 17:24:54 +0200 Subject: [PATCH 17/32] Update to scope-profiler 0.2.7 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8fc92f0c0..916651f51 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ dependencies = [ "pytest-testmon<=2.2.0", "ruff==0.15.0, <=0.16.0", "line_profiler<=5.0.2", - "scope-profiler==0.2.6, <=0.2.6", + "scope-profiler==0.2.7, <=0.2.7", ] [project.license] From 72b2c42ccbe1bcadffbe4b92de8137ff699509d5 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 17:41:17 +0200 Subject: [PATCH 18/32] Updated scope-profiler, set labels to each profiling simulation --- .../ToyDrift/periodic_slab/params_periodic_slab.py | 9 +++++++++ .../VlasovAmpereOneSpecies/bump_on/params_bump_on.py | 9 +++++++++ .../params_strong_Landau_damping.py | 9 +++++++++ .../two_stream/params_two_stream.py | 9 +++++++++ .../weak_Landau_damping/params_weak_Landau_damping.py | 9 +++++++++ src/struphy/io/options.py | 7 +++++++ src/struphy/simulation/sim.py | 4 ++-- 7 files changed, 54 insertions(+), 2 deletions(-) diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py index 64da7c98f..07d48ab18 100644 --- a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py +++ b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py @@ -71,12 +71,21 @@ ) args, _ = parser.parse_known_args() +# scope-profiler label: distinguishes solver/rank-count combinations in post-processing +# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. +from feectools.ddm.mpi import mpi as MPI + +_comm = MPI.COMM_WORLD +_num_ranks = _comm.Get_size() if _comm is not None else 1 +_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") + # Environment options env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), profiling_activated=True, profiling_trace=True, + profiling_label=_profiling_label, ) # Time stepping diff --git a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py index 36f08eb86..4b3cc063f 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py +++ b/profiling/examples/VlasovAmpereOneSpecies/bump_on/params_bump_on.py @@ -74,11 +74,20 @@ ) args, _ = parser.parse_known_args() +# scope-profiler label: distinguishes solver/rank-count combinations in post-processing +# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. +from feectools.ddm.mpi import mpi as MPI + +_comm = MPI.COMM_WORLD +_num_ranks = _comm.Get_size() if _comm is not None else 1 +_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") + env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), profiling_activated=True, profiling_trace=True, + profiling_label=_profiling_label, ) # Time stepping diff --git a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py index 77aa639ce..c89bd4aa9 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py +++ b/profiling/examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py @@ -73,12 +73,21 @@ ) args, _ = parser.parse_known_args() +# scope-profiler label: distinguishes solver/rank-count combinations in post-processing +# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. +from feectools.ddm.mpi import mpi as MPI + +_comm = MPI.COMM_WORLD +_num_ranks = _comm.Get_size() if _comm is not None else 1 +_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") + # Environment options env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), profiling_activated=True, profiling_trace=True, + profiling_label=_profiling_label, ) # Time stepping diff --git a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py index 0d8565982..240722552 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py +++ b/profiling/examples/VlasovAmpereOneSpecies/two_stream/params_two_stream.py @@ -74,11 +74,20 @@ ) args, _ = parser.parse_known_args() +# scope-profiler label: distinguishes solver/rank-count combinations in post-processing +# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. +from feectools.ddm.mpi import mpi as MPI + +_comm = MPI.COMM_WORLD +_num_ranks = _comm.Get_size() if _comm is not None else 1 +_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") + env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), profiling_activated=True, profiling_trace=True, + profiling_label=_profiling_label, ) # Time stepping diff --git a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py index 88e60a37b..fb8b46792 100644 --- a/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py +++ b/profiling/examples/VlasovAmpereOneSpecies/weak_Landau_damping/params_weak_Landau_damping.py @@ -73,11 +73,20 @@ ) args, _ = parser.parse_known_args() +# scope-profiler label: distinguishes solver/rank-count combinations in post-processing +# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. +from feectools.ddm.mpi import mpi as MPI + +_comm = MPI.COMM_WORLD +_num_ranks = _comm.Get_size() if _comm is not None else 1 +_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") + env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), profiling_activated=True, profiling_trace=True, + profiling_label=_profiling_label, ) # Time stepping diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index ab1dc2cfa..a78ea8cfa 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -330,6 +330,12 @@ class EnvironmentOptions(OptionsBase): profiling_trace: bool, optional Save time-trace of each profiling region (default=False) + + profiling_label: str, optional + Short name for this run, forwarded to ``ProfileManager.setup(label=...)`` (default=None, + i.e. scope-profiler falls back to the output file's stem). Used by scope-profiler's + post-processing (chart legends, summary heading, ``scope-profiler inspect``, JSON + statistics) to distinguish several runs being compared, e.g. ``"petsc, 4 ranks"``. """ out_folders: str = os.getcwd() @@ -341,6 +347,7 @@ class EnvironmentOptions(OptionsBase): num_clones: int = 1 profiling_activated: bool = False profiling_trace: bool = False + profiling_label: str | None = None def __post_init__(self): self.path_out: str = os.path.join(self.out_folders, self.sim_folder) diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 0c452016b..420b88958 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -214,14 +214,14 @@ def __init__( def _setup_profiling(self): # setup profiling agent ProfileManager.setup( - profiling_activated=self.env.profiling_activated, - time_trace=self.env.profiling_trace, + deactivate_profiling=not self.env.profiling_activated, use_likwid=False, file_path=os.path.join( self.env.out_folders, self.env.sim_folder, "profiling_data.h5", ), + label=self.env.profiling_label, ) def show_parameters(self): From 9226b96a2246e60ac5794e7ccf165a1d097a5e0c Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 21:43:29 +0200 Subject: [PATCH 19/32] Added profiling/submit_strong_landau_damping_petsc.py --- .../submit_strong_landau_damping_petsc.py | 52 +++++++++++++++++++ 1 file changed, 52 insertions(+) create mode 100644 profiling/submit_strong_landau_damping_petsc.py diff --git a/profiling/submit_strong_landau_damping_petsc.py b/profiling/submit_strong_landau_damping_petsc.py new file mode 100644 index 000000000..7854be74e --- /dev/null +++ b/profiling/submit_strong_landau_damping_petsc.py @@ -0,0 +1,52 @@ +import argparse +from pathlib import Path + +from profiling_job import ProfilingCase + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "VlasovAmpereOneSpecies" / "strong_Landau_damping" + + profiling_case = ProfilingCase( + label="strong_landau_damping_petsc", + name="Strong Landau damping, initial Poisson solve: pcg vs. PETSc", + description=( + "Strong (nonlinear) Landau damping test case for VlasovAmpereOneSpecies, solving the " + "one-time initial Poisson problem with either feectools' native preconditioned CG " + "(--solver pcg) or PETScSolver (KSP=cg, PC=gamg, --solver petsc)." + ), + physics_problem="Nonlinear Landau damping in a uniform, collisionless plasma.", + struphy_model_used="VlasovAmpereOneSpecies", + params_source=params_dir / "params_strong_Landau_damping.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + profiling_case.use_slurm = False + + # Launch one run per (rank count, solver) combination, same rank counts for both solvers so + # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` + # output folder), so the two solvers never collide even at the same rank count. + for num_tasks in (1, ): + for solver in ("pcg", "petsc"): + profiling_case.launch(num_tasks, param_flags=["--solver", solver]) + + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() From cb69555bdd2d7b153a1f6ff7dda69e5f42485690 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 21:46:39 +0200 Subject: [PATCH 20/32] Moved argparse to utils._get_profiling_args --- profiling/submit_diocotron_strong_scaling.py | 13 ++----------- profiling/submit_poisson_strong_scaling.py | 13 ++----------- profiling/submit_strong_landau_damping_petsc.py | 13 ++----------- profiling/utils.py | 14 ++++++++++++++ 4 files changed, 20 insertions(+), 33 deletions(-) diff --git a/profiling/submit_diocotron_strong_scaling.py b/profiling/submit_diocotron_strong_scaling.py index 449f76fd0..8fbcc7d6f 100644 --- a/profiling/submit_diocotron_strong_scaling.py +++ b/profiling/submit_diocotron_strong_scaling.py @@ -15,20 +15,11 @@ from pathlib import Path from profiling_job import ProfilingCase - +from utils import _get_profiling_args def main() -> None: - # Parse arguments, do not remove --upload - parser = argparse.ArgumentParser( - description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), - ) - parser.add_argument( - "--upload", - action="store_true", - help="Upload the packaged profiling results to the profiling-data repo.", - ) - args = parser.parse_args() + args = _get_profiling_args() # Paths relative to this script's location, so it can be run from anywhere. script_dir = Path(__file__).resolve().parent diff --git a/profiling/submit_poisson_strong_scaling.py b/profiling/submit_poisson_strong_scaling.py index 8929e5946..d3a31391e 100644 --- a/profiling/submit_poisson_strong_scaling.py +++ b/profiling/submit_poisson_strong_scaling.py @@ -13,20 +13,11 @@ from pathlib import Path from profiling_job import ProfilingCase - +from utils import _get_profiling_args def main() -> None: - # Parse arguments, do not remove --upload - parser = argparse.ArgumentParser( - description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), - ) - parser.add_argument( - "--upload", - action="store_true", - help="Upload the packaged profiling results to the profiling-data repo.", - ) - args = parser.parse_args() + args = _get_profiling_args() # Paths relative to this script's location, so it can be run from anywhere. script_dir = Path(__file__).resolve().parent diff --git a/profiling/submit_strong_landau_damping_petsc.py b/profiling/submit_strong_landau_damping_petsc.py index 7854be74e..71a616e16 100644 --- a/profiling/submit_strong_landau_damping_petsc.py +++ b/profiling/submit_strong_landau_damping_petsc.py @@ -2,20 +2,11 @@ from pathlib import Path from profiling_job import ProfilingCase - +from utils import _get_profiling_args def main() -> None: - # Parse arguments, do not remove --upload - parser = argparse.ArgumentParser( - description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), - ) - parser.add_argument( - "--upload", - action="store_true", - help="Upload the packaged profiling results to the profiling-data repo.", - ) - args = parser.parse_args() + args = _get_profiling_args() # Paths relative to this script's location, so it can be run from anywhere. script_dir = Path(__file__).resolve().parent diff --git a/profiling/utils.py b/profiling/utils.py index 8c99c3003..5d31e285f 100644 --- a/profiling/utils.py +++ b/profiling/utils.py @@ -13,6 +13,20 @@ from typing import Any +def _get_profiling_args() -> argparse.Namespace: + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + args = parser.parse_args() + + return args + def _slug(value: str) -> str: return re.sub(r"[^A-Za-z0-9._-]+", "_", value).strip("._-") or "unknown" From 7ca1a35c922419f6aa22937ba26df5342df3a4c1 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 10 Aug 2026 08:53:10 +0200 Subject: [PATCH 21/32] Added new example --- .../params_periodic_slab_hires.py | 193 ++++++++++++++++ .../params_weibel_instability.py | 211 ++++++++++++++++++ profiling/submit_toydrift_hires_petsc.py | 84 +++++++ profiling/submit_toydrift_petsc.py | 2 +- profiling/submit_vlasov_maxwell_petsc.py | 44 ++++ src/struphy/io/options.py | 1 + src/struphy/linear_algebra/schur_solver.py | 5 + src/struphy/linear_algebra/solver.py | 11 + src/struphy/propagators/implicit_diffusion.py | 1 + 9 files changed, 551 insertions(+), 1 deletion(-) create mode 100644 profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py create mode 100644 profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py create mode 100644 profiling/submit_toydrift_hires_petsc.py create mode 100644 profiling/submit_vlasov_maxwell_petsc.py diff --git a/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py b/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py new file mode 100644 index 000000000..2c90749f3 --- /dev/null +++ b/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py @@ -0,0 +1,193 @@ +import os + +# ----------------------------- +# Description of the simulation +# ----------------------------- + +description = """ +Higher-resolution variant of ToyDrift/periodic_slab (see that case's own description for the +model/domain background): a 32^3 grid with degree-3 splines (32768 dofs vs. periodic_slab's +24^3 grid's 13824) with PETSc's preconditioner set explicitly to algebraic multigrid +(`pc_type="gamg"`, via `SolverParameters.pc_type` -- see struphy.linear_algebra.solver. +SolverParameters), instead of the default `"jacobi"`. + +Measured via struphy.linear_algebra.petsc_examples_benchmark's repeated-solve methodology (which +isolates just the solve, amortizing one-time matrix/preconditioner setup across several calls -- +see that module's docstring): feectools' unpreconditioned CG took ~10.9 s/solve (190 iterations) +against PETSc+gamg's ~0.29 s/solve (2 iterations) here -- a ~38x difference. This is essentially +the same *ratio* periodic_slab already shows (~36x at 13824 dofs): the near-singular Poisson +system these ToyDrift cases solve (`stab_eps` clamped to ~1e-14 by ImplicitDiffusion's "always +stabilize" logic, see PETScSolver's docstring) is ill-conditioned mainly through its weakly +constrained constant/DC mode, not primarily through raw grid resolution, so pcg's iteration count +does not grow much further with size in this regime -- while gamg's convergence stays +essentially grid-independent regardless. What *does* grow with size is the absolute cost: at this +larger, more realistic problem size, pcg's ~11 seconds per Poisson solve is the practically +relevant "huge difference" -- multiplied over the many timesteps of a real simulation, it is the +difference between a run finishing in minutes and one taking hours. + +Like periodic_slab, gc_poisson runs as a *regular per-step propagator*, so a single +`sim.run(one_time_step=True)` call already times one full, representative solve. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import ToyDrift + +# Units +base_units = BaseUnits(kBT=1.0) + +# Model instance +model = ToyDrift(base_units=base_units) + +# List all variables and decide whether to save their data +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." +) +args, _ = parser.parse_known_args() + +# scope-profiler label: distinguishes solver/rank-count combinations in post-processing +# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. +from feectools.ddm.mpi import mpi as MPI + +_comm = MPI.COMM_WORLD +_num_ranks = _comm.Get_size() if _comm is not None else 1 +_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") + +# Environment options +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, + profiling_label=_profiling_label, +) + +# Time stepping +time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid() + +# Fluid equilibrium: straight B field, homogeneous density (default n0=1.0) +equil = equils.HomogenSlab(B0z=1.0, n0=1.0) + +# Grid -- 32^3, ~2x periodic_slab's 24^3 per direction (~2.4x the dofs at fixed degree): large +# enough for pcg's iteration count (hence cost) to blow up while PETSc+gamg stays flat. +grid = grids.TensorProductGrid(num_elements=(32, 32, 32)) + +# Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble +# DirectionalDerivativeOperator along a non-periodic axis, see +# struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix) +derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None)) + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters(ppc=5, seed=42) +weights_params = WeightsParameters(control_variate=True) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) + +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + bufsize=0.4, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.gc_poisson.options.solver = args.solver +# pc_type only affects solver="petsc" (ignored by pcg, see SolverParameters.pc_type); "gamg" is +# what makes PETSc's advantage large here. Note this only matters for the *steady-state* +# per-solve cost (see the module docstring): a single sim.run(one_time_step=True) call, as this +# file's __main__ performs, still pays gamg's one-time multigrid setup cost up front, so it will +# not by itself reproduce the headline speedup above -- that requires several solves at the same +# dt to amortize setup, exactly what petsc_examples_benchmark.py's repeated-solve timing does. +model.propagators.gc_poisson.options.solver_params = SolverParameters(tol=1e-10, maxiter=5_000, pc_type="gamg") +model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( + algo="explicit", + evaluate_e_field=True, +) + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). + +# Background for kinetic species +background = maxwellians.GyroMaxwellian2D( + n=(1.0, None), + vth_para=(1.0, None), + vth_perp=(1.0, None), + equil=equil, +) +model.kinetic_ions.var.add_background(background) + +# Perturbation, matching the Landau-damping style used elsewhere in this benchmark suite +perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) +init = maxwellians.GyroMaxwellian2D( + n=(1.0, perturbation), + vth_para=(1.0, None), + vth_perp=(1.0, None), + equil=equil, +) +model.kinetic_ions.var.add_initial_condition(init) + +if __name__ == "__main__": + # one_time_step=True isolates the (still per-step, unlike VlasovAmpereOneSpecies) + # gc_poisson solve for a single-step timing snapshot. + sim.run(one_time_step=True) diff --git a/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py b/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py new file mode 100644 index 000000000..d63b3920e --- /dev/null +++ b/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py @@ -0,0 +1,211 @@ +import os + +# ----------------------------- +# 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. + +description = """ +Weibel instability: A linear test case for the VlasovMaxwellOneSpecies model. This test considers +a plasma with an anisotropic velocity distribution, where temperature differs between directions. +Small magnetic perturbations grow due to the anisotropy, leading to the generation of transverse +magnetic fields. + +Like VlasovAmpereOneSpecies (see strong_Landau_damping/weak_Landau_damping/two_stream/bump_on in +this same profiling suite), VlasovMaxwellOneSpecies exposes its Poisson solve as a one-time +`model.initial_poisson` (not a regular per-step propagator: the fields then evolve via +MaxwellWeakAmpere, PushVxB and VlasovAmpereCoupling instead), so `solver=` only affects this +initial solve -- see `struphy.linear_algebra.petsc_examples_benchmark`'s module docstring for why +that benchmark re-invokes it directly rather than relying on a single sim.run(). + +Plain copy of examples/VlasovMaxwellOneSpecies/weibel_instability, with `num_elements` scaled up +from the original's tiny, highly-anisotropic 1D-style default of `(32, 1, 1)` cells to a proper +`(16, 16, 16)` 3D grid (PETSc's advantage only shows up above roughly 5,000 dofs -- see +struphy.linear_algebra.petsc_poisson_benchmark's module docstring), particle count reduced to +match, and a fixed seed (the original doesn't fix ppc/seed the same way) so pcg/petsc draw the +same particles. +""" + +# ------------------ +# Import Struphy API +# ------------------ + +from struphy import ( + BaseUnits, + BinningPlot, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + KernelDensityPlot, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import VlasovMaxwellOneSpecies + +# Units +base_units = BaseUnits() + +# Model instance +model = VlasovMaxwellOneSpecies( + base_units=base_units, + alpha=1.0, + epsilon=-1.0, + measure_gauss_law=True, +) + +# --------------------- +# Parameters setup +# --------------------- + +import cunumpy as xp + +k = 1.25 +B_pert_amp = -1e-4 + +vth1_background_val = 0.02 / xp.sqrt(2) +vth2_background_val = vth1_background_val * xp.sqrt(12) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.phi.save_data = True +model.kinetic_ions.var.save_data = True + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# `--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. +import argparse + +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." +) +args, _ = parser.parse_known_args() + +# scope-profiler label: distinguishes solver/rank-count combinations in post-processing +# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. +from feectools.ddm.mpi import mpi as MPI + +_comm = MPI.COMM_WORLD +_num_ranks = _comm.Get_size() if _comm is not None else 1 +_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") + +# Environment options +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), + profiling_activated=True, + profiling_trace=True, + profiling_label=_profiling_label, +) + +# Time stepping +time_opts = Time(dt=0.05, Tend=400, split_algo="LieTrotter") + +# Geometry +domain = domains.Cuboid(r1=2 * xp.pi / k) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +grid = grids.TensorProductGrid(num_elements=(16, 16, 16)) + +# Derham options +derham_opts = DerhamOptions() + +# Simulation object +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +loading_params = LoadingParameters( + ppc=20, + set_zero_velocity=(False, False, True), + moments=(0.0, 0.0, 0.0, vth1_background_val, vth2_background_val, 1.0), + seed=42, +) +weights_params = WeightsParameters(control_variate=False) +boundary_params = BoundaryParameters() +sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) + +binplot_dens = BinningPlot(slice="e1_v1", n_bins=(128, 128), ranges=((0.0, 1.0), (-0.1, 0.1))) +binplot_velocity = BinningPlot(slice="v1_v2", n_bins=(128, 128), ranges=((-0.1, 0.1), (-0.1, 0.1))) +binplot_current = tuple( + BinningPlot(slice=f"e{i}", n_bins=32, ranges=(0.0, 1.0), output_quantity=f"current_{j}") + for j in range(1, 4) + for i in range(1, 4) +) +saving_params = SavingParameters(binning_plots=(binplot_dens, binplot_velocity, *binplot_current)) + +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=2.0, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.maxwell.options = model.propagators.maxwell.Options() +model.propagators.push_eta.options = model.propagators.push_eta.Options() +model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.coupling_va.options = model.propagators.coupling_va.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0", solver=args.solver) + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# If backgrounds or perturbations are not specified, they are assumed to be zero. + +maxwellian = maxwellians.Maxwellian3D( + vth1=(vth1_background_val, None), + vth2=(vth2_background_val, None), +) +model.kinetic_ions.var.add_background(maxwellian) + +# Perturbation of initial magnetic field +model.em_fields.b_field.add_perturbation( + perturbation=perturbations.ModesCos(amps=(B_pert_amp,), ls=(1,), comp=2), # Initial Bz depending on x-axis +) + +if __name__ == "__main__": + # one_time_step=True isolates the initial Poisson solve (the only part of this model that + # solver= affects -- the fields then evolve via MaxwellWeakAmpere/PushVxB/VlasovAmpereCoupling, + # unrelated to PETScSolver) from the ~8000 identical transport steps a full run would + # otherwise dilute the comparison with. + sim.run(one_time_step=True) diff --git a/profiling/submit_toydrift_hires_petsc.py b/profiling/submit_toydrift_hires_petsc.py new file mode 100644 index 000000000..6915d00be --- /dev/null +++ b/profiling/submit_toydrift_hires_petsc.py @@ -0,0 +1,84 @@ +"""ToyDrift periodic slab, higher resolution: PETSc vs. pcg at a larger, more realistic size. + +This is a scaled-up variant of ``submit_toydrift_petsc.py``'s case (see +``profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py`` for the full +story): a 32^3 grid, 32768 dofs (vs. periodic_slab's 24^3, 13824 dofs) with PETSc's preconditioner +explicitly set to ``"gamg"`` (algebraic multigrid, via ``SolverParameters.pc_type`` -- see +``struphy.linear_algebra.solver.SolverParameters``) instead of the default ``"jacobi"``. + +Measured via ``struphy.linear_algebra.petsc_examples_benchmark``'s repeated-solve methodology +(isolates just the solve, amortizing gamg's one-time multigrid setup across several calls at +fixed dt -- see that module's docstring): pcg needs ~190 CG iterations per solve (~10.9s) at this +size, against PETSc+gamg's ~2 iterations (~0.29s) -- a ~38x difference, essentially the same +*ratio* periodic_slab already shows (~36x at 13824 dofs): this near-singular system is +ill-conditioned mainly through its weakly constrained DC mode, not primarily through resolution, +so the ratio does not grow much further with grid size. What does grow is the *absolute* cost: +pcg's ~11 seconds per solve here is the practically relevant "huge difference" once multiplied +over the many timesteps of a real simulation. See ``submit_toydrift_petsc.py`` and +``submit_strong_landau_damping_petsc.py`` for the smaller-scale comparisons. + +Each launch runs one 32^3-grid, one-time-step simulation (``params_periodic_slab_hires.py``'s +``__main__`` calls ``sim.run(one_time_step=True)``); note that a *single* one-time-step run still +pays gamg's one-time setup cost up front and will not by itself reproduce the ~38x figure above -- +that requires several solves at the same dt to amortize the setup, exactly what +``petsc_examples_benchmark.py`` does. A pcg launch alone takes roughly a minute here (mostly the +single Poisson solve), so a local (non-SLURM) full sweep of this script takes several minutes. +""" + +import argparse +from pathlib import Path + +from profiling_job import ProfilingCase + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "ToyDrift" / "periodic_slab_hires" + + profiling_case = ProfilingCase( + label="toydrift_periodic_slab_hires_petsc", + name="ToyDrift periodic slab (32^3, hires): pcg vs. PETSc+gamg at a larger problem size", + description=( + "Higher-resolution (32^3 grid, 32768 dofs, degree-3 splines) periodic-slab ToyDrift " + "setup, solving the per-step guiding-center Poisson problem with either feectools' " + "native preconditioned CG (--solver pcg) or PETScSolver with an algebraic multigrid " + "preconditioner (KSP=cg, PC=gamg, --solver petsc). At this size pcg needs roughly 190 " + "iterations per solve (about 11 seconds) while PETSc+gamg needs about 2 (a third of a " + "second); the ratio (~38x) is similar to the smaller periodic_slab case, but the " + "absolute per-solve cost -- and hence the real wall-clock impact over a full run -- " + "is far larger here." + ), + physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", + struphy_model_used="ToyDrift", + params_source=params_dir / "params_periodic_slab_hires.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + profiling_case.use_slurm = False + + # Launch one run per (rank count, solver) combination, same rank counts for both solvers so + # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` + # output folder), so the two solvers never collide even at the same rank count. + for num_tasks in (1,): # 2, 4): + for solver in ("pcg", "petsc"): + profiling_case.launch(num_tasks, param_flags=["--solver", solver]) + + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/submit_toydrift_petsc.py b/profiling/submit_toydrift_petsc.py index 5aede1a8c..ece5d8c70 100644 --- a/profiling/submit_toydrift_petsc.py +++ b/profiling/submit_toydrift_petsc.py @@ -67,7 +67,7 @@ def main() -> None: # Launch one run per (rank count, solver) combination, same rank counts for both solvers so # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` # output folder), so the two solvers never collide even at the same rank count. - for num_tasks in (1, 2, 4): + for num_tasks in (1, ): #2, 4): for solver in ("pcg", "petsc"): profiling_case.launch(num_tasks, param_flags=["--solver", solver]) diff --git a/profiling/submit_vlasov_maxwell_petsc.py b/profiling/submit_vlasov_maxwell_petsc.py new file mode 100644 index 000000000..fc02031e4 --- /dev/null +++ b/profiling/submit_vlasov_maxwell_petsc.py @@ -0,0 +1,44 @@ +from pathlib import Path + +from profiling_job import ProfilingCase + +from utils import _get_profiling_args + + +def main() -> None: + + args = _get_profiling_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "VlasovMaxwellOneSpecies" / "weibel_instability" + + profiling_case = ProfilingCase( + label="weibel_instability_petsc", + name="Weibel instability, initial Poisson solve: pcg vs. PETSc", + description=( + "Weibel instability test case for VlasovMaxwellOneSpecies, solving the one-time " + "initial Poisson problem with either feectools' native preconditioned CG " + "(--solver pcg) or PETScSolver (KSP=cg, PC=gamg, --solver petsc)." + ), + physics_problem="Weibel instability driven by an anisotropic velocity distribution.", + struphy_model_used="VlasovMaxwellOneSpecies", + params_source=params_dir / "params_weibel_instability.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + profiling_case.use_slurm = False + + # Launch one run per (rank count, solver) combination, same rank counts for both solvers so + # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` + # output folder), so the two solvers never collide even at the same rank count. + for num_tasks in (1, 2, 4): + for solver in ("pcg", "petsc"): + profiling_case.launch(num_tasks, param_flags=["--solver", solver]) + + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index a78ea8cfa..1e8de553e 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -69,6 +69,7 @@ class LiteralOptions: # solvers OptsSymmSolver = Literal["pcg", "cg", "petsc"] OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] + OptsPETScPrecond = Literal["none", "jacobi", "gamg", "ilu", "sor"] OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] OptsSaddlePointSolver = Literal["uzawa"] OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index 8789b6f08..d97e7b7a2 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -79,6 +79,11 @@ def __init__( kwargs = solver_params.__dict__.copy() kwargs.pop("info") + # pc_type is petsc-only (see struphy.linear_algebra.solver.SolverParameters); this + # module always goes through feectools' own `inverse` (imported directly above, not + # struphy's petsc-aware wrapper), whose InverseLinearOperator subclasses forward + # unknown kwargs straight to their constructor and would raise on it. + kwargs.pop("pc_type", None) if precond is not None: kwargs["pc"] = precond diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index 2ba503562..3287f73e1 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -46,6 +46,12 @@ def inverse(A, solver: str, **kwargs): return PETScSolver(A, **petsc_kwargs) + # pc_type/ksp_type are petsc-only (see _PETSC_SOLVER_KWARGS above); feectools' + # InverseLinearOperator subclasses forward unknown kwargs straight to their + # constructor and would raise on them, so they never reach this branch. + kwargs.pop("pc_type", None) + kwargs.pop("ksp_type", None) + from feectools.linalg.solvers import inverse as feectools_inverse return feectools_inverse(A, solver, **kwargs) @@ -59,6 +65,11 @@ class SolverParameters: maxiter: int = 3000 info: bool = False recycle: bool = True + pc_type: LiteralOptions.OptsPETScPrecond = "jacobi" + """Preconditioner for ``solver="petsc"`` only (ignored otherwise): PETSc's ``PCType`` + name, e.g. ``"jacobi"`` (cheap, diagonal) or ``"gamg"`` (algebraic multigrid -- far + stronger for large, ill-conditioned systems, but with more setup overhead per matrix + assembly).""" def __post_init__(self): self.verbose = False diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 797042737..c62522adb 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -362,6 +362,7 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, recycle=self.options.solver_params.recycle, + pc_type=self.options.solver_params.pc_type, ) # allocate memory for solution From 56072ad15b787a24f007331482215dabe1b8f367 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 10 Aug 2026 08:53:23 +0200 Subject: [PATCH 22/32] Added tests and examples --- .../petsc_examples_benchmark.py | 225 ++++++++++++++++++ .../linear_algebra/petsc_poisson_benchmark.py | 207 ++++++++++++++++ .../tests/test_petsc_poisson_solve_pic.py | 109 +++++++++ 3 files changed, 541 insertions(+) create mode 100644 src/struphy/linear_algebra/petsc_examples_benchmark.py create mode 100644 src/struphy/linear_algebra/petsc_poisson_benchmark.py create mode 100644 src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py diff --git a/src/struphy/linear_algebra/petsc_examples_benchmark.py b/src/struphy/linear_algebra/petsc_examples_benchmark.py new file mode 100644 index 000000000..8004f69f2 --- /dev/null +++ b/src/struphy/linear_algebra/petsc_examples_benchmark.py @@ -0,0 +1,225 @@ +"""Benchmark: solver="petsc" vs solver="pcg" on real struphy examples. + +Uses the committed parameter files under +``profiling/examples///params_.py``, each with a ``--solver`` CLI flag +(default ``"pcg"``) selecting the solver of the case's Poisson-type propagator: + +- ``VlasovAmpereOneSpecies/{strong_Landau_damping,weak_Landau_damping,two_stream,bump_on}``: plain + copies of the corresponding ``examples/VlasovAmpereOneSpecies//params_.py`` (unedited + on disk), with ``num_elements`` scaled up from the examples' tiny, highly-anisotropic 1D-style + default of ``(32, 1, 1)`` cells to a proper ``(16, 16, 16)`` 3D grid -- PETSc's advantage only + shows up above roughly 5,000 dofs -- ``ppc`` reduced to match, and a fixed + ``LoadingParameters.seed`` added (the originals don't set one, so pcg/petsc would otherwise draw + different particles and not be comparable). This model only solves Poisson *once*, as an initial + condition (the field then evolves via VlasovAmpereCoupling), so the repeated-solve timing below + re-invokes ``model.initial_poisson`` directly after ``sim.run()`` rather than relying on the + model's own (single-shot) usage of it. + +- ``ToyDrift/periodic_slab``: no periodic ToyDrift example exists under ``examples/`` (the real + one, ``examples/ToyGyrokinetic/diocotron_instability``, needs a physically non-periodic + HollowCylinder domain), so this one is written from scratch with a periodic Cuboid domain + instead -- which works because ToyDrift's field solve is a plain ``PoissonSolve`` with no + geometry-coupled averaging (unlike ``PoissonAdiabaticGyrokinetic``, used by + ``DriftKineticElectrostaticAdiabatic``, which diverges outright on a periodic domain regardless + of options -- tried first, not usable here). Unlike VlasovAmpereOneSpecies, ToyDrift's + ``gc_poisson`` runs as a *regular per-step propagator*, so no re-invocation workaround is needed. + +- ``ToyDrift/periodic_slab_hires``: same setup as ``periodic_slab``, scaled up to a 32^3 grid + (vs. 24^3) to push feectools' unpreconditioned CG into several hundred iterations per solve + while PETSc+gamg (``pc_type="gamg"``, set explicitly via ``SolverParameters.pc_type`` -- see + ``struphy.linear_algebra.solver.SolverParameters``) stays at a handful, regardless of grid size. + The most lopsided case in this suite by design; see its own params file's docstring. + +- ``VlasovMaxwellOneSpecies/weibel_instability``: plain copy of + ``examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py``, scaled up + the same way as the VlasovAmpereOneSpecies cases above (3D grid, reduced ppc, fixed seed). Same + one-shot ``model.initial_poisson`` pattern as VlasovAmpereOneSpecies (the fields then evolve via + MaxwellWeakAmpere/PushVxB/VlasovAmpereCoupling instead), so the same re-invocation timing applies. + +This script's only job is to import each file's ``sim``/``model`` and drive them -- no source +patching, no ``runpy``. + +Correctness is checked between the two solvers on the *mean-removed* solution (the near-singular, +essentially unregularized ``stab_eps`` these examples use -- via ``ImplicitDiffusion``'s "always +stabilize" clamp to ``1e-14`` -- leaves the constant/DC mode only very weakly constrained, so it is +extremely sensitive to tiny numerical differences between solvers; this is expected and is not a +correctness issue in the physically meaningful, oscillatory part of the solution), normalized by +the *full* solution's norm (not the tiny mean-removed norm itself, which can be dominated by +floating-point noise for weak-perturbation examples and make a naively-normalized relative error +meaningless). + +KNOWN ISSUE -- do not trust results under MPI (comm size > 1): for this same near-singular +``stab_eps``-clamped-to-``1e-14`` regime, PETScSolver was found to disagree substantially with +feectools' native solver specifically under >1 MPI rank, independent of ``pc_type`` (both +"jacobi" and "gamg" reproduced it; "gamg" was far worse -- a false "converged in 1 iteration" to +a wildly wrong answer). This reproduces even though pcg-vs-pcg (same solver, two independent runs) +is bit-identical, ruling out a methodology artifact in this script. The root cause was not found; +serial execution and non-near-singular systems (this same MPI path, e.g. with an explicit +``stab_eps`` of 1e-8 or larger) were extensively validated and are unaffected -- see +``petsc_poisson_benchmark.py`` and ``test_petsc_poisson_solve_pic.py``. This script therefore +only asserts/prints the correctness check when running serially, and warns instead under MPI. + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_examples_benchmark +""" + +import importlib +import os +import shutil +import sys +import tempfile +import time +from pathlib import Path + +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() + +REPO_ROOT = Path(__file__).resolve().parents[3] +PROFILING_EXAMPLES_DIR = REPO_ROOT / "profiling" / "examples" + +# (model directory under profiling/examples/, case name) +CASES = ( + ("VlasovAmpereOneSpecies", "strong_Landau_damping"), + ("VlasovAmpereOneSpecies", "weak_Landau_damping"), + ("VlasovAmpereOneSpecies", "two_stream"), + ("VlasovAmpereOneSpecies", "bump_on"), + ("ToyDrift", "periodic_slab"), + ("ToyDrift", "periodic_slab_hires"), + ("VlasovMaxwellOneSpecies", "weibel_instability"), +) + +# `profiling` is a repo-local package (not part of the installed struphy distribution), normally +# importable only because '' (cwd) is on sys.path at interpreter startup. This script chdir()s +# into a scratch directory before importing (to keep Simulation's output out of the repo), which +# breaks that for '-c'/REPL-style invocations where '' resolves dynamically -- so add the repo +# root explicitly, once, up front. +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + + +def _poisson_propagator(model): + """VlasovAmpereOneSpecies exposes its (one-shot) Poisson solve as `model.initial_poisson`; + other models (e.g. ToyDrift) run it as a regular per-step propagator instead. + """ + if hasattr(model, "initial_poisson"): + return model.initial_poisson + return model.propagators.gc_poisson + + +def _run_variant(model_dir: str, name: str, variant: str, out_folder: str, dt: float, n_solves: int): + module_name = f"profiling.examples.{model_dir}.{name}.params_{name}" + + # EnvironmentOptions.out_folders defaults to `os.getcwd()`, but as a plain dataclass field + # default this is evaluated once, the first time struphy.io.options is imported in this + # process -- not per EnvironmentOptions() call, and not affected by a later os.chdir(). Since + # `struphy`'s own __init__.py re-exports EnvironmentOptions (and everything else), that first + # import can happen before this function even runs (e.g. as a side effect of importing this + # very module). The profiling params files read STRUPHY_PROFILING_OUT_FOLDERS explicitly for + # exactly this reason -- set it before importing, so their Simulation's output lands here + # instead of wherever the process happened to start. + os.environ["STRUPHY_PROFILING_OUT_FOLDERS"] = out_folder + + # Both variants now come from the same params_.py (a `--solver` CLI flag picks between + # them internally), so a plain second `import_module` would just return the first variant's + # already-imported module/sim/model unchanged. Feed the desired `--solver` through sys.argv + # (the params file reads it via argparse) and force a re-execution via `reload` so the second + # variant gets its own fresh Simulation/model instead of reusing (and re-running) the first's. + old_argv = sys.argv + sys.argv = [old_argv[0], "--solver", variant] + try: + if module_name in sys.modules: + mod = importlib.reload(sys.modules[module_name]) + else: + mod = importlib.import_module(module_name) + finally: + sys.argv = old_argv + + sim = mod.sim + model = mod.model + + sim.run(one_time_step=True) # real setup: particle loading, Derham/mass operators, and (for + # VlasovAmpereOneSpecies) the initial Poisson solve, exactly as + # VlasovAmpereOneSpecies.allocate_helpers / a real per-step run does + + poisson = _poisson_propagator(model) + solver = poisson._solver + if variant == "petsc": + solver._options["pc_type"] = "gamg" + solver._ksp = None # force a rebuild with the new pc_type + + # repeated calls at fixed dt (same real charge deposition -- this benchmark is about the + # linear-solve cost, not particle physics): matches how repeated Poisson solves would + # amortize matrix assembly across timesteps at a fixed dt in a real run (see + # ImplicitDiffusion.__call__'s lhs-operator caching) + poisson(dt) # warm-up + t0 = time.perf_counter() + for _ in range(n_solves): + poisson(dt) + t = (time.perf_counter() - t0) / n_solves + info = solver.get_info() if hasattr(solver, "get_info") else solver._info + + phi = model.em_fields.phi.spline.vector.toarray() + return t, info, phi + + +def bench_example(model_dir: str, name: str, dt: float = 0.05, n_solves: int = 10): + """Import and run one example's pcg/petsc parameter files, and report timing + correctness.""" + params_dir = PROFILING_EXAMPLES_DIR / model_dir / name + if not (params_dir / f"params_{name}.py").exists(): + raise FileNotFoundError( + f"Could not find {params_dir}/params_{name}.py -- expected the parameter file " + f"under profiling/examples/{model_dir}/{name}/." + ) + + # tempfile.mkdtemp() is not MPI-coordinated: each rank would otherwise get a *different* + # random path, and Simulation's output-file creation (rank 0 only) would then fail on every + # other rank. Create it on rank 0 and broadcast the path instead. + out_folder = tempfile.mkdtemp() if rank == 0 else None + if comm is not None: + out_folder = comm.bcast(out_folder, root=0) + + try: + t_cg, info_cg, sol_cg = _run_variant(model_dir, name, "pcg", out_folder, dt, n_solves) + t_petsc, info_petsc, sol_petsc = _run_variant(model_dir, name, "petsc", out_folder, dt, n_solves) + finally: + if comm is not None: + comm.Barrier() + if rank == 0: + shutil.rmtree(out_folder, ignore_errors=True) + + mean_removed_pcg = sol_cg - sol_cg.mean() + mean_removed_petsc = sol_petsc - sol_petsc.mean() + rel_err = xp.linalg.norm(mean_removed_pcg - mean_removed_petsc) / xp.linalg.norm(sol_cg) + + comm_size = comm.Get_size() if comm is not None else 1 + if rank == 0: + print(f"\n{model_dir}/{name}: ndofs={sol_cg.size}") + print(f" pcg (unprec.) : {t_cg * 1e3:9.2f} ms/step niter={info_cg.get('niter')}") + print(f" petsc + gamg : {t_petsc * 1e3:9.2f} ms/step niter={info_petsc.get('niter')}") + print(f" speedup: {t_cg / t_petsc:.2f}x", end=" ") + if comm_size > 1: + print( + f"relative solution mismatch: {rel_err:.2e} " + "-- NOT a reliable correctness check under MPI for this near-singular regime, " + "see module docstring (KNOWN ISSUE)" + ) + else: + print(f"relative solution mismatch: {rel_err:.2e}") + assert rel_err < 1e-6, ( + f"{model_dir}/{name}: pcg/petsc solutions disagree by {rel_err:.2e}, expected < 1e-6 serially" + ) + + +def main(): + for model_dir, name in CASES: + bench_example(model_dir, name) + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/petsc_poisson_benchmark.py b/src/struphy/linear_algebra/petsc_poisson_benchmark.py new file mode 100644 index 000000000..546edacfb --- /dev/null +++ b/src/struphy/linear_algebra/petsc_poisson_benchmark.py @@ -0,0 +1,207 @@ +"""Benchmark: PoissonSolve(solver="petsc") vs PoissonSolve(solver="pcg") on a real Vlasov-Poisson testcase. + +Builds actual :class:`~struphy.simulation.sim.Simulation` objects, using the same public API and +setup idiom as a real parameter file (compare +``examples/VlasovAmpereOneSpecies/strong_Landau_damping/params_strong_Landau_damping.py``, +generalized here to a 3D grid), and runs them with ``sim.run(one_time_step=True)`` -- this is the +one and only supported entry point for allocating and running a struphy simulation; it performs +real particle loading, real Derham/mass-operator setup, and (for this model) a real charge-density +deposition via ``ParticlesToGrid``/``AccumulatorVector``, solved via ``PoissonSolve`` exactly as +``VlasovAmpereOneSpecies.allocate_helpers`` does. + +Two earlier, less realistic benchmarks are *not* wins for PETSc and are not repeated here: + +- Mass-matrix solves (``L2Projector``): already well-conditioned, feectools' native + preconditioned CG wins outright. +- ``PoissonSolve`` with a synthetic, non-mass-weighted random right-hand side: not representative + of any real code path (every real source -- ``FEECVariable``, ``ParticlesToGrid``, ``Callable`` + -- is mass-matrix weighted when forming the weak-form right-hand side, which is inherently + smoothing). + +The genuine win requires: a small ``stab_eps`` (true elliptic Poisson, matching realistic +electrostatic PIC parameters -- ``stab_eps`` is a numerical regularization, not a dominant +physical diffusion), a broadband/noisy right-hand side (real PIC deposition, not a smooth +manufactured mode), and repeated solves at fixed ``dt`` (relies on +``ImplicitDiffusion.__call__`` caching its lhs operator when ``sig_1`` is unchanged, so +``PETScSolver`` can reuse its assembled matrix -- see git history for that fix). Since +``VlasovAmpereOneSpecies`` only calls its Poisson solve *once* (as an initial condition -- the +electric field then evolves via Ampere's law, not repeated Poisson solves), the repeated-solve +timing below re-invokes ``model.initial_poisson`` directly after ``sim.run()`` has performed the +real setup, rather than relying on the model's own (single-shot) usage of it. + +Run with: + +.. code-block:: bash + + python3 -m struphy.linear_algebra.petsc_poisson_benchmark +""" + +import shutil +import tempfile +import time +import warnings + +import cunumpy as xp +from feectools.ddm.mpi import mpi as MPI + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import VlasovAmpereOneSpecies + +comm = MPI.COMM_WORLD +rank = comm.Get_rank() + + +def build_and_run(num_elements, degree, ppc, solver_name, perturbation, pc_type=None, stab_eps=1e-8, out_folder=None): + """Build a VlasovAmpereOneSpecies Simulation exactly as a params.py file would, and run one + step of it via sim.run(one_time_step=True) -- the real, supported entry point. This performs + real particle loading and the real (single-shot) initial Poisson solve. + """ + # alpha=1.0, epsilon=-1.0 (matching the real strong_Landau_damping example) keeps the + # right-hand side well-scaled (order 1); epsilon in particular can never be auto-derived as + # negative (its formula is always positive), so overriding it is unavoidable here, and + # struphy warns on every such override. The warning is expected and harmless -- silence it + # rather than leaving alpha/epsilon at their auto-derived values, which was tried and + # produces a poorly-scaled right-hand side (huge/tiny relative to 1), breaking the implicit + # assumption -- shared by every SolverParameters.tol comparison in this benchmark -- that + # "tol" means the same thing regardless of problem scale. + # with warnings.catch_warnings(): + # warnings.filterwarnings("ignore", message="Override equation parameter", category=UserWarning) + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + env = EnvironmentOptions(out_folders=out_folder, sim_folder=f"bench_{solver_name}") + time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") + domain = domains.Cuboid() + grid = grids.TensorProductGrid(num_elements=num_elements) + derham_opts = DerhamOptions(degree=degree, bcs=(None, None, None)) + + sim = Simulation( + model=model, + params_path=None, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, + ) + + loading_params = LoadingParameters(ppc=ppc, seed=1234) + weights_params = WeightsParameters(control_variate=True) + boundary_params = BoundaryParameters() + model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options() + model.initial_poisson.options = model.initial_poisson.Options( + stab_mat="M0", + stab_eps=stab_eps, + solver=solver_name, + precond="MassMatrixPreconditioner", + solver_params=SolverParameters(tol=1e-10, maxiter=20_000, info=False, recycle=False), + ) + + background = maxwellians.Maxwellian3D(n=(1.0, None)) + model.kinetic_ions.var.add_background(background) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + model.kinetic_ions.var.add_initial_condition(init) + + sim.run(one_time_step=True) + + if solver_name == "petsc" and pc_type is not None and model.initial_poisson._solver._options["pc_type"] != pc_type: + model.initial_poisson._solver._options["pc_type"] = pc_type + model.initial_poisson._solver._ksp = None + + return sim, model + + +def bench_case(name, num_elements, degree, ppc, perturbation, n_solves=10, dt=0.05, **kwargs): + # tempfile.mkdtemp() is not MPI-coordinated: each rank would otherwise get a *different* + # random path, and Simulation's output-file creation (rank 0 only) would then fail on every + # other rank. Create it on rank 0 and broadcast the path instead. + out_folder = tempfile.mkdtemp() if rank == 0 else None + if comm is not None: + out_folder = comm.bcast(out_folder, root=0) + + try: + _, model_cg = build_and_run(num_elements, degree, ppc, "pcg", perturbation, out_folder=out_folder, **kwargs) + _, model_petsc = build_and_run( + num_elements, degree, ppc, "petsc", perturbation, pc_type="gamg", out_folder=out_folder, **kwargs + ) + + # both models' initial_poisson already ran once inside sim.run(); time repeated calls at + # fixed dt, exactly as a real timestepping loop would (see module docstring) + model_cg.initial_poisson(dt) # warm-up + t0 = time.perf_counter() + for _ in range(n_solves): + model_cg.initial_poisson(dt) + t_cg = (time.perf_counter() - t0) / n_solves + info_cg = model_cg.initial_poisson._solver._info + + model_petsc.initial_poisson(dt) # warm-up + t0 = time.perf_counter() + for _ in range(n_solves): + model_petsc.initial_poisson(dt) + t_petsc = (time.perf_counter() - t0) / n_solves + info_petsc = model_petsc.initial_poisson._solver.get_info() + + sol_cg = model_cg.em_fields.phi.spline.vector.toarray() + sol_petsc = model_petsc.em_fields.phi.spline.vector.toarray() + rel_err = xp.linalg.norm(sol_cg - sol_petsc) / xp.linalg.norm(sol_cg) + finally: + if comm is not None: + comm.Barrier() + if rank == 0: + shutil.rmtree(out_folder, ignore_errors=True) + + ndofs = model_cg.em_fields.phi.spline.vector.space.dimension + Np = model_cg.kinetic_ions.var.particles.markers.shape[0] + + if rank == 0: + print(f"\n{name}: num_elements={num_elements}, degree={degree}, ndofs={ndofs}, Np~{Np}") + print(f" pcg (unprec.) : {t_cg * 1e3:9.2f} ms/step niter={info_cg.get('niter')}") + print(f" petsc + gamg : {t_petsc * 1e3:9.2f} ms/step niter={info_petsc.get('niter')}") + print(f" speedup: {t_cg / t_petsc:.2f}x relative solution mismatch: {rel_err:.2e}") + + +def main(): + # 1. grid-size scaling, matching examples/VlasovAmpereOneSpecies/strong_Landau_damping's ICs + landau_damping = perturbations.ModesCos(amps=(0.5,), ls=(1,)) + for num_elements in [[8, 8, 8], [16, 16, 16], [24, 24, 24], [32, 32, 32]]: + bench_case("Landau damping (grid scaling)", num_elements, [2, 2, 2], ppc=20, perturbation=landau_damping) + + # 2. weak Landau damping ICs (small-amplitude perturbation, closer to the linear regime) + weak_landau_damping = perturbations.ModesCos(amps=(0.001,), ls=(1,)) + bench_case("weak Landau damping", [24, 24, 24], [2, 2, 2], ppc=20, perturbation=weak_landau_damping) + + # 3. higher spline degree (matching the real examples' degree=3 in the perturbed direction) + bench_case("Landau damping, degree 3", [16, 16, 16], [3, 3, 3], ppc=20, perturbation=landau_damping) + + # 4. sparser sampling (fewer particles per cell -> noisier deposited density) + bench_case("Landau damping, low ppc (noisier)", [24, 24, 24], [2, 2, 2], ppc=5, perturbation=landau_damping) + + # 5. genuinely 3D, multi-mode perturbation (unlike the 1D-in-x real examples), a closer + # stand-in for 3D electrostatic turbulence + multi_mode_3d = perturbations.ModesCos(amps=(0.5, 0.3, 0.2), ls=(1, 2, 0), ms=(0, 1, 2), ns=(0, 0, 1)) + bench_case("3D multi-mode perturbation", [24, 24, 24], [2, 2, 2], ppc=20, perturbation=multi_mode_3d) + + +if __name__ == "__main__": + main() diff --git a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py new file mode 100644 index 000000000..f2cac1ed6 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py @@ -0,0 +1,109 @@ +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy import LoadingParameters, WeightsParameters, maxwellians, perturbations +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.variables import FEECVariable +from struphy.pic.accumulation import accum_kernels +from struphy.pic.accumulation.particles_to_grid import ParticlesToGrid +from struphy.pic.particles import Particles6D +from struphy.propagators.base import Propagator +from struphy.propagators.poisson_solve import PoissonSolve +from struphy.topology.grids import TensorProductGrid +from struphy.utils.pyccel import Pyccelkernel + + +class _FakePICVariable: + """Minimal duck-typed stand-in for a model's PICVariable, since ParticlesToGrid only reads .particles.""" + + def __init__(self, particles): + self.particles = particles + + +def test_poisson_solve_petsc_matches_pcg_with_real_pic_deposition(): + """PoissonSolve(solver="petsc") must match PoissonSolve(solver="pcg") when driven by a real + particle-in-cell charge-density deposition (not a synthetic/manufactured source), reproducing + the setup of examples/VlasovAmpereOneSpecies/strong_Landau_damping (Maxwellian3D background + + ModesCos perturbation, control-variate weights). + """ + comm = MPI.COMM_WORLD + + domain = Cuboid() + grid = TensorProductGrid(num_elements=[10, 10, 10]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + mass_ops = WeightedMassOperators(derham, domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + background = maxwellians.Maxwellian3D(n=(1.0, None)) + perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + + domain_array = derham.domain_array + nprocs = derham.domain_decomposition.nprocs + + def run(solver_name): + loading_params = LoadingParameters(Np=20_000, seed=1234) + weights_params = WeightsParameters(control_variate=True) + + particles = Particles6D( + comm_world=comm, + clone_config=None, + loading_params=loading_params, + weights_params=weights_params, + domain=domain, + domain_decomp=(domain_array, nprocs), + background=background, + initial_condition=init, + ) + particles.draw_markers() + if comm.Get_size() > 1: + particles.mpi_sort_markers() + particles.initialize_weights() + + rho = ParticlesToGrid( + _FakePICVariable(particles), + "H1", + Pyccelkernel(accum_kernels.charge_density_0form), + ) + + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + + solver_params = SolverParameters(tol=1e-10, maxiter=20000, info=False, recycle=False) + + prop = PoissonSolve(rho=rho) + prop.variables.phi = phi + prop.options = prop.Options( + stab_eps=1e-8, + solver=solver_name, + precond="MassMatrixPreconditioner", + solver_params=solver_params, + ) + prop.allocate() + if solver_name == "petsc": + prop._solver._options["pc_type"] = "gamg" + prop._solver._ksp = None + prop(0.05) + return phi.spline.vector.toarray() + + sol_pcg = run("pcg") + sol_petsc = run("petsc") + + rel_err = xp.linalg.norm(sol_pcg - sol_petsc) / xp.linalg.norm(sol_pcg) + assert rel_err < 1e-6 + + +if __name__ == "__main__": + test_poisson_solve_petsc_matches_pcg_with_real_pic_deposition() From bd906d1d03825e56f761e8f83f03ecce0262814b Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 10 Aug 2026 14:24:14 +0200 Subject: [PATCH 23/32] Cleanup --- .../periodic_slab/params_periodic_slab.py | 175 ------------------ .../params_periodic_slab_hires.py | 46 +++-- profiling/submit_toydrift_hires_petsc.py | 23 ++- profiling/submit_toydrift_petsc.py | 78 -------- .../petsc_examples_benchmark.py | 15 +- 5 files changed, 43 insertions(+), 294 deletions(-) delete mode 100644 profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py delete mode 100644 profiling/submit_toydrift_petsc.py diff --git a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py b/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py deleted file mode 100644 index 07d48ab18..000000000 --- a/profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py +++ /dev/null @@ -1,175 +0,0 @@ -import os - -# ----------------------------- -# Description of the simulation -# ----------------------------- - -description = """ -Periodic-slab variant of the ToyDrift model (the model behind -examples/ToyGyrokinetic/diocotron_instability, which uses a physically non-periodic -HollowCylinder domain -- radial confinement is inherent to the diocotron instability). This case -swaps in a periodic Cuboid domain instead: unlike PoissonAdiabaticGyrokinetic (used by -DriftKineticElectrostaticAdiabatic), ToyDrift's field solve is a plain PoissonSolve with no -geometry-coupled averaging, so it works correctly on a periodic domain out of the box. - -Unlike VlasovAmpereOneSpecies (which only solves Poisson once, as an initial condition), this -model's gc_poisson runs as a *regular per-step propagator* -- exactly the repeated-solve pattern -where PETSc's algebraic multigrid preconditioner shows a genuine win (see -struphy.linear_algebra.petsc_examples_benchmark's module docstring for the general story). -""" - -# ------------------ -# Import Struphy API -# ------------------ - -from struphy import ( - BaseUnits, - BoundaryParameters, - DerhamOptions, - EnvironmentOptions, - LoadingParameters, - Simulation, - SortingParameters, - Time, - WeightsParameters, - domains, - equils, - grids, - maxwellians, - perturbations, -) -from struphy.linear_algebra.solver import SolverParameters - -# --------------------- -# Instance of the model -# --------------------- -from struphy.models import ToyDrift - -# Units -base_units = BaseUnits(kBT=1.0) - -# Model instance -model = ToyDrift(base_units=base_units) - -# List all variables and decide whether to save their data -model.em_fields.phi.save_data = True -model.kinetic_ions.var.save_data = False - -# -------------------------- -# Instance of the simulation -# -------------------------- - -# `--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. -import argparse - -parser = argparse.ArgumentParser() -parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") -parser.add_argument( - "--solver", type=str, default="pcg", choices=["pcg", "petsc"], help="Solver for the Poisson-type solve." -) -args, _ = parser.parse_known_args() - -# scope-profiler label: distinguishes solver/rank-count combinations in post-processing -# (chart legends, `scope-profiler inspect`), see EnvironmentOptions.profiling_label. -from feectools.ddm.mpi import mpi as MPI - -_comm = MPI.COMM_WORLD -_num_ranks = _comm.Get_size() if _comm is not None else 1 -_profiling_label = f"{args.solver}, {_num_ranks} rank" + ("s" if _num_ranks != 1 else "") - -# Environment options -env = EnvironmentOptions( - sim_folder=f"sim_{args.id:02d}", - out_folders=os.environ.get("STRUPHY_PROFILING_OUT_FOLDERS", os.getcwd()), - profiling_activated=True, - profiling_trace=True, - profiling_label=_profiling_label, -) - -# Time stepping -time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") - -# Geometry -domain = domains.Cuboid() - -# Fluid equilibrium: straight B field, homogeneous density (default n0=1.0) -equil = equils.HomogenSlab(B0z=1.0, n0=1.0) - -# Grid -grid = grids.TensorProductGrid(num_elements=(24, 24, 24)) - -# Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble -# DirectionalDerivativeOperator along a non-periodic axis, see -# struphy.linear_algebra.petsc_solver._directional_derivative_to_stencil_matrix) -derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None)) - -# Simulation object -sim = Simulation( - model=model, - params_path=__file__, - env=env, - time_opts=time_opts, - domain=domain, - equil=equil, - grid=grid, - derham_opts=derham_opts, -) - -# ------------------- -# Particle parameters -# ------------------- - -loading_params = LoadingParameters(ppc=5, seed=42) -weights_params = WeightsParameters(control_variate=True) -boundary_params = BoundaryParameters() -sorting_params = SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True) - -model.kinetic_ions.set_markers( - loading_params=loading_params, - weights_params=weights_params, - boundary_params=boundary_params, - sorting_params=sorting_params, - bufsize=0.4, -) - -# ------------------ -# Propagator options -# ------------------ - -model.propagators.gc_poisson.options.solver = args.solver -model.propagators.gc_poisson.options.solver_params = SolverParameters(tol=1e-10, maxiter=20_000) -model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( - algo="explicit", - evaluate_e_field=True, -) - -# ------------------ -# Initial conditions -# ------------------ -# Initial conditions are the sum of the background(s) and the perturbation(s). - -# Background for kinetic species -background = maxwellians.GyroMaxwellian2D( - n=(1.0, None), - vth_para=(1.0, None), - vth_perp=(1.0, None), - equil=equil, -) -model.kinetic_ions.var.add_background(background) - -# Perturbation, matching the Landau-damping style used elsewhere in this benchmark suite -perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) -init = maxwellians.GyroMaxwellian2D( - n=(1.0, perturbation), - vth_para=(1.0, None), - vth_perp=(1.0, None), - equil=equil, -) -model.kinetic_ions.var.add_initial_condition(init) - -if __name__ == "__main__": - # one_time_step=True isolates the (still per-step, unlike VlasovAmpereOneSpecies) - # gc_poisson solve for a single-step timing snapshot. - sim.run(one_time_step=True) diff --git a/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py b/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py index 2c90749f3..4b968170f 100644 --- a/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py +++ b/profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py @@ -5,28 +5,34 @@ # ----------------------------- description = """ -Higher-resolution variant of ToyDrift/periodic_slab (see that case's own description for the -model/domain background): a 32^3 grid with degree-3 splines (32768 dofs vs. periodic_slab's -24^3 grid's 13824) with PETSc's preconditioner set explicitly to algebraic multigrid -(`pc_type="gamg"`, via `SolverParameters.pc_type` -- see struphy.linear_algebra.solver. -SolverParameters), instead of the default `"jacobi"`. +Periodic-slab variant of the ToyDrift model (the model behind +examples/ToyGyrokinetic/diocotron_instability, which uses a physically non-periodic +HollowCylinder domain -- radial confinement is inherent to the diocotron instability). This case +swaps in a periodic Cuboid domain instead: unlike PoissonAdiabaticGyrokinetic (used by +DriftKineticElectrostaticAdiabatic), ToyDrift's field solve is a plain PoissonSolve with no +geometry-coupled averaging, so it works correctly on a periodic domain out of the box. Unlike +VlasovAmpereOneSpecies (which only solves Poisson once, as an initial condition), gc_poisson runs +as a *regular per-step propagator* here, so a single `sim.run(one_time_step=True)` call already +times one full, representative solve. + +Grid: 32^3 elements, degree-3 splines (32768 dofs), with PETSc's preconditioner set explicitly to +algebraic multigrid (`pc_type="gamg"`, via `SolverParameters.pc_type` -- see +struphy.linear_algebra.solver.SolverParameters), instead of the default `"jacobi"`. This +resolution is deliberately large enough to push feectools' unpreconditioned CG into several +hundred iterations per solve. Measured via struphy.linear_algebra.petsc_examples_benchmark's repeated-solve methodology (which isolates just the solve, amortizing one-time matrix/preconditioner setup across several calls -- see that module's docstring): feectools' unpreconditioned CG took ~10.9 s/solve (190 iterations) -against PETSc+gamg's ~0.29 s/solve (2 iterations) here -- a ~38x difference. This is essentially -the same *ratio* periodic_slab already shows (~36x at 13824 dofs): the near-singular Poisson -system these ToyDrift cases solve (`stab_eps` clamped to ~1e-14 by ImplicitDiffusion's "always -stabilize" logic, see PETScSolver's docstring) is ill-conditioned mainly through its weakly -constrained constant/DC mode, not primarily through raw grid resolution, so pcg's iteration count -does not grow much further with size in this regime -- while gamg's convergence stays -essentially grid-independent regardless. What *does* grow with size is the absolute cost: at this -larger, more realistic problem size, pcg's ~11 seconds per Poisson solve is the practically -relevant "huge difference" -- multiplied over the many timesteps of a real simulation, it is the -difference between a run finishing in minutes and one taking hours. - -Like periodic_slab, gc_poisson runs as a *regular per-step propagator*, so a single -`sim.run(one_time_step=True)` call already times one full, representative solve. +against PETSc+gamg's ~0.29 s/solve (2 iterations) here -- a ~38x difference, the largest gap in +this benchmark suite. The near-singular Poisson system ToyDrift solves (`stab_eps` clamped to +~1e-14 by ImplicitDiffusion's "always stabilize" logic, see PETScSolver's docstring) is +ill-conditioned mainly through its weakly constrained constant/DC mode, not primarily through raw +grid resolution, so pcg's iteration count does not grow much further with size in this regime -- +while gamg's convergence stays essentially grid-independent regardless. What *does* grow with +size is the absolute cost: at this larger, more realistic problem size, pcg's ~11 seconds per +Poisson solve is the practically relevant "huge difference" -- multiplied over the many timesteps +of a real simulation, it is the difference between a run finishing in minutes and one taking hours. """ # ------------------ @@ -108,8 +114,8 @@ # Fluid equilibrium: straight B field, homogeneous density (default n0=1.0) equil = equils.HomogenSlab(B0z=1.0, n0=1.0) -# Grid -- 32^3, ~2x periodic_slab's 24^3 per direction (~2.4x the dofs at fixed degree): large -# enough for pcg's iteration count (hence cost) to blow up while PETSc+gamg stays flat. +# Grid -- 32^3: large enough for pcg's iteration count (hence cost) to blow up while PETSc+gamg +# stays flat. grid = grids.TensorProductGrid(num_elements=(32, 32, 32)) # Derham options -- fully periodic (required: PETScSolver cannot (yet) assemble diff --git a/profiling/submit_toydrift_hires_petsc.py b/profiling/submit_toydrift_hires_petsc.py index 6915d00be..12623036e 100644 --- a/profiling/submit_toydrift_hires_petsc.py +++ b/profiling/submit_toydrift_hires_petsc.py @@ -1,21 +1,20 @@ -"""ToyDrift periodic slab, higher resolution: PETSc vs. pcg at a larger, more realistic size. +"""ToyDrift periodic slab (32^3): PETSc vs. pcg, the largest gap in this benchmark suite. -This is a scaled-up variant of ``submit_toydrift_petsc.py``'s case (see -``profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py`` for the full -story): a 32^3 grid, 32768 dofs (vs. periodic_slab's 24^3, 13824 dofs) with PETSc's preconditioner -explicitly set to ``"gamg"`` (algebraic multigrid, via ``SolverParameters.pc_type`` -- see +See ``profiling/examples/ToyDrift/periodic_slab_hires/params_periodic_slab_hires.py`` for the +full story: a 32^3 grid (32768 dofs) with PETSc's preconditioner explicitly set to ``"gamg"`` +(algebraic multigrid, via ``SolverParameters.pc_type`` -- see ``struphy.linear_algebra.solver.SolverParameters``) instead of the default ``"jacobi"``. Measured via ``struphy.linear_algebra.petsc_examples_benchmark``'s repeated-solve methodology (isolates just the solve, amortizing gamg's one-time multigrid setup across several calls at fixed dt -- see that module's docstring): pcg needs ~190 CG iterations per solve (~10.9s) at this -size, against PETSc+gamg's ~2 iterations (~0.29s) -- a ~38x difference, essentially the same -*ratio* periodic_slab already shows (~36x at 13824 dofs): this near-singular system is -ill-conditioned mainly through its weakly constrained DC mode, not primarily through resolution, -so the ratio does not grow much further with grid size. What does grow is the *absolute* cost: -pcg's ~11 seconds per solve here is the practically relevant "huge difference" once multiplied -over the many timesteps of a real simulation. See ``submit_toydrift_petsc.py`` and -``submit_strong_landau_damping_petsc.py`` for the smaller-scale comparisons. +size, against PETSc+gamg's ~2 iterations (~0.29s) -- a ~38x difference. This near-singular system +is ill-conditioned mainly through its weakly constrained DC mode, not primarily through +resolution, so the ratio does not grow much further with grid size -- but the *absolute* cost +does: pcg's ~11 seconds per solve here is the practically relevant "huge difference" once +multiplied over the many timesteps of a real simulation. See +``submit_strong_landau_damping_petsc.py`` and ``submit_vlasov_maxwell_petsc.py`` for the +smaller-scale (3-10x) comparisons. Each launch runs one 32^3-grid, one-time-step simulation (``params_periodic_slab_hires.py``'s ``__main__`` calls ``sim.run(one_time_step=True)``); note that a *single* one-time-step run still diff --git a/profiling/submit_toydrift_petsc.py b/profiling/submit_toydrift_petsc.py deleted file mode 100644 index ece5d8c70..000000000 --- a/profiling/submit_toydrift_petsc.py +++ /dev/null @@ -1,78 +0,0 @@ -"""ToyDrift periodic-slab: PETSc vs. pcg for a real per-step Poisson solve. - -This file defines a single profiling case built from the periodic-slab ToyDrift setup (see -`profiling/examples/ToyDrift/periodic_slab/params_periodic_slab.py` -- there is no periodic -ToyDrift example under `examples/`: the real one, `examples/ToyGyrokinetic/diocotron_instability`, -uses a physically non-periodic HollowCylinder domain, since radial confinement is inherent to the -diocotron instability; this case swaps in a periodic Cuboid domain instead, which works because -ToyDrift's field solve is a plain `PoissonSolve` with no geometry-coupled averaging -- unlike -`PoissonAdiabaticGyrokinetic`, used by `DriftKineticElectrostaticAdiabatic`, which was tried first -and diverges outright on a periodic domain regardless of options), launched once per (rank count, -solver) combination via `--solver pcg`/`--solver petsc` (`param_flags`), which sets -`model.propagators.gc_poisson.options.solver`. - -Unlike VlasovAmpereOneSpecies (whose Poisson solve only runs once, as an initial condition), -ToyDrift's `gc_poisson` runs as a *regular per-step propagator* -- exactly the repeated-solve -pattern where PETSc's algebraic multigrid preconditioner shows a genuine win, without needing the -initial-Poisson-only benchmark's workaround of re-invoking the propagator by hand after -`sim.run()`. Each generated script runs the simulation itself by invoking -`params_periodic_slab.py` directly (its `__main__` block calls `sim.run(one_time_step=True)`, a -single-step timing snapshot). - -See `submit_strong_landau_damping_petsc.py` for the same comparison on VlasovAmpereOneSpecies, and -its module docstring / `struphy.linear_algebra.petsc_examples_benchmark`'s for the general PETSc -vs. pcg story (including the known MPI correctness caveat for this near-singular stab_eps regime). -""" - -import argparse -from pathlib import Path - -from profiling_job import ProfilingCase - - -def main() -> None: - - # Parse arguments, do not remove --upload - parser = argparse.ArgumentParser( - description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), - ) - parser.add_argument( - "--upload", - action="store_true", - help="Upload the packaged profiling results to the profiling-data repo.", - ) - args = parser.parse_args() - - # Paths relative to this script's location, so it can be run from anywhere. - script_dir = Path(__file__).resolve().parent - params_dir = script_dir / "examples" / "ToyDrift" / "periodic_slab" - - profiling_case = ProfilingCase( - label="toydrift_periodic_slab_petsc", - name="ToyDrift periodic slab, per-step Poisson solve: pcg vs. PETSc", - description=( - "Periodic-slab ToyDrift setup, solving the per-step guiding-center Poisson problem " - "with either feectools' native preconditioned CG (--solver pcg) or PETScSolver " - "(KSP=cg, PC=gamg, --solver petsc)." - ), - physics_problem="Electrostatic E x B drift of a single ion species in a periodic slab.", - struphy_model_used="ToyDrift", - params_source=params_dir / "params_periodic_slab.py", - language="fortran", - compiler="GNU", - upload=args.upload, - ) - profiling_case.use_slurm = False - - # Launch one run per (rank count, solver) combination, same rank counts for both solvers so - # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` - # output folder), so the two solvers never collide even at the same rank count. - for num_tasks in (1, ): #2, 4): - for solver in ("pcg", "petsc"): - profiling_case.launch(num_tasks, param_flags=["--solver", solver]) - - profiling_case.finalize_run() - - -if __name__ == "__main__": - main() diff --git a/src/struphy/linear_algebra/petsc_examples_benchmark.py b/src/struphy/linear_algebra/petsc_examples_benchmark.py index 8004f69f2..2192276cb 100644 --- a/src/struphy/linear_algebra/petsc_examples_benchmark.py +++ b/src/struphy/linear_algebra/petsc_examples_benchmark.py @@ -15,20 +15,18 @@ re-invokes ``model.initial_poisson`` directly after ``sim.run()`` rather than relying on the model's own (single-shot) usage of it. -- ``ToyDrift/periodic_slab``: no periodic ToyDrift example exists under ``examples/`` (the real - one, ``examples/ToyGyrokinetic/diocotron_instability``, needs a physically non-periodic +- ``ToyDrift/periodic_slab_hires``: no periodic ToyDrift example exists under ``examples/`` (the + real one, ``examples/ToyGyrokinetic/diocotron_instability``, needs a physically non-periodic HollowCylinder domain), so this one is written from scratch with a periodic Cuboid domain instead -- which works because ToyDrift's field solve is a plain ``PoissonSolve`` with no geometry-coupled averaging (unlike ``PoissonAdiabaticGyrokinetic``, used by ``DriftKineticElectrostaticAdiabatic``, which diverges outright on a periodic domain regardless of options -- tried first, not usable here). Unlike VlasovAmpereOneSpecies, ToyDrift's ``gc_poisson`` runs as a *regular per-step propagator*, so no re-invocation workaround is needed. - -- ``ToyDrift/periodic_slab_hires``: same setup as ``periodic_slab``, scaled up to a 32^3 grid - (vs. 24^3) to push feectools' unpreconditioned CG into several hundred iterations per solve - while PETSc+gamg (``pc_type="gamg"``, set explicitly via ``SolverParameters.pc_type`` -- see - ``struphy.linear_algebra.solver.SolverParameters``) stays at a handful, regardless of grid size. - The most lopsided case in this suite by design; see its own params file's docstring. + Uses a 32^3 grid to push feectools' unpreconditioned CG into several hundred iterations per + solve while PETSc+gamg (``pc_type="gamg"``, set explicitly via ``SolverParameters.pc_type`` -- + see ``struphy.linear_algebra.solver.SolverParameters``) stays at a handful, regardless of grid + size. The most lopsided case in this suite by design; see its own params file's docstring. - ``VlasovMaxwellOneSpecies/weibel_instability``: plain copy of ``examples/VlasovMaxwellOneSpecies/weibel_instability/params_weibel_instability.py``, scaled up @@ -89,7 +87,6 @@ ("VlasovAmpereOneSpecies", "weak_Landau_damping"), ("VlasovAmpereOneSpecies", "two_stream"), ("VlasovAmpereOneSpecies", "bump_on"), - ("ToyDrift", "periodic_slab"), ("ToyDrift", "periodic_slab_hires"), ("VlasovMaxwellOneSpecies", "weibel_instability"), ) From 885589ca531e3d1ea6e8a714d496f18843bd0c17 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 10 Aug 2026 10:49:18 +0200 Subject: [PATCH 24/32] Use xp.asarray --- src/struphy/models/variables.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 2fa50e6a3..d943369ca 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -581,7 +581,7 @@ def allocate( self.particles.draw_markers(sort=sort) # set zero velocity according to loading_params - zero_index = xp.nonzero(self.particles.loading_params.set_zero_velocity)[0].flatten() + zero_index = xp.nonzero(xp.asarray(self.particles.loading_params.set_zero_velocity))[0].flatten() self.particles.set_velocities_comp(velocity=0.0, comp=zero_index) self.particles.initialize_weights() From 61f0482dda99768bb2d001230a90e2d3e26cd2a8 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 10 Aug 2026 15:25:16 +0200 Subject: [PATCH 25/32] Removed hardcode use_slurm=False --- profiling/submit_diocotron_strong_scaling.py | 1 - profiling/submit_strong_landau_damping_petsc.py | 1 - profiling/submit_toydrift_hires_petsc.py | 1 - profiling/submit_vlasov_maxwell_petsc.py | 1 - 4 files changed, 4 deletions(-) diff --git a/profiling/submit_diocotron_strong_scaling.py b/profiling/submit_diocotron_strong_scaling.py index 8fbcc7d6f..0a421ad14 100644 --- a/profiling/submit_diocotron_strong_scaling.py +++ b/profiling/submit_diocotron_strong_scaling.py @@ -37,7 +37,6 @@ def main() -> None: upload=args.upload, ) - profiling_case.use_slurm = False # Launch one run per rank count for num_tasks in (2, 4): # , 8, 16, 32, 64, 128, 256): diff --git a/profiling/submit_strong_landau_damping_petsc.py b/profiling/submit_strong_landau_damping_petsc.py index 71a616e16..d0ac95cae 100644 --- a/profiling/submit_strong_landau_damping_petsc.py +++ b/profiling/submit_strong_landau_damping_petsc.py @@ -27,7 +27,6 @@ def main() -> None: compiler="GNU", upload=args.upload, ) - profiling_case.use_slurm = False # Launch one run per (rank count, solver) combination, same rank counts for both solvers so # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` diff --git a/profiling/submit_toydrift_hires_petsc.py b/profiling/submit_toydrift_hires_petsc.py index 12623036e..a84c6f344 100644 --- a/profiling/submit_toydrift_hires_petsc.py +++ b/profiling/submit_toydrift_hires_petsc.py @@ -67,7 +67,6 @@ def main() -> None: compiler="GNU", upload=args.upload, ) - profiling_case.use_slurm = False # Launch one run per (rank count, solver) combination, same rank counts for both solvers so # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` diff --git a/profiling/submit_vlasov_maxwell_petsc.py b/profiling/submit_vlasov_maxwell_petsc.py index fc02031e4..7108f7701 100644 --- a/profiling/submit_vlasov_maxwell_petsc.py +++ b/profiling/submit_vlasov_maxwell_petsc.py @@ -28,7 +28,6 @@ def main() -> None: compiler="GNU", upload=args.upload, ) - profiling_case.use_slurm = False # Launch one run per (rank count, solver) combination, same rank counts for both solvers so # they are directly comparable. Each launch gets its own `--id` (hence its own `sim_` From e2c44781018c18e80264b1a59a4292a717adc4c1 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 10 Aug 2026 11:44:49 +0200 Subject: [PATCH 26/32] Added petsc module to pitagora --- setup/modules.pitagora.sh | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/setup/modules.pitagora.sh b/setup/modules.pitagora.sh index 795dca625..46b0f59b2 100644 --- a/setup/modules.pitagora.sh +++ b/setup/modules.pitagora.sh @@ -3,5 +3,10 @@ intel-oneapi-mkl/2024.0.0--intel-oneapi-mpi--2021.12.1 \ python/3.11.7" MODULES_GCC="gcc/12.3.0 \ -openmpi/4.1.6--gcc--12.3.0 \ +openmpi/4.1.6--gcc--12.3.0-ucx1.20 \ +petsc/3.22.1--openmpi--4.1.6--gcc--12.3.0-ucx1.20-complex-mumps \ python/3.11.7" + +# The petsc module above is built with CUDA support, so petsc4py's import dlopens +# libcuda.so.1 even on nodes without a GPU driver. Point at a stub so it doesn't fail. +export LD_LIBRARY_PATH="$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)/.venv/petsc_cuda_stub:${LD_LIBRARY_PATH:-}" From 31562c68a0a143455dfe8b938aa99d90dec24fc3 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 10 Aug 2026 15:24:37 +0200 Subject: [PATCH 27/32] rename --- profiling/submit_toydrift_hires_petsc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/profiling/submit_toydrift_hires_petsc.py b/profiling/submit_toydrift_hires_petsc.py index a84c6f344..bff109e08 100644 --- a/profiling/submit_toydrift_hires_petsc.py +++ b/profiling/submit_toydrift_hires_petsc.py @@ -49,7 +49,7 @@ def main() -> None: profiling_case = ProfilingCase( label="toydrift_periodic_slab_hires_petsc", - name="ToyDrift periodic slab (32^3, hires): pcg vs. PETSc+gamg at a larger problem size", + name="ToyDrift periodic slab: PCG vs. PETSc+GAMG", description=( "Higher-resolution (32^3 grid, 32768 dofs, degree-3 splines) periodic-slab ToyDrift " "setup, solving the per-step guiding-center Poisson problem with either feectools' " From 763c706aaa0e0392557319b3f4d22c2524902889 Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 10 Aug 2026 22:47:10 +0200 Subject: [PATCH 28/32] Added petsc solver in the schur solver --- src/struphy/linear_algebra/schur_solver.py | 62 +++++++++++++++++----- 1 file changed, 49 insertions(+), 13 deletions(-) diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index d97e7b7a2..832ec6cff 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -4,6 +4,7 @@ from line_profiler import profile from struphy.linear_algebra.solver import SolverParameters +from struphy.linear_algebra.solver import inverse as struphy_inverse class SchurSolver: @@ -69,9 +70,30 @@ def __init__( # linear operators self._A = A self._BC = BC + # Set by the A/BC property setters whenever a caller reassigns either (e.g. + # VlasovAmpereCoupling/EfieldWeightsCoupling rebuild `.BC` from a fresh, particle-dependent + # operator every call); only consulted for petsc, see below. + self._schur_dirty = True + + self._is_petsc = solver_name == "petsc" + + if self._is_petsc: + # PETScSolver caches its assembled PETSc.Mat by the *object identity* of `linop` + # (see PETScSolver._get_ksp), rebuilding only when that identity changes -- exactly + # the mechanism ImplicitDiffusion relies on (see its lhs-operator caching). The + # in-place-mutated `self._schur` buffer below defeats that: it is the same Python + # object on every call, so PETScSolver would keep the *first* call's matrix forever, + # silently going stale if `dt`, `A` or `BC` ever change. So for petsc, `self._schur` + # is instead a fresh composite operator, rebuilt only when `dt` changes or `.A`/`.BC` + # were reassigned (`self._schur_dirty`) -- see __call__. Callers that mutate an + # operator obtained via the `A`/`BC` getters in place, without reassigning it through + # the setter, would not be detected -- no current caller does this. + self._schur = None + self._schur_dt = None + else: + # Allocate memory for matrices used in solving the Schur system + self._schur = A.copy() - # Allocate memory for matrices used in solving the Schur system - self._schur = A.copy() self._rhs_m = A.copy() # initialize solver with dummy matrix A @@ -79,15 +101,16 @@ def __init__( kwargs = solver_params.__dict__.copy() kwargs.pop("info") - # pc_type is petsc-only (see struphy.linear_algebra.solver.SolverParameters); this - # module always goes through feectools' own `inverse` (imported directly above, not - # struphy's petsc-aware wrapper), whose InverseLinearOperator subclasses forward - # unknown kwargs straight to their constructor and would raise on it. - kwargs.pop("pc_type", None) if precond is not None: kwargs["pc"] = precond - self._solver = inverse(A, solver_name, **kwargs) + if self._is_petsc: + # struphy's inverse() dispatches "petsc" to PETScSolver and forwards pc_type; the + # dummy operator here is just to build the solver object -- __call__ always assigns + # the real one via `self._solver.linop` before solving. + self._solver = struphy_inverse(A, solver_name, **kwargs) + else: + self._solver = inverse(A, solver_name, **kwargs) # right-hand side vector (avoids temporary memory allocation!) self._rhs = A.codomain.zeros() @@ -106,11 +129,17 @@ def BC(self): def A(self, a): """Upper left block from [[A B], [C Id]].""" self._A = a + # e.g. VlasovAmpereCoupling/EfieldWeightsCoupling reassign `.A`/`.BC` to a fresh + # (possibly particle-dependent) operator every call; `x.A *= y`-style augmented + # assignment also lands here (Python always re-invokes the setter). See the petsc + # cache-invalidation note in __init__/__call__. + self._schur_dirty = True @BC.setter def BC(self, bc): """Product from [[A B], [C Id]].""" self._BC = bc + self._schur_dirty = True @profile def __call__(self, xn, Byn, dt, out=None): @@ -144,12 +173,19 @@ def __call__(self, xn, Byn, dt, out=None): assert xn.space == self._A.domain assert Byn.space == self._A.codomain - # left- and right-hand side operators - self._schur *= 0.0 - self._schur += self._BC - self._schur *= -(dt**2) - self._schur += self._A + # left-hand side operator + if self._is_petsc: + if self._schur is None or dt != self._schur_dt or self._schur_dirty: + self._schur = self._A - (dt**2) * self._BC + self._schur_dt = dt + self._schur_dirty = False + else: + self._schur *= 0.0 + self._schur += self._BC + self._schur *= -(dt**2) + self._schur += self._A + # right-hand side operator self._rhs_m *= 0.0 self._rhs_m += self._BC self._rhs_m *= dt**2 From 8fe9a3bf0c1c434a92b8583f05a12bc2f950710d Mon Sep 17 00:00:00 2001 From: Max Date: Mon, 10 Aug 2026 22:47:24 +0200 Subject: [PATCH 29/32] Added src/struphy/linear_algebra/petsc_speedup_example.py --- .../linear_algebra/petsc_speedup_example.py | 142 ++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 src/struphy/linear_algebra/petsc_speedup_example.py diff --git a/src/struphy/linear_algebra/petsc_speedup_example.py b/src/struphy/linear_algebra/petsc_speedup_example.py new file mode 100644 index 000000000..beb81e96d --- /dev/null +++ b/src/struphy/linear_algebra/petsc_speedup_example.py @@ -0,0 +1,142 @@ +"""Standalone, self-contained example: PETSc beats pcg on a real (non-toy) Poisson solve. + +Run directly -- no profiling infrastructure, no submission, nothing to set up beyond an +environment with petsc4py installed (``pip install -e ".[petsc]"``): + +.. code-block:: bash + + python src/struphy/linear_algebra/petsc_speedup_example.py + +What it does +------------ +Builds the same ToyDrift periodic-slab setup used by +``profiling/examples/ToyDrift/periodic_slab_hires`` (the largest PETSc-vs-pcg gap in struphy's +profiling suite) directly in this script, so you can read top to bottom exactly what is being +solved and how it is timed -- no need to trace through ``ProfilingCase``/submit-script machinery. + +ToyDrift's ``gc_poisson`` is a *regular per-step propagator* (unlike e.g. VlasovAmpereOneSpecies, +which only solves Poisson once as an initial condition), so it is called once per timestep with +the left-hand-side operator reused across calls at fixed ``dt`` (see +``ImplicitDiffusion.__call__``'s lhs-operator caching). That means the *first* solve pays for +PETSc's one-time matrix assembly and, with ``pc_type="gamg"``, multigrid hierarchy construction -- +this script does one untimed warm-up call for exactly that reason, matching how the cost would +amortize over the many timesteps of a real simulation, before timing several further calls. + +What to expect +--------------- +At this problem size (32768 dofs), feectools' unpreconditioned CG needs on the order of a few +hundred iterations per solve (several seconds), while PETSc with an algebraic multigrid +preconditioner (``pc_type="gamg"``) needs only a handful (well under a second) -- typically a +30-40x speedup. The two solutions are also checked against each other (mean-removed, since the +near-singular constant/DC mode is only weakly constrained and not physically meaningful here -- +see ``PoissonSolve``'s stabilization) and should agree to within floating-point noise. +""" + +import time + +import cunumpy as xp + +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import ToyDrift + + +def _build_and_run(solver: str, pc_type: str, n_solves: int): + """Build a fresh ToyDrift periodic-slab simulation and time repeated Poisson solves.""" + model = ToyDrift(base_units=BaseUnits(kBT=1.0)) + + env = EnvironmentOptions(sim_folder=f"sim_petsc_speedup_example_{solver}") + time_opts = Time(dt=0.05, Tend=0.05, split_algo="LieTrotter") + domain = domains.Cuboid() # periodic: required for PETScSolver's DirectionalDerivativeOperator + equil = equils.HomogenSlab(B0z=1.0, n0=1.0) + grid = grids.TensorProductGrid(num_elements=(32, 32, 32)) + derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(None, None, None)) # fully periodic + + sim = Simulation( + model=model, + params_path=None, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, + ) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=5, seed=42), + weights_params=WeightsParameters(control_variate=True), + boundary_params=BoundaryParameters(), + sorting_params=SortingParameters(boxes_per_dim=(4, 4, 4), do_sort=True), + bufsize=0.4, + ) + + model.propagators.gc_poisson.options.solver = solver + model.propagators.gc_poisson.options.solver_params = SolverParameters( + tol=1e-10, + maxiter=5_000, + pc_type=pc_type, # ignored by pcg, see SolverParameters.pc_type + ) + model.propagators.push_gc_bxe.options = model.propagators.push_gc_bxe.Options( + algo="explicit", + evaluate_e_field=True, + ) + + background = maxwellians.GyroMaxwellian2D(n=(1.0, None), vth_para=(1.0, None), vth_perp=(1.0, None), equil=equil) + model.kinetic_ions.var.add_background(background) + perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) + init = maxwellians.GyroMaxwellian2D(n=(1.0, perturbation), vth_para=(1.0, None), vth_perp=(1.0, None), equil=equil) + model.kinetic_ions.var.add_initial_condition(init) + + # real setup: particle loading, Derham/mass operators, and one (untimed) Poisson solve + sim.run(one_time_step=True) + + poisson = model.propagators.gc_poisson + dt = time_opts.dt + + poisson(dt) # warm-up: pays for one-time matrix assembly / gamg hierarchy construction + t0 = time.perf_counter() + for _ in range(n_solves): + poisson(dt) + elapsed = (time.perf_counter() - t0) / n_solves + + info = poisson._solver.get_info() if hasattr(poisson._solver, "get_info") else poisson._solver._info + phi = model.em_fields.phi.spline.vector.toarray() + return elapsed, info, phi + + +def main(n_solves: int = 5): + print(f"Timing {n_solves} repeated Poisson solves per solver (after one warm-up call) ...\n") + + t_pcg, info_pcg, phi_pcg = _build_and_run("pcg", pc_type="jacobi", n_solves=n_solves) + t_petsc, info_petsc, phi_petsc = _build_and_run("petsc", pc_type="gamg", n_solves=n_solves) + + mean_removed_pcg = phi_pcg - phi_pcg.mean() + mean_removed_petsc = phi_petsc - phi_petsc.mean() + rel_err = xp.linalg.norm(mean_removed_pcg - mean_removed_petsc) / xp.linalg.norm(phi_pcg) + + print(f"pcg (unpreconditioned) : {t_pcg * 1e3:9.2f} ms/solve niter={info_pcg.get('niter')}") + print(f"petsc + gamg : {t_petsc * 1e3:9.2f} ms/solve niter={info_petsc.get('niter')}") + print(f"\nspeedup: {t_pcg / t_petsc:.2f}x") + print(f"relative solution mismatch (mean-removed): {rel_err:.2e}") + assert rel_err < 1e-6, f"pcg/petsc solutions disagree by {rel_err:.2e}, expected < 1e-6" + print("\nSolutions agree -- the speedup above is not at the cost of correctness.") + + +if __name__ == "__main__": + main() From 74bd130d14b3fb710a143b0132b81da4c4df82e3 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 11 Aug 2026 09:58:37 +0200 Subject: [PATCH 30/32] SchurSolver now supports solver=petsc --- src/struphy/linear_algebra/petsc_solver.py | 48 +++-- src/struphy/linear_algebra/schur_solver.py | 6 + .../tests/test_petsc_poisson_solve_pic.py | 4 +- .../tests/test_petsc_schur_solver.py | 185 ++++++++++++++++++ 4 files changed, 229 insertions(+), 14 deletions(-) create mode 100644 src/struphy/linear_algebra/tests/test_petsc_schur_solver.py diff --git a/src/struphy/linear_algebra/petsc_solver.py b/src/struphy/linear_algebra/petsc_solver.py index 559c2b783..3b4f5a414 100644 --- a/src/struphy/linear_algebra/petsc_solver.py +++ b/src/struphy/linear_algebra/petsc_solver.py @@ -20,7 +20,7 @@ def _directional_derivative_to_stencil_matrix(op): - """ Build a :class:`~feectools.linalg.stencil.StencilMatrix` equivalent to a + """Build a :class:`~feectools.linalg.stencil.StencilMatrix` equivalent to a (matrix-free) :class:`~feectools.feec.derivatives.DirectionalDerivativeOperator`, so it can be handed to :func:`feectools.linalg.topetsc.mat_topetsc`. @@ -84,36 +84,60 @@ def off(o): return M +def _materialize_block(block): + """Turn a ``BlockLinearOperator`` block entry into a concrete matrix (``StencilMatrix`` or + ``None``) that :func:`feectools.linalg.topetsc.mat_topetsc` can handle directly. + + Blocks of a topological operator such as ``derham.curl`` are not always plain + ``StencilMatrix``: sign conventions are sometimes expressed via ``ScaledLinearOperator`` + wrapping a ``StencilMatrix``/``DirectionalDerivativeOperator`` rather than baking the sign + into the matrix data (observed e.g. for ``derham.curl.T``, whose transposed blocks land on + this path). ``mat_topetsc`` calls ``.update_ghost_regions()`` on every block, which only + concrete matrix types implement -- so any such wrapper must be resolved to a concrete matrix + first. Recurses through nested ``ScaledLinearOperator``s. + """ + if block is None: + return None + if isinstance(block, DirectionalDerivativeOperator): + return _directional_derivative_to_stencil_matrix(block) + if isinstance(block, ScaledLinearOperator): + inner = _materialize_block(block.operator) + if inner is None: + return None + scaled = inner.copy() + scaled *= block.scalar + return scaled + return block + + def _assemble_leaf_operator(A): - """ Return an operator equivalent to `A` that is directly convertible via + """Return an operator equivalent to `A` that is directly convertible via :func:`feectools.linalg.topetsc.mat_topetsc` (i.e. a ``StencilMatrix`` or a ``BlockLinearOperator`` whose blocks are all ``StencilMatrix``), replacing any - ``DirectionalDerivativeOperator`` (block or bare) by its assembled equivalent. + ``DirectionalDerivativeOperator``/``ScaledLinearOperator`` (block or bare) by its assembled + equivalent -- see :func:`_materialize_block`. """ - if isinstance(A, DirectionalDerivativeOperator): - return _directional_derivative_to_stencil_matrix(A) + if isinstance(A, (DirectionalDerivativeOperator, ScaledLinearOperator)): + return _materialize_block(A) if isinstance(A, BlockLinearOperator): out = BlockLinearOperator(A.domain, A.codomain) for i, j in A.nonzero_block_indices: - block = A[i, j] - out[i, j] = _directional_derivative_to_stencil_matrix(block) if isinstance( - block, DirectionalDerivativeOperator - ) else block + out[i, j] = _materialize_block(A[i, j]) return out return A def _comm_of(space): - """ MPI communicator of a StencilVectorSpace/BlockVectorSpace, matching mat_topetsc's convention. """ + """MPI communicator of a StencilVectorSpace/BlockVectorSpace, matching mat_topetsc's convention.""" if isinstance(space, BlockVectorSpace): return space.spaces[0].cart.global_comm return space.cart.global_comm def _identity_petsc_mat(space): - """ Build a PETSc.Mat representing the identity operator on `space`. """ + """Build a PETSc.Mat representing the identity operator on `space`.""" from petsc4py import PETSc comm = _comm_of(space) @@ -134,7 +158,7 @@ def _identity_petsc_mat(space): def _assemble_petsc_matrix(A): - """ Recursively assemble a ``PETSc.Mat`` for a (possibly composite) feectools + """Recursively assemble a ``PETSc.Mat`` for a (possibly composite) feectools ``LinearOperator``, by converting every assembled leaf via :func:`feectools.linalg.topetsc.mat_topetsc` and combining the pieces with PETSc's own matrix algebra (``matMult`` for composition, ``axpy`` for sums, ``scale`` for scalar diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index 832ec6cff..19595d5b7 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -110,6 +110,12 @@ def __init__( # the real one via `self._solver.linop` before solving. self._solver = struphy_inverse(A, solver_name, **kwargs) else: + # pc_type is petsc-only (see struphy.linear_algebra.solver.SolverParameters); this + # branch goes straight to feectools' own `inverse` (imported directly above, not + # struphy's petsc-aware wrapper, which would otherwise strip it), whose + # InverseLinearOperator subclasses forward unknown kwargs straight to their + # constructor and would raise on it. + kwargs.pop("pc_type", None) self._solver = inverse(A, solver_name, **kwargs) # right-hand side vector (avoids temporary memory allocation!) diff --git a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py index f2cac1ed6..9d2f06fcc 100644 --- a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py +++ b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py @@ -1,5 +1,6 @@ import cunumpy as xp import pytest +from cunumpy import PyccelKernel pytest.importorskip("petsc4py") @@ -18,7 +19,6 @@ from struphy.propagators.base import Propagator from struphy.propagators.poisson_solve import PoissonSolve from struphy.topology.grids import TensorProductGrid -from struphy.utils.pyccel import Pyccelkernel class _FakePICVariable: @@ -75,7 +75,7 @@ def run(solver_name): rho = ParticlesToGrid( _FakePICVariable(particles), "H1", - Pyccelkernel(accum_kernels.charge_density_0form), + PyccelKernel(accum_kernels.charge_density_0form), ) phi = FEECVariable(space="H1") diff --git a/src/struphy/linear_algebra/tests/test_petsc_schur_solver.py b/src/struphy/linear_algebra/tests/test_petsc_schur_solver.py new file mode 100644 index 000000000..63e71917b --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_schur_solver.py @@ -0,0 +1,185 @@ +"""Regression tests for SchurSolver's solver="petsc" support (struphy.linear_algebra.schur_solver). + +SchurSolver previously always imported feectools' own `inverse`, which does not recognize the +name "petsc" -- so `solver="petsc"` would raise for every propagator built on it (MaxwellWeakAmpere, +VlasovAmpereCoupling, EfieldWeightsCoupling, CurlCurlSolve, ...), even though those propagators' +Options all declare `LiteralOptions.OptsSymmSolver` (which includes "petsc") as their solver type. + +Wiring petsc in was not just an import swap: PETScSolver caches its assembled PETSc.Mat by the +*object identity* of the operator assigned to `.linop` (see PETScSolver._get_ksp), rebuilding only +when that identity changes. SchurSolver's non-petsc path mutates its `self._schur` buffer *in +place* every call (`self._schur *= 0.0; += ...`) -- the same Python object every time, which would +make PETScSolver silently reuse a stale matrix forever after the first call. Two call patterns +exist among current callers, and both need to be correct: + +- MaxwellWeakAmpere never reassigns `.A`/`.BC` after construction (both are geometric, constant + operators) -- caching by `dt` alone is correct and safe there. +- VlasovAmpereCoupling (and EfieldWeightsCoupling) reassign `.BC` to a fresh, particle-dependent + operator via the property setter on *every* call -- caching there must be invalidated every + time, tracked via a dirty flag set in the `A`/`BC` property setters (Python routes both + `x.BC = y` and the augmented `x.BC *= y` through the setter). + +These tests exercise both patterns directly (not synthetically) via the real propagators. +""" + +import shutil +import tempfile + +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.mpi import mpi as MPI + +from struphy import ( + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + Simulation, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.domains import Cuboid +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import VlasovAmpereOneSpecies +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.propagators.maxwell_weak_ampere import MaxwellWeakAmpere +from struphy.topology.grids import TensorProductGrid + + +def test_maxwell_weak_ampere_petsc_matches_pcg(): + """MaxwellWeakAmpere(solver="petsc") must match solver="pcg" over several implicit timesteps. + + Exercises SchurSolver's petsc dt-only cache-invalidation path (see module docstring): + MaxwellWeakAmpere never reassigns `.A`/`.BC` after allocate(), so the same lhs operator must + be correctly reused across all calls at fixed dt. + """ + comm = MPI.COMM_WORLD + + domain = Cuboid() + grid = TensorProductGrid(num_elements=[6, 6, 6]) + derham_opts = DerhamOptions(degree=[2, 2, 2], bcs=(None, None, None)) + derham = Derham(grid, derham_opts, comm=comm) + mass_ops = WeightedMassOperators(derham, domain) + + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + def run(solver_name): + e_field = FEECVariable(space="Hcurl") + e_field.add_perturbation(perturbations.ModesCos(amps=(0.1,), ls=(1,), comp=0)) + e_field.allocate(derham=derham, domain=domain) + + b_field = FEECVariable(space="Hdiv") + b_field.add_perturbation(perturbations.ModesCos(amps=(0.05,), ls=(1,), comp=1)) + b_field.allocate(derham=derham, domain=domain) + + prop = MaxwellWeakAmpere() + prop.variables.e = e_field + prop.variables.b = b_field + prop.options = prop.Options( + algo="implicit", + solver=solver_name, + solver_params=SolverParameters(tol=1e-11, maxiter=3000), + ) + prop.allocate() + + dt = 0.02 + for _ in range(3): + prop(dt) + + return e_field.spline.vector.toarray(), b_field.spline.vector.toarray() + + e_pcg, b_pcg = run("pcg") + e_petsc, b_petsc = run("petsc") + + rel_err_e = xp.linalg.norm(e_pcg - e_petsc) / xp.linalg.norm(e_pcg) + rel_err_b = xp.linalg.norm(b_pcg - b_petsc) / xp.linalg.norm(b_pcg) + assert rel_err_e < 1e-6, f"e-field mismatch: {rel_err_e:.2e}" + assert rel_err_b < 1e-6, f"b-field mismatch: {rel_err_b:.2e}" + + +def test_vlasov_ampere_coupling_petsc_matches_pcg_with_real_pic_deposition(): + """VlasovAmpereCoupling(solver="petsc") must match solver="pcg" over several real timesteps, + driven by real particle-in-cell deposition (not a synthetic source). + + Exercises SchurSolver's petsc dirty-flag cache-invalidation path (see module docstring): + VlasovAmpereCoupling reassigns `.BC` to a fresh, particle-dependent operator every call, which + the dt-only cache used by MaxwellWeakAmpere's test above would get wrong if applied here. + Goes through the full model/Simulation machinery (unlike the other petsc regression tests in + this directory, which build propagators directly) because VlasovAmpereCoupling requires a real + PICVariable/ParticleSpecies (species.equation_params, weights_params) that is impractical to + duck-type -- see test_petsc_poisson_solve_pic.py, whose ParticlesToGrid-based fake works + because ParticlesToGrid does not check isinstance, unlike VlasovAmpereCoupling.Variables.ions. + """ + comm = MPI.COMM_WORLD + + def run(solver_name, out_folder): + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + env = EnvironmentOptions(out_folders=out_folder, sim_folder=f"sim_{solver_name}") + time_opts = Time(dt=0.02, Tend=0.06, split_algo="LieTrotter") + domain = domains.Cuboid(r1=12.56) + grid = grids.TensorProductGrid(num_elements=(8, 8, 8)) + derham_opts = DerhamOptions(degree=(2, 2, 2), bcs=(None, None, None)) + + sim = Simulation( + model=model, + params_path=None, + env=env, + time_opts=time_opts, + domain=domain, + equil=None, + grid=grid, + derham_opts=derham_opts, + ) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(Np=5_000, seed=1234), + weights_params=WeightsParameters(control_variate=True), + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options( + solver=solver_name, + solver_params=SolverParameters(tol=1e-10, maxiter=5000), + ) + model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + + background = maxwellians.Maxwellian3D(n=(1.0, None)) + perturbation = perturbations.ModesCos(amps=(0.5,), ls=(1,)) + init = maxwellians.Maxwellian3D(n=(1.0, perturbation)) + model.kinetic_ions.var.add_background(background) + model.kinetic_ions.var.add_initial_condition(init) + + sim.run() # several real steps, not one_time_step: coupling_va.BC changes every call + + return model.em_fields.e_field.spline.vector.toarray() + + out_folder = tempfile.mkdtemp() if comm.Get_rank() == 0 else None + out_folder = comm.bcast(out_folder, root=0) + + try: + e_pcg = run("pcg", out_folder) + e_petsc = run("petsc", out_folder) + finally: + comm.Barrier() + if comm.Get_rank() == 0: + shutil.rmtree(out_folder, ignore_errors=True) + + rel_err = xp.linalg.norm(e_pcg - e_petsc) / xp.linalg.norm(e_pcg) + assert rel_err < 1e-6, f"e-field mismatch: {rel_err:.2e}" + + +if __name__ == "__main__": + test_maxwell_weak_ampere_petsc_matches_pcg() + test_vlasov_ampere_coupling_petsc_matches_pcg_with_real_pic_deposition() From 0cbc39e433f2837cc37c30b41f38e45b020e7542 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 11 Aug 2026 10:54:24 +0200 Subject: [PATCH 31/32] Added src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py --- src/struphy/linear_algebra/schur_solver.py | 10 +- .../tests/test_petsc_schur_solver_full.py | 152 ++++++++++++++++++ 2 files changed, 160 insertions(+), 2 deletions(-) create mode 100644 src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index 19595d5b7..6bf7d2a90 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -269,7 +269,12 @@ def __init__(self, M, solver_name, **solver_params): self._S = self._A - self._B @ self._C - self._solver = inverse(self._S, solver_name, **solver_params) + # struphy_inverse dispatches solver_name="petsc" to PETScSolver and safely strips + # petsc-only kwargs (e.g. pc_type) for every other solver -- see SchurSolver, which needs + # this same dispatch but (unlike this class) also has to handle a stale-cache hazard from + # in-place operator mutation; no such hazard here since callers rebuild this whole object + # fresh each call rather than mutating `self._S` in place. + self._solver = struphy_inverse(self._S, solver_name, **solver_params) # right-hand side vector (avoids temporary memory allocation!) self._rhs = self._A.codomain.zeros() @@ -386,7 +391,8 @@ def __init__(self, M, solver_name, **solver_params): self._S = self._A - self._B @ self._C - self._D @ self._E - self._solver = inverse(self._S, solver_name, **solver_params) + # see SchurSolverFull.__init__'s note on struphy_inverse + self._solver = struphy_inverse(self._S, solver_name, **solver_params) # right-hand side vector (avoids temporary memory allocation!) self._rhs = self._A.codomain.zeros() diff --git a/src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py b/src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py new file mode 100644 index 000000000..1ec7e17ce --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_petsc_schur_solver_full.py @@ -0,0 +1,152 @@ +"""Regression test for SchurSolverFull's solver="petsc" support. + +Like SchurSolver (see test_petsc_schur_solver.py), SchurSolverFull/SchurSolverFull3 always +imported feectools' own `inverse` directly, which does not recognize "petsc" -- so +solver="petsc" would raise for the variational MHD propagators built on them +(VariationalPBEvolve, VariationalEntropyEvolve, VariationalMagFieldEvolve, VariationalQBEvolve). + +Unlike SchurSolver, there is no in-place-mutation caching hazard here: `self._S` is built once +in __init__ and never mutated afterwards -- callers (e.g. VariationalQBEvolve) rebuild the whole +SchurSolverFull3 object fresh every Newton iteration rather than reusing one across calls (see +those propagators' "local version to avoid creating new version of LinearOperator every time" +comment, which refers to the *Jacobian's blocks*, not to reusing the Schur solver object itself). +So the fix here is the dispatch alone: use struphy.linear_algebra.solver.inverse (which knows +"petsc" and safely strips petsc-only kwargs for every other solver) instead of feectools' own. + +This test exercises that dispatch directly on a small synthetic block system (matching +test_petsc_solver.py's style), not through a real variational MHD model: those models' Jacobian +blocks involve operator types (BasisProjectionOperator-derived, nonlinear-model-specific) that +have not been checked against _assemble_petsc_matrix's supported set, and building one from +scratch without an existing example/test to adapt was judged too failure-prone to do blind. If a +real variational-MHD case is wired up to use solver="petsc" later, verify it separately. +""" + +import cunumpy as xp +import pytest + +pytest.importorskip("petsc4py") + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import IdentityOperator +from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace +from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace + +from struphy.linear_algebra.schur_solver import SchurSolverFull, SchurSolverFull3 + + +def _make_space(n, p): + domain_decomposition = DomainDecomposition([n - p], [False], comm=MPI.COMM_WORLD) + cart = CartDecomposition(domain_decomposition, [n], [xp.array([0])], [xp.array([n - 1])], [p], [1]) + return StencilVectorSpace(cart) + + +def _spd_tridiagonal(V, p, scale): + """Banded, symmetric positive-definite StencilMatrix with 2p+1 diagonals on space `V`.""" + A = StencilMatrix(V, V) + A[:, -p:0] = -scale + A[:, 0:1] = 2 * p * scale + A[:, 1 : p + 1] = -scale + A.remove_spurious_entries() + return A + + +def test_schur_solver_full_petsc_matches_pcg(): + """SchurSolverFull(solver_name="petsc") must solve [[A B],[C Id]] to the same accuracy as + solver_name="pcg", for a small synthetic system [[A B],[C Id]] x = v. + """ + n, p = 12, 1 + xp.random.seed(0) + + V = _make_space(n, p) + A = _spd_tridiagonal(V, p, scale=1.0) + B = _spd_tridiagonal(V, p, scale=0.01) + C = _spd_tridiagonal(V, p, scale=0.01) + + domain = BlockVectorSpace(V, V) + M = BlockLinearOperator(domain, domain) + M[0, 0] = A + M[0, 1] = B + M[1, 0] = C + M[1, 1] = IdentityOperator(V) + + s = V.starts[0] + e = V.ends[0] + bx = StencilVector(V) + bx[s : e + 1] = xp.random.random(e + 1 - s) + by = StencilVector(V) + by[s : e + 1] = xp.random.random(e + 1 - s) + + v = BlockVector(domain) + v[0] = bx + v[1] = by + + solver_kwargs = {"pc": None, "tol": 1e-13, "maxiter": 2000, "verbose": False, "recycle": False} + solver_pcg = SchurSolverFull(M, "pcg", **solver_kwargs) + solver_petsc = SchurSolverFull(M, "petsc", **solver_kwargs) + + x_pcg = solver_pcg.dot(v) + x_petsc = solver_petsc.dot(v) + + err_x = xp.linalg.norm((x_pcg[0] - x_petsc[0]).toarray()) + err_y = xp.linalg.norm((x_pcg[1] - x_petsc[1]).toarray()) + assert err_x < 1e-8, f"x-block mismatch: {err_x:.2e}" + assert err_y < 1e-8, f"y-block mismatch: {err_y:.2e}" + + +def test_schur_solver_full3_petsc_matches_pcg(): + """SchurSolverFull3(solver_name="petsc") must solve [[A B D],[C Id 0],[E 0 Id]] to the same + accuracy as solver_name="pcg", for a small synthetic system. + """ + n, p = 12, 1 + xp.random.seed(1) + + V = _make_space(n, p) + A = _spd_tridiagonal(V, p, scale=1.0) + B = _spd_tridiagonal(V, p, scale=0.01) + C = _spd_tridiagonal(V, p, scale=0.01) + D = _spd_tridiagonal(V, p, scale=0.01) + E = _spd_tridiagonal(V, p, scale=0.01) + + domain = BlockVectorSpace(V, V, V) + M = BlockLinearOperator(domain, domain) + M[0, 0] = A + M[0, 1] = B + M[1, 0] = C + M[1, 1] = IdentityOperator(V) + M[0, 2] = D + M[2, 0] = E + M[2, 2] = IdentityOperator(V) + + s = V.starts[0] + e = V.ends[0] + bx = StencilVector(V) + bx[s : e + 1] = xp.random.random(e + 1 - s) + by = StencilVector(V) + by[s : e + 1] = xp.random.random(e + 1 - s) + bz = StencilVector(V) + bz[s : e + 1] = xp.random.random(e + 1 - s) + + v = BlockVector(domain) + v[0] = bx + v[1] = by + v[2] = bz + + solver_kwargs = {"pc": None, "tol": 1e-13, "maxiter": 2000, "verbose": False, "recycle": False} + solver_pcg = SchurSolverFull3(M, "pcg", **solver_kwargs) + solver_petsc = SchurSolverFull3(M, "petsc", **solver_kwargs) + + x_pcg = solver_pcg.dot(v) + x_petsc = solver_petsc.dot(v) + + err_x = xp.linalg.norm((x_pcg[0] - x_petsc[0]).toarray()) + err_y = xp.linalg.norm((x_pcg[1] - x_petsc[1]).toarray()) + err_z = xp.linalg.norm((x_pcg[2] - x_petsc[2]).toarray()) + assert err_x < 1e-8, f"x-block mismatch: {err_x:.2e}" + assert err_y < 1e-8, f"y-block mismatch: {err_y:.2e}" + assert err_z < 1e-8, f"z-block mismatch: {err_z:.2e}" + + +if __name__ == "__main__": + test_schur_solver_full_petsc_matches_pcg() + test_schur_solver_full3_petsc_matches_pcg() From 67576d0e7997116a56ab5bf338390809624671e2 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 11 Aug 2026 13:47:43 +0200 Subject: [PATCH 32/32] bugfix for near singular systems --- src/struphy/linear_algebra/petsc_solver.py | 23 +++++++++++++++++++ src/struphy/linear_algebra/solver.py | 7 +++--- .../tests/test_petsc_poisson_solve_pic.py | 18 ++++++++++++++- src/struphy/propagators/implicit_diffusion.py | 11 +++++++++ 4 files changed, 55 insertions(+), 4 deletions(-) diff --git a/src/struphy/linear_algebra/petsc_solver.py b/src/struphy/linear_algebra/petsc_solver.py index 3b4f5a414..408c44480 100644 --- a/src/struphy/linear_algebra/petsc_solver.py +++ b/src/struphy/linear_algebra/petsc_solver.py @@ -260,6 +260,22 @@ class PETScSolver(InverseLinearOperator): pc_type : str, default="none" PETSc preconditioner type, see ``petsc4py.PETSc.PC.Type``. E.g. ``"gamg"`` (algebraic multigrid) for large, ill-conditioned elliptic systems. + + near_null_space : {"none", "constant"}, default="none" + Registers a null space with the assembled matrix via ``Mat.setNullSpace`` when + ``"constant"`` (the all-ones vector), so KSP removes any inconsistent component from the + right-hand side instead of letting it pollute the solve. This is not just a minor + robustness tweak: for a *near*-singular operator whose kernel is (numerically) the + constant vector -- e.g. ``ImplicitDiffusion``'s ``grad.T @ M @ grad + sigma_1 * stab_mat`` + on a periodic domain with tiny ``sigma_1`` -- omitting it was found to make ``"gamg"`` + silently converge (small reported residual) to a solution that disagrees substantially + with feectools' own solver, worse and MPI-rank-count-dependent as rank count grows (up to + ~180% relative error at 4 ranks in testing), while reporting success throughout; with it, + the same cases match to ~1e-15 in 1 iteration, independent of rank count. Only pass + ``"constant"`` for operators whose kernel is actually (near) the constant vector -- + forcing it on an operator that is not near-singular there is unlikely to help, and forcing + it on one that is near-singular along some *other* direction would silently corrupt the + answer as this option does not check its own applicability. """ def __init__( @@ -273,8 +289,10 @@ def __init__( recycle=False, ksp_type="cg", pc_type="none", + near_null_space="none", ): assert isinstance(A, LinearOperator), f"PETScSolver requires a LinearOperator, got {type(A)}." + assert near_null_space in ("none", "constant"), f"Unsupported {near_null_space = }" self._options = { "x0": x0, @@ -284,6 +302,7 @@ def __init__( "recycle": recycle, "ksp_type": ksp_type, "pc_type": pc_type, + "near_null_space": near_null_space, } super().__init__(A, **self._options) @@ -301,6 +320,10 @@ def _get_ksp(self): if self._ksp is None or self._ksp_linop is not A: gmat = _assemble_petsc_matrix(A) + if self._options["near_null_space"] == "constant": + nullspace = PETSc.NullSpace().create(constant=True, comm=gmat.getComm()) + gmat.setNullSpace(nullspace) + if self._ksp is None: self._ksp = PETSc.KSP().create(comm=gmat.getComm()) diff --git a/src/struphy/linear_algebra/solver.py b/src/struphy/linear_algebra/solver.py index 3287f73e1..06230321d 100644 --- a/src/struphy/linear_algebra/solver.py +++ b/src/struphy/linear_algebra/solver.py @@ -6,7 +6,7 @@ logger = logging.getLogger("struphy") # kwargs accepted by struphy.linear_algebra.petsc_solver.PETScSolver.__init__ -_PETSC_SOLVER_KWARGS = ("x0", "tol", "maxiter", "verbose", "recycle", "ksp_type", "pc_type") +_PETSC_SOLVER_KWARGS = ("x0", "tol", "maxiter", "verbose", "recycle", "ksp_type", "pc_type", "near_null_space") def inverse(A, solver: str, **kwargs): @@ -46,11 +46,12 @@ def inverse(A, solver: str, **kwargs): return PETScSolver(A, **petsc_kwargs) - # pc_type/ksp_type are petsc-only (see _PETSC_SOLVER_KWARGS above); feectools' - # InverseLinearOperator subclasses forward unknown kwargs straight to their + # pc_type/ksp_type/near_null_space are petsc-only (see _PETSC_SOLVER_KWARGS above); + # feectools' InverseLinearOperator subclasses forward unknown kwargs straight to their # constructor and would raise on them, so they never reach this branch. kwargs.pop("pc_type", None) kwargs.pop("ksp_type", None) + kwargs.pop("near_null_space", None) from feectools.linalg.solvers import inverse as feectools_inverse diff --git a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py index 9d2f06fcc..c67379f40 100644 --- a/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py +++ b/src/struphy/linear_algebra/tests/test_petsc_poisson_solve_pic.py @@ -33,6 +33,20 @@ def test_poisson_solve_petsc_matches_pcg_with_real_pic_deposition(): particle-in-cell charge-density deposition (not a synthetic/manufactured source), reproducing the setup of examples/VlasovAmpereOneSpecies/strong_Landau_damping (Maxwellian3D background + ModesCos perturbation, control-variate weights). + + Compared *mean-removed* (matching struphy.linear_algebra.petsc_examples_benchmark's + methodology): this case's background is a large uniform density (n=1.0 everywhere), so the + charge density's mean/DC component is large, and the near-singular stab_eps=1e-8 regularizes + the constant/DC mode only very weakly -- feectools' pcg divides that large DC charge by the + tiny stab_eps, landing on an essentially arbitrary large DC potential offset (~-168 in + testing) that is numerical-noise-amplification, not a physically meaningful answer. PETScSolver + now registers the constant mode as a near null space for exactly this operator (see + PETScSolver's near_null_space docstring and ImplicitDiffusion.allocate's near_null_space="constant" + comment) -- the fix for a real MPI-rank-dependent correctness bug this same near-singular + regime caused under >1 rank -- which makes it correctly and robustly discard that + inconsistent DC component instead (mean exactly 0) rather than reproducing pcg's arbitrary + noise-amplified one. The physically meaningful, oscillatory part of the solution still needs + to match to near machine precision, which is what this test actually checks. """ comm = MPI.COMM_WORLD @@ -101,7 +115,9 @@ def run(solver_name): sol_pcg = run("pcg") sol_petsc = run("petsc") - rel_err = xp.linalg.norm(sol_pcg - sol_petsc) / xp.linalg.norm(sol_pcg) + mean_removed_pcg = sol_pcg - sol_pcg.mean() + mean_removed_petsc = sol_petsc - sol_petsc.mean() + rel_err = xp.linalg.norm(mean_removed_pcg - mean_removed_petsc) / xp.linalg.norm(sol_pcg) assert rel_err < 1e-6 diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index c62522adb..e1e45d1ad 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -363,6 +363,17 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: verbose=self.options.solver_params.verbose, recycle=self.options.solver_params.recycle, pc_type=self.options.solver_params.pc_type, + # self._diffusion_op = grad.T @ diffusion_mat @ grad structurally has the constant + # function in its kernel on a periodic domain (grad(constant) = 0), regardless of + # diffusion_mat or how small/large sigma_1 (stab_eps) is -- and PETScSolver only + # supports this operator on periodic domains to begin with (see + # _directional_derivative_to_stencil_matrix), so this is always a valid hint where it + # applies at all. Ignored for solver != "petsc". Without it, PETSc+gamg was found to + # silently converge (small reported residual) to a solution that disagrees with + # feectools' own solver -- worse, and more MPI-rank-count-dependent, as rank count + # grows -- for exactly this near-singular regime; see PETScSolver's near_null_space + # docstring. + near_null_space="constant", ) # allocate memory for solution