Scalars fetch allocated objects again after each allocation: no stale scalars in a second run (fixes #852)

src/struphy/models/linear_vlasov_ampere_one_species.py

@@ -131,22 +131,19 @@ def post_allocate(self):
         # the periodic Poisson problem below singular.
         self._solve_initial_poisson()
 
+        # background f0 and its values at the markers, for en_f (the particles are new at each allocation)
+        backgrounds = self.kinetic_ions.var.backgrounds
+        if isinstance(backgrounds, list):
+            self._f0 = backgrounds[0]
+        else:
+            self._f0 = backgrounds
+        assert isinstance(self._f0, Maxwellian3D)
+        self._f0_values = xp.zeros(self.kinetic_ions.var.particles.n_markers_max, dtype=float)
+
     def _compute_en_f(self):
         particles = self.kinetic_ions.var.particles
 
         # evaluate f0
-        if not hasattr(self, "_f0"):
-            backgrounds = self.kinetic_ions.var.backgrounds
-            if isinstance(backgrounds, list):
-                self._f0 = backgrounds[0]
-            else:
-                self._f0 = backgrounds
-            self._f0_values = xp.zeros(
-                self.kinetic_ions.var.particles.n_markers_max,
-                dtype=float,
-            )
-            assert isinstance(self._f0, Maxwellian3D)
-
         self._f0_values[particles.valid_mks] = self._f0(*particles.phasespace_coords)
 
         # alpha^2 * v_th^2 * N / 2 * sum_p s_0 * w_p^2 / f_{0,p}  (weights are w_p = delta f_p / (N * s_{0,p}))

src/struphy/models/scalars.py

@@ -33,13 +33,31 @@ class Scalar(metaclass=ABCMeta):
     Parameters
     ----------
     variables : Variable or Scalar
-        The variable(s) associated with the scalar, or scalars for summation."""
+        The variable(s) associated with the scalar, or scalars for summation.
+
+    Notes
+    -----
+    Objects created by :meth:`Simulation.allocate <struphy.simulation.sim.Simulation.allocate>` (spline vectors,
+    mass operators, the projected equilibrium, particles, ...) must not be fetched in ``__init__`` (the scalars are
+    created before the allocation) and must not be cached in :meth:`_local_update`: each
+    :meth:`Simulation.run <struphy.simulation.sim.Simulation.run>` allocates them again, and a cached object would
+    keep the values of the first run. Fetch them in :meth:`allocate` instead, which :meth:`update` calls once after
+    each allocation (see :meth:`Scalars.reset_allocation`).
+    """
 
     def __init__(self, *variables: Union[Variable, "Scalar"]):
         self.variables = variables
         self.local_value = xp.empty(1, dtype=float)
         self.value = xp.empty(1, dtype=float)
         self.uptodate = False
+        self.allocated = False
+
+    def allocate(self):
+        """Fetch the objects created by the allocation of the simulation that :meth:`_local_update` needs.
+
+        Called by :meth:`update` before the first update after each allocation. Does nothing by default.
+        """
+        pass
 
     @abstractmethod
     def _local_update(self):
@@ -53,6 +71,9 @@ def _mpi_sum(self):
 
     def update(self):
         """Update the scalar quantity by performing local update and then summing over MPI processes."""
+        if not self.allocated:
+            self.allocate()
+            self.allocated = True
         if not self.uptodate:
             self._local_update()
             self._mpi_sum()
@@ -99,6 +120,9 @@ def _local_update(self):
             "Subclasses of PICScalar must implement _local_update to compute self.local_value[0]."
         )
 
+    def allocate(self):
+        self.clone_config = self.variables[0].particles.clone_config
+
     def _mpi_sum(self):
         self.value[0] = self.local_value[0]
 
@@ -111,9 +135,6 @@ def _mpi_sum(self):
             )
 
         # sum between clones
-        if not hasattr(self, "clone_config"):
-            self.clone_config = self.variables[0].particles.clone_config
-
         if self.clone_config is not None:
             self.clone_config.inter_comm.Allreduce(
                 MPI.IN_PLACE,
@@ -179,6 +200,15 @@ def __init__(self, **scalars: dict[str, Scalar]):
     def dct(self) -> dict[str, Scalar]:
         return self._dct
 
+    def reset_allocation(self):
+        """Make every scalar (including the summands of sums) fetch the allocated objects again at its next update.
+
+        Called by :meth:`Simulation.allocate <struphy.simulation.sim.Simulation.allocate>`, which creates new spline
+        vectors, mass operators, particles, ... at each run.
+        """
+        for scalar in self.dct.values():
+            _reset_allocation(scalar)
+
     def update(self):
         for scalar in self.dct.values():
             scalar.update()
@@ -201,6 +231,14 @@ def host_value(self, name: str) -> np.ndarray:
         return self._host_values[i : i + 1]
 
 
+def _reset_allocation(scalar: Scalar):
+    scalar.allocated = False
+    scalar.uptodate = False
+    for variable in scalar.variables:
+        if isinstance(variable, Scalar):
+            _reset_allocation(variable)
+
+
 def _mark_outdated(scalar: Scalar):
     scalar.uptodate = False
     for variable in scalar.variables:
@@ -247,18 +285,17 @@ def __init__(
         self.bilinear_form_name = bilinear_form_name
         self.normalization = normalization
 
-    def _local_update(self):
-        if not hasattr(self, "left_vec"):
-            self.left_vec = self.variables[0].spline.vector
-            if isinstance(self.variables[1], str):
-                self.right_vec = getattr(Propagator.projected_equil, self.variables[1])
-            else:
-                self.right_vec = self.variables[1].spline.vector
-            self.vec_space = self.left_vec.space
-        if not hasattr(self, "bilinear_form"):
-            self.bilinear_form: WeightedMassOperator = getattr(Propagator.mass_ops, self.bilinear_form_name)
-            assert self.bilinear_form.codomain == self.vec_space, "bilinear_form codomain must match variable space"
+    def allocate(self):
+        self.left_vec = self.variables[0].spline.vector
+        if isinstance(self.variables[1], str):
+            self.right_vec = getattr(Propagator.projected_equil, self.variables[1])
+        else:
+            self.right_vec = self.variables[1].spline.vector
+        self.vec_space = self.left_vec.space
+        self.bilinear_form: WeightedMassOperator = getattr(Propagator.mass_ops, self.bilinear_form_name)
+        assert self.bilinear_form.codomain == self.vec_space, "bilinear_form codomain must match variable space"
 
+    def _local_update(self):
         value = self.normalization * 0.5 * self.bilinear_form.dot_inner(self.right_vec, self.left_vec)
         self.local_value[0] = value
 
@@ -288,16 +325,16 @@ def __init__(
         super().__init__(feec_variable)
         self.normalization = normalization
 
-    def _local_update(self):
-        if not hasattr(self, "vec"):
-            self.vec = self.variables[0].spline.vector
-            if isinstance(self.vec, PolarVector):
-                self.ones = Propagator.derham.V3pol.zeros()
-                self.ones.tp[:] = 1.0
-            else:
-                self.ones = Propagator.derham.V3.zeros()
-                self.ones[:] = 1.0
+    def allocate(self):
+        self.vec = self.variables[0].spline.vector
+        if isinstance(self.vec, PolarVector):
+            self.ones = Propagator.derham.V3pol.zeros()
+            self.ones.tp[:] = 1.0
+        else:
+            self.ones = Propagator.derham.V3.zeros()
+            self.ones[:] = 1.0
 
+    def _local_update(self):
         self.local_value[0] = self.normalization * self.vec.inner(self.ones)
 
     def _mpi_sum(self):

src/struphy/models/tests/test_scalars.py

@@ -36,6 +36,50 @@ def test_nested_sum_follows_its_summands():
     assert scalars.dct["total"].value[0] == 60.0
 
 
+class _AllocatedConstant(Scalar):
+    """A scalar that reads its value from an object of `holder` created by the allocation, as BilinearEnergyFEEC
+    reads the spline vector of a variable."""
+
+    def __init__(self, holder):
+        super().__init__()
+        self.holder = holder
+
+    def allocate(self):
+        self.data = self.holder.data
+
+    def _local_update(self):
+        self.local_value[0] = self.data[0]
+
+    def _mpi_sum(self):
+        self.value[0] = self.local_value[0]
+
+
+def test_scalars_follow_a_new_allocation():
+    """A second Simulation.run() allocates new objects; the scalars (and the summands of sums) must use those,
+    not the ones of the first run (issue #852)."""
+    from types import SimpleNamespace
+
+    holder = SimpleNamespace(data=xp.array([1.0]))
+    a, b = _AllocatedConstant(holder), _Constant(2.0)
+    scalars = Scalars(a=a, total=a + b)
+
+    scalars.update()
+    assert scalars.host_value("a")[0] == 1.0
+    assert scalars.host_value("total")[0] == 3.0
+
+    # second run: the allocation replaces the object, the old one keeps its last value
+    holder.data = xp.array([10.0])
+    scalars.reset_allocation()
+    scalars.update()
+    assert scalars.host_value("a")[0] == 10.0
+    assert scalars.host_value("total")[0] == 12.0
+
+    # within a run, the object fetched at the allocation is used
+    holder.data[0] = 20.0
+    scalars.update()
+    assert scalars.host_value("a")[0] == 20.0
+
+
 class _MovingParticles:
     """Markers whose accessors return copies, as `Particles.velocities` does (it is fancy-indexed)."""
 

src/struphy/models/variational_compressible_fluid.py

@@ -113,6 +113,10 @@ def f(e1, e2, e3):
 
         self._energy_evaluator = InternalEnergyEvaluator(Propagator.derham, self.propagators.variat_ent.options.gamma)
 
+        if self.params["divdiv_alpha"] > 0.0:
+            self._divUv = Propagator.derham.div @ Propagator.basis_ops.Uv
+            self._tmp_divu = self._divUv.codomain.zeros()
+
     # default parameters
     def generate_default_parameter_file(self, path=None, prompt=True):
         params_path = super().generate_default_parameter_file(path=path, prompt=prompt)
@@ -138,9 +142,6 @@ def _compute_divdiv_energy(self):
 
         :meta private:
         """
-        if not hasattr(self, "_divUv"):
-            self._divUv = Propagator.derham.div @ Propagator.basis_ops.Uv
-            self._tmp_divu = self._divUv.codomain.zeros()
         d = self._divUv.dot(self.fluid.velocity.spline.vector, out=self._tmp_divu)
         return self.params["divdiv_alpha"] * Propagator.mass_ops.WMM3new.dot_inner(d, d)
 

src/struphy/models/variational_mhd.py

@@ -171,14 +171,15 @@ def f(e1, e2, e3):
 
         self._tmp_div_B = Propagator.derham.V3pol.zeros()
 
+        if self.params["divdiv_alpha"] > 0.0:
+            self._divUv = Propagator.derham.div @ Propagator.basis_ops.Uv
+            self._tmp_divu = self._divUv.codomain.zeros()
+
     def _compute_divdiv_energy(self):
         """alpha int rho (div u)^2 dx, the div-div part of the kinetic energy (see kinetic_metric).
 
         :meta private:
         """
-        if not hasattr(self, "_divUv"):
-            self._divUv = Propagator.derham.div @ Propagator.basis_ops.Uv
-            self._tmp_divu = self._divUv.codomain.zeros()
         d = self._divUv.dot(self.mhd.velocity.spline.vector, out=self._tmp_divu)
         return self.params["divdiv_alpha"] * Propagator.mass_ops.WMM3new.dot_inner(d, d)
 

src/struphy/simulation/sim.py

@@ -262,6 +262,10 @@ def allocate(self):
             with ProfileManager.profile_region("model post allocate", functions=[self.model.post_allocate]):
                 self.model.post_allocate()
 
+            # the scalars fetch the newly allocated objects at their next update (a second run allocates again)
+            if self.model.scalars is not None:
+                self.model.scalars.reset_allocation()
+
         logger.debug("... Done.")
 
     def kernels(self) -> tuple[Kernel, ...]: