From 4125aaa87b91b24357a61ea27e27cd4e793b18dd Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 8 Oct 2026 23:00:42 +0200 Subject: [PATCH 1/4] Fix total energy at t=0 in Vlasov-Ampere/Maxwell models (ghost regions, control variate) Two bugs made total_energy at t=0 inconsistent with the later time steps: 1. The initial field e = -grad(phi) set in post_allocate did not have its ghost regions updated (grad.dot(..., out=e) does not fill them). The first VlasovAmpereCoupling push then evaluated a wrong field near the periodic boundary, so the Crank-Nicolson step did not conserve energy in step 1 only (full-f jump of ~2e-4 relative on 8 elements). Now synced in VlasovAmpereOneSpecies, VlasovMaxwellOneSpecies, ColdPlasmaVlasov and LinearVlasovAmpereOneSpecies. 2. With the control variate, post_allocate reset the weights to the full-f weights0 after the initial Poisson solve, so kinetic_energy at t=0 was the full-f energy and from step 1 on the delta-f energy (total_energy dropped from ~19 to ~0.1). The weights are now reset to weights0 only without the control variate (VlasovAmpereOneSpecies, VlasovMaxwellOneSpecies incl. its Gauss-law diagnostic). ColdPlasmaVlasov never reset them, so full-f runs used delta-f weights; it now follows the same rule. Adds test_energy_conservation (full-f and control variate) to the VlasovAmpereOneSpecies verification tests. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/cold_plasma_vlasov.py | 10 +- .../linear_vlasov_ampere_one_species.py | 3 + src/struphy/models/scalars.py | 4 + .../test_verif_VlasovAmpereOneSpecies.py | 97 +++++++++++++++++++ .../models/vlasov_ampere_one_species.py | 11 ++- .../models/vlasov_maxwell_one_species.py | 17 ++-- 6 files changed, 132 insertions(+), 10 deletions(-) diff --git a/src/struphy/models/cold_plasma_vlasov.py b/src/struphy/models/cold_plasma_vlasov.py index d6fe0986c..0abf622cc 100644 --- a/src/struphy/models/cold_plasma_vlasov.py +++ b/src/struphy/models/cold_plasma_vlasov.py @@ -204,7 +204,7 @@ def post_allocate(self): """ logger.info("\nINITIAL POISSON SOLVE:") - # use control variate method + # charge of f - f0 for the Poisson right-hand side (see below for the weights of the time stepping) particles = self.hot_elec.var.particles particles.update_weights() @@ -216,8 +216,16 @@ def post_allocate(self): phi = self.initial_poisson.variables.phi.spline.vector Propagator.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + # The pushers evaluate e at the markers, including in the ghost cells: grad.dot does not fill them, + # so without this sync the first time step pushes with a wrong field (and breaks energy conservation). + self.em_fields.e_field.spline.vector.update_ghost_regions() logger.info("... Done.") + # The Poisson right-hand side is the charge of f - f0 in any case; the time stepping uses the delta-f + # weights only with the control variate (and the scalars must be computed with the same weights at t=0). + if not particles.control_variate: + particles.weights = particles.weights0.copy() + ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): params_path = super().generate_default_parameter_file(path=path, prompt=prompt) diff --git a/src/struphy/models/linear_vlasov_ampere_one_species.py b/src/struphy/models/linear_vlasov_ampere_one_species.py index 52db389a3..8644784de 100644 --- a/src/struphy/models/linear_vlasov_ampere_one_species.py +++ b/src/struphy/models/linear_vlasov_ampere_one_species.py @@ -186,6 +186,9 @@ def post_allocate(self): phi = self.initial_poisson.variables.phi.spline.vector Propagator.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + # The pushers evaluate e at the markers, including in the ghost cells: grad.dot does not fill them, + # so without this sync the first time step pushes with a wrong field (and breaks energy conservation). + self.em_fields.e_field.spline.vector.update_ghost_regions() logger.info("... Done.") def _compute_en_w(self): diff --git a/src/struphy/models/scalars.py b/src/struphy/models/scalars.py index 065709a12..ced4a7351 100644 --- a/src/struphy/models/scalars.py +++ b/src/struphy/models/scalars.py @@ -294,6 +294,10 @@ class KineticEnergyPIC(PICScalar): estimate (:meth:`~struphy.pic.base.Particles.initialize_weights` sets :math:`w_i = f_i / (s_i N_p)`), so the sum must not be divided by :math:`N_p` again. + With the control variate (``WeightsParameters(control_variate=True)``) the weights are the delta-f weights + :math:`w_i = (f_i - f_{0,i}) / (s_i N_p)` (:meth:`~struphy.pic.base.Particles.update_weights`), at all times + including :math:`t=0`, so this is the kinetic energy of :math:`f - f_0`. + For :class:`~struphy.pic.particles.Particles5D` the velocity coordinates are :math:`(v_\parallel, \mu)`, so only the parallel kinetic energy :math:`v_i^2 = v_{\parallel,i}^2` is summed; the magnetic-moment part :math:`\mu_i |B_0|` is tracked by a separate scalar in the models (e.g. ``en_fB``). diff --git a/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py index 0d978a5a1..037ba04c0 100644 --- a/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py +++ b/src/struphy/models/tests/verification/test_verif_VlasovAmpereOneSpecies.py @@ -4,6 +4,7 @@ import cunumpy as xp import h5py +import pytest from matplotlib import pyplot as plt from maybempi import MPI @@ -24,6 +25,7 @@ maxwellians, perturbations, ) +from struphy.linear_algebra.solver import SolverParameters from struphy.models import VlasovAmpereOneSpecies logger = logging.getLogger("struphy") @@ -166,5 +168,100 @@ def E_exact(t): shutil.rmtree(test_folder) +@pytest.mark.parametrize("control_variate", [False, True]) +def test_energy_conservation(control_variate: bool, exit_before_run: bool = False): + """``total_energy`` is conserved from :math:`t=0` on, with full-f and with delta-f (control variate) weights. + + A time step is :class:`~struphy.propagators.push_eta.PushEta` followed by the Crank-Nicolson step + :class:`~struphy.propagators.vlasov_ampere_coupling.VlasovAmpereCoupling` (no B0). This checks two things that + went wrong at :math:`t=0` only: + + * the initial field ``e = -grad(phi)`` must have its ghost regions in sync, otherwise the first push evaluates a + wrong field and the total energy jumps in the first step (by ~1e-6 relative here, ~2e-4 with 8 elements); + * with the control variate, ``kinetic_energy`` at :math:`t=0` must be computed with the delta-f weights, like + at all later times (it was computed with the full-f weights, i.e. the energy of the background too). + + With full-f weights the scheme conserves the total energy to the solver tolerance. With the control variate + the weights are updated after the Crank-Nicolson step, :math:`w_p = w_{0,p} - f_0(\\mathbf v_p) / (s_{0,p} N)`, + which changes the delta-f kinetic energy by a Monte-Carlo estimate of zero plus an :math:`O(\\Delta t^2)` term + per step (the energy that the field gives to the background markers). So the total delta-f energy is + conserved only up to noise and time discretization error, and is checked with a loose tolerance and more + markers. + """ + amplitude = 0.1 + ppc = 400 if control_variate else 40 + + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + test_folder = os.path.join(os.getcwd(), "struphy_verification_tests") + out_folders = os.path.join(test_folder, "VlasovAmpereOneSpecies") + env = EnvironmentOptions(out_folders=out_folders, sim_folder=f"energy_conservation_cv_{control_variate}") + + time_opts = Time(dt=0.1, Tend=0.3) + + domain = domains.Cuboid(r1=12.56) + # 16 elements, so that every rank has at least p=3 of them with 4 ranks (CI runs with clones on 4 ranks) + grid = grids.TensorProductGrid(num_elements=(16, 1, 1)) + derham_opts = DerhamOptions(degree=(3, 1, 1)) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=ppc, loading="sobol_standard"), + weights_params=WeightsParameters(control_variate=control_variate), + boundary_params=BoundaryParameters(), + sorting_params=SortingParameters(), + saving_params=SavingParameters(), + bufsize=0.4, + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options( + solver_params=SolverParameters(tol=1e-12), + ) + model.initial_poisson.options = model.initial_poisson.Options( + stab_mat="M0", + stab_eps=1e-6, + solver_params=SolverParameters(tol=1e-12), + ) + + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + perturbation = perturbations.ModesCos(ls=(1,), amps=(amplitude,)) + model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + sim = Simulation( + model=model, + env=env, + time_opts=time_opts, + domain=domain, + grid=grid, + derham_opts=derham_opts, + ) + + if exit_before_run: + logger.info("Exiting before running simulation.") + return sim + + sim.run() + + comm = MPI.COMM_WORLD + + if comm.Get_rank() == 0: + with h5py.File(os.path.join(env.path_out, "data", "data_proc0.hdf5"), "r") as f: + en_E = f["scalar"]["electric_energy"][()] + en_tot = f["scalar"]["total_energy"][()] + + # energy is exchanged between field and particles ... + assert abs(en_E[-1] - en_E[0]) > 1e-2 * en_E[0], f"{en_E =}" + + # ... and the total is conserved, including the first step + drift = float(xp.max(xp.abs(en_tot - en_tot[0])) / abs(en_tot[0])) + tol = 0.1 if control_variate else 1e-10 + assert drift < tol, f"Total energy not conserved: {en_tot =}, {drift =}." + logger.info(f"Assertion for energy conservation passed ({control_variate =}, {drift =}).") + + shutil.rmtree(env.path_out) + + if __name__ == "__main__": test_weak_Landau(do_plot=True) + test_energy_conservation(control_variate=False) + test_energy_conservation(control_variate=True) diff --git a/src/struphy/models/vlasov_ampere_one_species.py b/src/struphy/models/vlasov_ampere_one_species.py index 34b24dd06..ca63b26fb 100644 --- a/src/struphy/models/vlasov_ampere_one_species.py +++ b/src/struphy/models/vlasov_ampere_one_species.py @@ -166,7 +166,7 @@ def post_allocate(self): """ logger.info("\nINITIAL POISSON SOLVE:") - # use control variate method (reset weights after Poisson solve) + # charge of f - f0 for the Poisson right-hand side (see below for the weights of the time stepping) particles = self.kinetic_ions.var.particles particles.update_weights() @@ -178,10 +178,15 @@ def post_allocate(self): phi = self.initial_poisson.variables.phi.spline.vector Propagator.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + # The pushers evaluate e at the markers, including in the ghost cells: grad.dot does not fill them, + # so without this sync the first time step pushes with a wrong field (and breaks energy conservation). + self.em_fields.e_field.spline.vector.update_ghost_regions() logger.info("... Done.") - # reset particle weights - particles.weights = particles.weights0.copy() + # The Poisson right-hand side is the charge of f - f0 in any case; the time stepping uses the delta-f + # weights only with the control variate (and the scalars must be computed with the same weights at t=0). + if not particles.control_variate: + particles.weights = particles.weights0.copy() ## default parameters def generate_default_parameter_file(self, path=None, prompt=True): diff --git a/src/struphy/models/vlasov_maxwell_one_species.py b/src/struphy/models/vlasov_maxwell_one_species.py index 00ab7cdd9..055e6b445 100644 --- a/src/struphy/models/vlasov_maxwell_one_species.py +++ b/src/struphy/models/vlasov_maxwell_one_species.py @@ -181,7 +181,7 @@ def post_allocate(self): logger.info("\nINITIAL POISSON SOLVE:") - # use control variate method (reset weights after Poisson solve) + # charge of f - f0 for the Poisson right-hand side (see below for the weights of the time stepping) particles.update_weights() self.initial_poisson.allocate() @@ -195,10 +195,15 @@ def post_allocate(self): phi = self.initial_poisson.variables.phi.spline.vector Propagator.derham.grad.dot(-phi, out=self.em_fields.e_field.spline.vector) + # The pushers evaluate e at the markers, including in the ghost cells: grad.dot does not fill them, + # so without this sync the first time step pushes with a wrong field (and breaks energy conservation). + self.em_fields.e_field.spline.vector.update_ghost_regions() logger.info("... Done.") - # reset particle weights - particles.weights = particles.weights0.copy() + # The Poisson right-hand side is the charge of f - f0 in any case; the time stepping uses the delta-f + # weights only with the control variate (and the scalars must be computed with the same weights at t=0). + if not particles.control_variate: + particles.weights = particles.weights0.copy() def calculate_gauss_error(self): r"""Maximum norm of the weak Gauss-law residual @@ -215,10 +220,10 @@ def calculate_gauss_error(self): particles.update_weights() self.charge_accum() rho = self.charge_accum.vectors[0] - # reset particle weights - particles.weights = particles.weights0.copy() + # restore the weights of the time stepping (full-f without the control variate) + if not particles.control_variate: + particles.weights = particles.weights0.copy() - # non control variate method e = self.em_fields.e_field.spline.vector residual = self.op.dot(e) residual += self.initial_poisson.coeffs[0] * rho From 2677a28cf1acaf56600130d939f26300f0e871ae Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:02:42 +0200 Subject: [PATCH 2/4] Send Sobol-loaded markers to their MPI process after drawing With loading="sobol_standard"/"sobol_antithetic" every process draws its markers on the whole unit cube (pseudo-random loading keeps only those on the process domain). They were only sent to their owning process when box sorting was enabled (SortingParameters(do_sort=True, boxes_per_dim=...)). Without it, the initial charge accumulation in the Poisson solve of VlasovAmpereOneSpecies wrote outside the local stencil arrays (fill_vec index -2 / 14 for a local range 0..13, seen with pyccel --debug), corrupting the heap. On Linux this aborted the 2-rank verification job (Fatal Python error: Aborted) in the new test_energy_conservation, which is the only MPI test loading Sobol markers without box sorting. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 1d08e9116..3a806e916 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1402,6 +1402,11 @@ def draw_markers( self.mpi_sort_markers() self.do_sort() logger.info("Done.") + elif self.mpi_comm is not None and self.loading in ("sobol_standard", "sobol_antithetic"): + # Sobol markers are drawn on the whole domain on every process (pseudo-random ones only on the + # process domain): send them to their process, otherwise the first accumulation (e.g. the initial + # Poisson solve) writes outside the local arrays and corrupts the heap. + self.mpi_sort_markers() def initialize_weights( self, From 4d4b51f35db87fee555282df8364bb38ed165db7 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 08:18:35 +0200 Subject: [PATCH 3/4] Diagonal scaling: include band entries that wrap onto the diagonal In a periodic direction with fewer points than the band width (2p + 1 > n, e.g. a single element in a 1D run), StencilMatrix.diagonal() returns only the entry at offset 0, while the entries at offsets k = m n also lie on the main diagonal. The diagonal D of the mass matrix was therefore wrong, and MassMatrixPreconditioner with diagonal_scaling=True (the default since #739) was no longer exact on the logical cube; in VlasovAmpereCoupling this broke energy conservation in 1D runs. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 30 +++++++++++++++-- src/struphy/feec/tests/test_mass_matrices.py | 34 ++++++++++++++++++++ 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index d64838dc2..a3a67c70b 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -318,7 +318,8 @@ def weight_1d(c: int, d: int) -> Callable | xp.ndarray: return _kronecker_approximation(mass_operator, bc, weight_1d) def _core_diagonal(self): - return self._mass_operator.matrix.diagonal() + # the matrix itself: _local_diagonal also adds the band entries that wrap onto the diagonal + return self._mass_operator.matrix @property def dim_reduce(self) -> int | None: @@ -765,7 +766,7 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: if isinstance(op, StencilDiagonalMatrix): return [op._data] if isinstance(op, StencilMatrix): - return [op.diagonal()._data] + return [_stencil_local_diagonal(op)] if isinstance(op, KroneckerStencilMatrix): # outer product of the diagonals of the (process-local) factors on the local rows diags = [] @@ -788,6 +789,31 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: raise NotImplementedError(f"Diagonal of {type(op).__name__} is not supported.") +def _stencil_local_diagonal(M: StencilMatrix) -> xp.ndarray: + """ + Local diagonal of a StencilMatrix, including the band entries that wrap onto the main diagonal. + + In a periodic direction with fewer points than the band width (``2p + 1 > n``, e.g. a single + element), the band entries at the offsets ``k = m n`` (``m != 0``) are on the main diagonal too; + :meth:`StencilMatrix.diagonal` returns only the entry at ``k = 0``. + """ + V = M.codomain + assert all(m == 1 for m in V.shifts) and all(m == 1 for m in M.domain.shifts), "Only shifts 1 are supported." + offsets = [ + [k for k in range(-p, p + 1) if k % n == 0] if periodic else [0] + for n, p, periodic in zip(V.npts, M.pads, V.periods) + ] + diag = M.diagonal()._data + if all(o == [0] for o in offsets): + return diag + idx = M._get_diagonal_indices() + rows, cols = idx[: V.ndim], idx[V.ndim :] + for ks in itertools.product(*offsets): + if any(ks): + diag = diag + M._data[rows + tuple(j + k for j, k in zip(cols, ks))] + return diag + + def _diagonal_operator(space: VectorSpace, diags: list[xp.ndarray]) -> LinearOperator: """Diagonal operator on ``space`` (Stencil- or BlockVectorSpace) with the given local diagonals.""" if isinstance(space, BlockVectorSpace): diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index 23af76b82..799646edf 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -1742,6 +1742,40 @@ def test_transpose_and_copy(): assert xp.allclose(A.T.toarray(), -A_arr, atol=1e-14) +@pytest.mark.parametrize( + "num_elements, degree", [((16, 1, 1), (3, 1, 1)), ((16, 2, 1), (2, 2, 1)), ((8, 6, 4), (2, 2, 2))] +) +@pytest.mark.parametrize("space", ["M0", "M1", "M2", "M3"]) +def test_mass_preconditioner_diagonal_scaling_periodic_exact(num_elements, degree, space): + """On the periodic logical cube the mass matrices are Kronecker products, so MassMatrixPreconditioner with + diagonal scaling inverts them exactly; also with fewer points than the band width in a direction (one element), + where band entries wrap onto the main diagonal.""" + + import numpy as np + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import MassMatrixPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + domain = domains.Cuboid(r1=12.56) + derham = Derham(TensorProductGrid(num_elements=num_elements), DerhamOptions(degree=degree), comm=MPI.COMM_WORLD) + M = getattr(WeightedMassOperators(derham, domain), space) + + rng = np.random.default_rng(0) + v = M.domain.zeros() + for block in v.blocks if hasattr(v, "blocks") else [v]: + block._data[:] = rng.random(block._data.shape) + v.update_ghost_regions() + + w = MassMatrixPreconditioner(M, diagonal_scaling=True).dot(M.dot(v)) + err = ((w - v).inner(w - v) / v.inner(v)) ** 0.5 + assert err < 1e-12, f"{space}, {num_elements=}, {degree=}: |P M v - v| / |v| = {err}" + + if __name__ == "__main__": # test_mass( # num_elements=(32, 32, 32), From a0f1cf35558b5c985728fbad4b60915e85aa9bed Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 08:23:55 +0200 Subject: [PATCH 4/4] VlasovAmpereCoupling: update the ghost regions of the mid-point field before the push The pusher kernel reads the ghost regions of e_sum directly. They were only correct if the Schur solver happened to return its solution with valid ghost regions, which is no longer the case with the diagonal scaling of MassMatrixPreconditioner (default since #739): its last step is a StencilDiagonalMatrix.dot that writes only the owned entries. In 1D runs (one element in e2, e3) every evaluation reads ghost values, and the total energy was no longer conserved (drift ~1e-4 instead of ~1e-12). The other couplings already update the ghost regions before pushing. Co-Authored-By: Claude Opus 5.5 --- src/struphy/propagators/vlasov_ampere_coupling.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/struphy/propagators/vlasov_ampere_coupling.py b/src/struphy/propagators/vlasov_ampere_coupling.py index 0c8c76414..77a129ae0 100644 --- a/src/struphy/propagators/vlasov_ampere_coupling.py +++ b/src/struphy/propagators/vlasov_ampere_coupling.py @@ -214,6 +214,8 @@ def __call__(self, dt): self._e_sum += self.variables.e.spline.vector self._e_sum += self._e_tmp self._e_sum *= 0.5 + # the pusher kernel reads the ghost regions directly (the solver output may leave them stale) + self._e_sum.update_ghost_regions() # Update velocities self._pusher(dt)