Fix total energy at t=0 in Vlasov-Ampère/Maxwell models (ghost regions, control variate)

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)

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):

src/struphy/models/scalars.py

@@ -333,6 +333,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``).

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)

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):

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

src/struphy/pic/base.py

@@ -1423,6 +1423,11 @@ def draw_markers(
             logger.info("\nSorting the markers into boxes after initial draw")
             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,