Faster setup of the Hiptmair-Xu preconditioner (TAE example): exact stiffness diagonal, preconditioners built in the setup

examples/ShearAlfven/itpa_tae_benchmark/params_ShearAlfven.py

@@ -182,7 +182,7 @@
 derham_opts = DerhamOptions(degree=(3, 3, 3), bcs=(("dirichlet", "dirichlet"), None, None))
 
 # Profiling options
-profiling_opts = ProfilingOptions(use_line_profiler=True)
+profiling_opts = ProfilingOptions(use_line_profiler=False)
 
 # Simulation object
 sim = Simulation(
@@ -238,4 +238,4 @@
 model.mhd.velocity.add_perturbation(perturbation=perturbation_poloidal)
 
 if __name__ == "__main__":
-    sim.run()
+    sim.run(profiling_activated=True)

pyproject.toml

@@ -50,7 +50,7 @@ dependencies = [
     "pytest-testmon<=2.2.0",
     "ruff==0.16.10",
     "line_profiler<=5.0.2",
-    "scope-profiler>=0.5.0, <=0.7.2",
+    "scope-profiler>=0.5.0, <=0.7.3",
 ]
 
 [project.license]

src/struphy/feec/tests/test_stiffness_preconditioner.py

@@ -74,6 +74,53 @@ def test_stiffness_approximation_unit_cube(derivative, sigma, num_elements):
     assert _max_abs(diffs, comm) < 1e-12 * _max_abs(_local_diagonal(P.matrix), comm)
 
 
+@pytest.mark.parametrize("derivative", ["grad", "curl", "div"])
+@pytest.mark.parametrize(
+    "num_elements, degree, bcs",
+    [
+        ([6, 5, 4], [2, 3, 2], (("dirichlet", "free"), None, None)),
+        # periodic directions with fewer points than the band width (entries wrap onto the band)
+        ([6, 4, 3], [3, 3, 3], (("dirichlet", "dirichlet"), None, None)),
+    ],
+)
+def test_stiffness_local_diagonal(derivative, num_elements, degree, bcs):
+    """The diagonal of d^T M d from the stencil data of M equals the probed one (sum of mass terms, full metric)."""
+
+    from maybempi import MPI
+
+    from struphy import domains
+    from struphy.feec.mass import WeightedMassOperators
+    from struphy.feec.preconditioner import (
+        _mass_terms,
+        _probe_diagonal,
+        _stiffness_local_diagonal,
+        _terms_matrix,
+    )
+    from struphy.feec.psydac_derham import Derham
+    from struphy.io.options import DerhamOptions
+    from struphy.topology.grids import TensorProductGrid
+
+    comm = MPI.COMM_WORLD
+    derham = Derham(TensorProductGrid(num_elements=num_elements), DerhamOptions(degree=degree, bcs=bcs), comm=comm)
+    mass_ops = WeightedMassOperators(derham, domains.HollowTorus(a1=0.1, a2=1.0, R0=3.0, tor_period=3))
+
+    form = FORMS[derivative]
+    d = getattr(derham, f"{derivative}_bcfree")
+    mid = getattr(mass_ops, f"M{int(form) + 1}")
+    space, metric = {"0": ("Hcurl", ("Ginv", "sqrt_g")), "1": ("Hdiv", ("G", "1/sqrt_g")), "2": ("L2", ("1/sqrt_g",))}[
+        form
+    ]
+    weighted = mass_ops.create_weighted_mass(
+        space, space, weights=metric + (lambda e1, e2, e3: 1.0 + e1 * e2,), name="M_test", assemble=True
+    )
+    terms = _mass_terms(mid - 0.3 * weighted)
+
+    diag = _stiffness_local_diagonal(d, terms)
+    ref = _probe_diagonal(d.T @ _terms_matrix(terms) @ d, d.domain)
+    assert len(diag) == len(ref)
+    assert _max_abs([a - b for a, b in zip(diag, ref)], comm) < 1e-12 * _max_abs(ref, comm)
+
+
 @pytest.mark.parametrize("mapping", ["Cuboid", "HollowCylinder"])
 @pytest.mark.parametrize("sigma", [0.0, 1.0])
 def test_stiffness_preconditioner_grad(mapping, sigma):

src/struphy/feec/utilities.py

@@ -47,10 +47,10 @@ def __init__(self, *vec_fun):
         self._funs = vec_fun
 
     def __call__(self, e1, e2, e3):
+        # each component evaluated once (the equilibrium callables can be expensive)
+        a = [fun(e1, e2, e3) for fun in self._funs]
         # array from 2d list gives 3x3 array is in the first two indices
-        tmp = xp.array(
-            [[self._funs[m](e1, e2, e3) * self._funs[n](e1, e2, e3) for n in range(3)] for m in range(3)],
-        )
+        tmp = xp.array([[a[m] * a[n] for n in range(3)] for m in range(3)])
 
         # numpy operates on the last two indices with @
         return xp.transpose(tmp, axes=(2, 3, 4, 0, 1))
@@ -76,24 +76,18 @@ def __init__(self, *vec_fun):
         assert len(vec_fun) == 3
         assert all([callable(fun) for fun in vec_fun])
 
-        self._cross_mask = [
-            [0, -1, 1],
-            [1, 0, -1],
-            [-1, 1, 0],
-        ]
-
-        self._funs = [
-            [lambda e1, e2, e3: 0 * e1, vec_fun[2], vec_fun[1]],
-            [vec_fun[2], lambda e1, e2, e3: 0 * e2, vec_fun[0]],
-            [vec_fun[1], vec_fun[0], lambda e1, e2, e3: 0 * e3],
-        ]
+        self._funs = vec_fun
 
     def __call__(self, e1, e2, e3):
+        # each component evaluated once (the equilibrium callables can be expensive)
+        a1, a2, a3 = (fun(e1, e2, e3) for fun in self._funs)
+        zero = 0.0 * a1
         # array from 2d list gives 3x3 array is in the first two indices
         tmp = xp.array(
             [
-                [self._cross_mask[m][n] * fun(e1, e2, e3) for n, fun in enumerate(row)]
-                for m, row in enumerate(self._funs)
+                [zero, -a3, a2],
+                [a3, zero, -a1],
+                [-a2, a1, zero],
             ],
         )
 

src/struphy/propagators/base.py

@@ -101,6 +101,17 @@ def __call__(self, dt: float):
             Time step size.
         """
 
+    def prepare_time_step(self, dt: float):
+        """Set up the data that depends on the time step (e.g. preconditioners) for ``dt``, before the first call
+        with ``dt``, so that it is part of the setup rather than of the first time step (see
+        :meth:`~struphy.simulation.sim.Simulation.run`). Optional; does nothing by default.
+
+        Parameters
+        ----------
+        dt : float
+            Time step size of the first call.
+        """
+
     @property
     def _solve_region(self) -> str:
         """Name of the profiling region for the linear solve(s) of this propagator."""

src/struphy/propagators/shear_alfven_curl_curl_propagator.py

@@ -6,6 +6,7 @@
 from feectools.linalg.solvers import inverse
 from line_profiler import profile
 from maybempi import MPI
+from scope_profiler import ProfileManager
 
 from struphy.feec import preconditioner
 from struphy.feec.kronecker_block import basis_projection_operator, derivative_operator
@@ -131,7 +132,8 @@ class Options(OptionsBase):
         precond : LiteralOptions.OptsShearAlfvenCurlCurlPrecond, default="HiptmairXu"
             Preconditioner of the Schur system: Hiptmair-Xu for the curl-curl operator with ``M2gyro`` in the middle
             (``"HiptmairXu"``, implicit algorithm with ``solver="pcg"`` only, no polar splines; it depends on the time
-            step, so it is built at the first time step and rebuilt when the time step changes), or a preconditioner
+            step, so it is built in the setup (:meth:`prepare_time_step`) and rebuilt when the time step changes), or a
+            preconditioner
             of the mass matrix ``M1`` (``"MassMatrixPreconditioner"``, ``"PolarMassMatrixPreconditioner"``). With
             ``algo="explicit"``, ``"HiptmairXu"`` is replaced by ``"MassMatrixPreconditioner"`` (for ``M1``).
         aux_correction : bool, default=False
@@ -341,38 +343,46 @@ def weight(e1, e2, e3, m, n):
             (Pi_par, Pi_par.T, [kron_par, kron_curl @ kron_par]),
         ]
 
+    @profile
     def _build_precond(self, dt: float):
         """Hiptmair-Xu preconditioner of S ~ dt^2/4 (C^T M2gyro C + 4/dt^2 M1) (sigma = 4/dt^2 with the unscaled M2gyro in
         the middle, as the vector auxiliary space correction of Hiptmair-Xu assumes a unit-scale curl-curl term), plus
         the field-aligned auxiliary space corrections N P N^T, P the Kronecker preconditioner of
         N^T (M1 + dt^2/4 C^T M2gyro C) N on V0."""
         sigma = 4.0 / dt**2
-        hx = HiptmairXuPreconditioner(
-            None, self.mass_ops, sigma=sigma, mid="M2gyro", aux_weight=self.options.hx_aux_weight
-        )
+        with ProfileManager.profile_region("init HiptmairXu"):
+            hx = HiptmairXuPreconditioner(
+                None, self.mass_ops, sigma=sigma, mid="M2gyro", aux_weight=self.options.hx_aux_weight
+            )
         pc = sigma * hx
         if self.options.aux_correction:
-            mid = (dt**2 / 4) * self.mass_ops.M2gyro
-            corrections = [
-                (N, NT, preconditioner.StiffnessPreconditioner(self.mass_ops, ops, sigma=0.0, mid=["M1", mid]))
-                for N, NT, ops in self._aux_spaces
-            ]
-            pc = preconditioner.AuxiliarySpacePreconditioner(pc, corrections)
+            with ProfileManager.profile_region("add auxiliary space corrections"):
+                mid = (dt**2 / 4) * self.mass_ops.M2gyro
+                corrections = [
+                    (N, NT, preconditioner.StiffnessPreconditioner(self.mass_ops, ops, sigma=0.0, mid=["M1", mid]))
+                    for N, NT, ops in self._aux_spaces
+                ]
+                pc = preconditioner.AuxiliarySpacePreconditioner(pc, corrections)
         return pc
 
+    def prepare_time_step(self, dt: float):
+        """Hiptmair-Xu preconditioner of S ~ dt^2/4 (C^T M2gyro C + 4/dt^2 M1) for the time step ``dt``: built in the
+        setup, rebuilt when dt changes (sigma = 4/dt^2 with the unscaled M2gyro in the middle, as the vector auxiliary
+        space correction of Hiptmair-Xu assumes a unit-scale curl-curl term)."""
+        if self.options.algo == "implicit" and self.options.precond == "HiptmairXu" and dt != self._hx_dt:
+            with ProfileManager.profile_region("build_HiptmairXu_precond"):
+                self._schur_solver.precond = self._build_precond(dt)
+            self._hx_dt = dt
+
     @profile
     def __call__(self, dt):
         # current FE coeffs
         un = self.variables.u.spline.vector
         bn = self.variables.b.spline.vector
 
         if self.options.algo == "implicit":
-            # Hiptmair-Xu preconditioner of S ~ dt^2/4 (C^T M2gyro C + 4/dt^2 M1), for the current time step: built
-            # at the first time step, rebuilt when dt changes (sigma = 4/dt^2 with the unscaled M2gyro in the middle,
-            # as the vector auxiliary space correction of Hiptmair-Xu assumes a unit-scale curl-curl term)
-            if self.options.precond == "HiptmairXu" and dt != self._hx_dt:
-                self._schur_solver.precond = self._build_precond(dt)
-                self._hx_dt = dt
+            # Hiptmair-Xu preconditioner for the current time step (no-op unless dt changed since the setup)
+            self.prepare_time_step(dt)
 
             # solve for new b coeffs
             byn = self._B.dot(un, out=self._byn)

src/struphy/propagators/shear_alfven_propagator.py

@@ -105,7 +105,7 @@ class Options(OptionsBase):
             the velocity components, ``u_space="Hdiv"`` and no polar splines), with the auxiliary space correction of
             the :math:`\mathbf E \times \mathbf B` flows if ``exb_correction`` (see there). Both need the implicit
             algorithm with ``solver="pcg"``. The Schur operator depends on the time step, so these preconditioners are
-            built at the first time step (and updated if the time step changes).
+            built in the setup (:meth:`prepare_time_step`) and updated if the time step changes.
             The multigrid preconditioner is experimental: on the TAE example (HollowTorus, 8x32x4) it needs more PCG
             iterations than the mass matrix preconditioner and is much slower (the smoother does not treat the large
             kernel of the curl term on 2-forms). The Kronecker preconditioner (with ``exb_correction``) needs about 11
@@ -278,43 +278,51 @@ def f2(t, y1, y2, out: BlockVector = out2):
         self._u_tmp2 = self.variables.u.spline.vector.space.zeros()
         self._b_tmp1 = self.variables.b.spline.vector.space.zeros()
 
+    def prepare_time_step(self, dt: float):
+        """Preconditioner of the Schur operator for the time step ``dt`` (multigrid or Kronecker): built in the setup,
+        updated or rebuilt when dt changes."""
+        if self.options.algo != "implicit":
+            return
+        # the multigrid preconditioner of the Schur operator A - dt^2 B C: updated (re-using the coarse mass and
+        # derivative operators) when dt changes
+        if self.options.precond == "MultiGrid" and dt != self._multigrid_dt:
+            schur = self._schur_solver.A - dt**2 * self._schur_solver.BC
+            if self._schur_solver.precond is None:
+                self._schur_solver.precond = MultiGridPreconditioner(
+                    schur, self.derham, self.domain, self.options.multigrid, mass_ops=self.mass_ops
+                )
+            else:
+                self._schur_solver.precond.update(schur)
+            self._multigrid_dt = dt
+
+        # the Kronecker preconditioner of S = M + dt^2/4 T^T C^T M2 C T (plus the E x B correction), rebuilt when dt
+        # changes
+        if self.options.precond == "StiffnessPreconditioner" and dt != self._schur_precond_dt:
+            mid = (dt**2 / 4) * self.mass_ops.M2
+            pc = preconditioner.StiffnessPreconditioner(
+                self.mass_ops, self._kron_curl_T, sigma=1.0, mid=mid, stab=self._id_M
+            )
+            if self.options.exb_correction:
+                # N^T S N = N^T M N + dt^2/4 (C T N)^T M2 (C T N) on V0
+                pc_phi = preconditioner.StiffnessPreconditioner(
+                    self.mass_ops,
+                    [self._kron_exb, self._kron_curl_T_exb],
+                    sigma=0.0,
+                    mid=[self._id_M, mid],
+                )
+                pc = preconditioner.AuxiliarySpacePreconditioner(pc, [(self._exb_N, self._exb_NT, pc_phi)])
+            self._schur_solver.precond = pc
+            self._schur_precond_dt = dt
+
     @profile
     def __call__(self, dt):
         # current FE coeffs
         un = self.variables.u.spline.vector
         bn = self.variables.b.spline.vector
 
         if self.options.algo == "implicit":
-            # the multigrid preconditioner of the Schur operator A - dt^2 B C, for the current time step: built at the
-            # first time step, updated (re-using the coarse mass and derivative operators) when dt changes
-            if self.options.precond == "MultiGrid" and dt != self._multigrid_dt:
-                schur = self._schur_solver.A - dt**2 * self._schur_solver.BC
-                if self._schur_solver.precond is None:
-                    self._schur_solver.precond = MultiGridPreconditioner(
-                        schur, self.derham, self.domain, self.options.multigrid, mass_ops=self.mass_ops
-                    )
-                else:
-                    self._schur_solver.precond.update(schur)
-                self._multigrid_dt = dt
-
-            # the Kronecker preconditioner of S = M + dt^2/4 T^T C^T M2 C T (plus the E x B correction), rebuilt when dt
-            # changes
-            if self.options.precond == "StiffnessPreconditioner" and dt != self._schur_precond_dt:
-                mid = (dt**2 / 4) * self.mass_ops.M2
-                pc = preconditioner.StiffnessPreconditioner(
-                    self.mass_ops, self._kron_curl_T, sigma=1.0, mid=mid, stab=self._id_M
-                )
-                if self.options.exb_correction:
-                    # N^T S N = N^T M N + dt^2/4 (C T N)^T M2 (C T N) on V0
-                    pc_phi = preconditioner.StiffnessPreconditioner(
-                        self.mass_ops,
-                        [self._kron_exb, self._kron_curl_T_exb],
-                        sigma=0.0,
-                        mid=[self._id_M, mid],
-                    )
-                    pc = preconditioner.AuxiliarySpacePreconditioner(pc, [(self._exb_N, self._exb_NT, pc_phi)])
-                self._schur_solver.precond = pc
-                self._schur_precond_dt = dt
+            # preconditioner for the current time step (no-op unless dt changed since the setup)
+            self.prepare_time_step(dt)
 
             # solve for new u coeffs
             byn = self._B.dot(bn, out=self._byn)

src/struphy/propagators/tests/test_shear_alfven_curl_curl.py

@@ -101,6 +101,22 @@ def test_energy_dissipation(tmp_path):
     assert en[-1] > 0.9 * en[0], en
 
 
+def test_prepare_time_step(tmp_path):
+    """prepare_time_step (called in the setup of Simulation.run) builds the Hiptmair-Xu preconditioner, which a step
+    with the same dt reuses and one with another dt rebuilds."""
+    options = ShearAlfvenCurlCurlPropagator.Options(solver_params=SolverParameters(tol=1e-13, maxiter=1000))
+    model = allocate_model(tmp_path, "prepare", options)
+    shear_alf = model.propagators.shear_alf
+    assert shear_alf._hx_dt is None
+    shear_alf.prepare_time_step(0.5)
+    assert shear_alf._hx_dt == 0.5
+    pc = shear_alf._schur_solver.precond
+    shear_alf(0.5)
+    assert shear_alf._schur_solver.precond is pc
+    shear_alf(0.2)
+    assert shear_alf._hx_dt == 0.2 and shear_alf._schur_solver.precond is not pc
+
+
 def test_formulations_agree(tmp_path):
     """For a smooth mode, the steps with and without the projection in the curl-curl term differ by the projection
     error only."""

src/struphy/simulation/sim.py

@@ -727,6 +727,10 @@ def run(self, one_time_step: bool = False, profiling_activated: bool | None = No
 
                 # CUDA kernels (CuPy backend only): compile them now rather than in the first time step
                 self.compile_cuda_kernels()
+
+                # time step dependent data of the propagators (e.g. preconditioners), rather than in the first step
+                with ProfileManager.profile_region("setup: prepare time step", functions=[self._prepare_time_step]):
+                    self._prepare_time_step()
                 with ProfileManager.profile_region("setup: run metadata", functions=[self._write_run_metadata]):
                     self._write_run_metadata(
                         one_time_step=one_time_step,
@@ -1359,6 +1363,16 @@ def _allocate_variables(self):
 
         #             self._pointer[key] = val["obj"].vector
 
+    def _prepare_time_step(self):
+        """Call :meth:`~struphy.propagators.base.Propagator.prepare_time_step` of each propagator with the time step
+        it receives from :meth:`~struphy.models.base.StruphyModel.integrate` (``dt / 2`` for all but the last
+        propagator with Strang splitting)."""
+        dt = self.time_opts.dt
+        props = self.model.prop_list
+        for n, prop in enumerate(props):
+            half = self.time_opts.split_algo == "Strang" and n < len(props) - 1
+            prop.prepare_time_step(dt / 2 if half else dt)
+
     @profile
     def _allocate_propagators(self):
         """Allocate propagators and bind shared FEEC/domain operators.