From a59388aec673014c4573e9d92d44111e5f57d60e Mon Sep 17 00:00:00 2001 From: Nicolo' Minichelli Date: Fri, 23 Oct 2026 17:05:43 +0200 Subject: [PATCH 1/4] Variational MHD: split sound step for the q model, Kronecker Schur preconditioner for H1vec velocities - VariationalQBEvolve(coupling="bq"|"b"|"q"): advance only the coupling of the velocity to B or to q, the other variable and its auxiliary variable fixed (2x2 Newton system). VariationalPerturbativeMHD_with_q(split_sound=True) uses two substeps, variat_qb (b) and variat_sound (q), each energy conserving; the q TAE example uses it (sound solve 8-9 PCG iterations, 2.5 s instead of 4.2 s per step). - Options schur_precond and exb_correction of VariationalPBEvolve and VariationalQBEvolve: ShearAlfvenSchurPreconditioner (now also for u_space="H1vec", with the E x B map into H1vec) for the shear-Alfven Schur operator of the linear p model and of the Alfven substep of the split q model, built in prepare_time_step; fallback to the mass-matrix preconditioner elsewhere. SchurSolverFull/Full3 get a precond setter. - StiffnessPreconditioner (general operators): separable weight shapes per diagonal block (block_shapes=True) instead of one common shape. Exact for separable weights with different profiles per component (curvilinear metric); no change in Cartesian geometry. LinearMHD shear-Alfven PCG iterations on the low-mode TAE: 23-25 -> 14-17. - Tests: energy conservation of the split step and first-order splitting error (verification), same solution with both Schur preconditioners and fallback (test_variational_schur_precond.py). Co-Authored-By: Claude Opus 5.5 --- .../with_q/lowmode_tae/params_lowmode_tae.py | 10 +- src/struphy/feec/preconditioner.py | 108 +++++++-- src/struphy/feec/variational_utilities.py | 57 +++++ src/struphy/linear_algebra/schur_solver.py | 28 +++ ...verif_VariationalPerturbativeMHD_with_q.py | 37 ++- .../verification/variational_mhd_utils.py | 34 ++- .../variational_perturbative_mhd_with_q.py | 40 ++- .../tests/test_variational_schur_precond.py | 64 +++++ .../propagators/variational_pb_evolve.py | 32 ++- .../propagators/variational_qb_evolve.py | 227 ++++++++++-------- 10 files changed, 504 insertions(+), 133 deletions(-) create mode 100644 src/struphy/propagators/tests/test_variational_schur_precond.py diff --git a/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py b/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py index c6e695dc4..bdfc1e62b 100644 --- a/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py +++ b/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py @@ -49,7 +49,9 @@ The variational propagators run with model="linear_q" and their default (Newton/Picard) solver options, except implicit_weights=False: the weights of the transport operators at t^n, i.e. first order in time like the p model (no midpoint option) and LinearMHD (Crank-Nicolson substeps, Lie-Trotter -splitting). The step stays energy conserving and needs a single Picard iteration. +splitting). With split_sound=True the coupling of the velocity to B (shear-Alfven/fast waves, variat_qb) and to q +(sound waves, variat_sound) are advanced in two implicit substeps, like the shear-Alfven and magnetosonic steps of +LinearMHD: each substep is energy conserving and needs a single Picard iteration, with a smaller linear system. Expected result. With n = 1 and m = 1 the gap is at q(r_c) = 3/2, i.e. r_c = 0.5, where n = 1 and |B| = 1.0 (flux-surface average), so v_A(r_c) = 1.0 v_unit and @@ -135,8 +137,9 @@ base_units = BaseUnits() # Model instance: linear and ideal (no viscosity, no resistivity), as LinearMHD +# split_sound: shear-Alfven/fast and sound waves in two substeps, as in LinearMHD model = VariationalPerturbativeMHD_with_q( - base_units=base_units, linear=True, with_viscosity=False, with_resistivity=False + base_units=base_units, linear=True, with_viscosity=False, with_resistivity=False, split_sound=True ) # List all variables and decide whether to save their data @@ -201,6 +204,9 @@ model.propagators.variat_qb.options = model.propagators.variat_qb.Options( model="linear_q", gamma=gamma, implicit_weights=False ) +model.propagators.variat_sound.options = model.propagators.variat_sound.Options( + model="linear_q", gamma=gamma, implicit_weights=False +) # ------------------ # Initial conditions diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index c3c9da0c7..5fd7078ff 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -679,8 +679,8 @@ class StiffnessPreconditioner(KroneckerPreconditioner): A = \mathbb L^T \mathbb M_{\textrm{mid}} \mathbb L + \sigma \, \mathbb S \,, with ``mid`` (default: the mass operator of the codomain of :math:`\mathbb L`) and ``stab`` (default: of its - domain). The weights of ``mid`` and ``stab`` are approximated by separable shapes (separately), times the - mean of each block; the off-diagonal blocks of ``mid`` are neglected. Each diagonal block of :math:`A` is then a + domain). The weights of ``mid`` and ``stab`` are approximated by separable shapes (separately, and per diagonal + block with ``block_shapes``), times the mean of each block; the off-diagonal blocks of ``mid`` are neglected. Each diagonal block of :math:`A` is then a sum of Kronecker products :math:`\sum_t c_t F_{t,1} \otimes F_{t,2} \otimes F_{t,3} + \sigma_c M_1 \otimes M_2 \otimes M_3` (block Jacobi in the components of the domain). It is inverted in the basis :math:`U_1 \otimes U_2 \otimes U_3` of the generalized eigenproblems :math:`\hat S_d U_d = M_d U_d \Lambda_d` of @@ -740,6 +740,13 @@ class StiffnessPreconditioner(KroneckerPreconditioner): Gradient only: the directions averaged in the stabilization :math:`\mathbb S (\mathbb I - \mathbb A)`, e.g. ``"eta3"``, ``"eta2 eta3"`` or ``"eta1 eta2 eta3"`` (global mean); None (default) for no averaging. + block_shapes : bool + General operators only: approximate the weights of each diagonal block of ``mid`` and ``stab`` by their own + separable shape (default), instead of one shape common to all blocks. The blocks are inverted independently + (block Jacobi), so this costs nothing and is exact for separable weights with different profiles per + component, as the metric coefficients of curvilinear maps (e.g. :math:`G_{22} \sqrt g \propto r^3` for the + poloidal component of a vector field in a torus); in a Cartesian geometry all blocks have the same shape. + References ---------- R. E. Lynch, J. R. Rice and D. H. Thomas, Direct solution of partial difference equations by tensor product @@ -769,8 +776,10 @@ def __init__( mid: str | LinearOperator | None = None, stab: str | LinearOperator | None = None, stab_average: str | None = None, + block_shapes: bool = True, ): derham = mass_ops.derham + self._block_shapes = block_shapes ops = list(derivative) if isinstance(derivative, (list, tuple)) else [derivative] self._general = all(isinstance(op, KroneckerBlockOperator) for op in ops) if self._general: @@ -1059,12 +1068,14 @@ def _build_general_approximation(self, bc): # separable weights of stab (shapes per direction, means per block); unit weights keep the coefficients of the # terms (e.g. of a scaled mass operator) + # (per diagonal block with block_shapes, else one common shape) + separable = _separable_weights_per_block if self._block_shapes else _common_shapes_per_block if self._stab_is_identity: - stab_shapes, dom_means = [_ones_1d] * 3, [1.0] * n_dom + stab_shapes, dom_means = [[_ones_1d] * 3] * n_dom, [1.0] * n_dom elif self._weights == "average": - stab_shapes, dom_means = _separable_weight(_block_weights(self._stab_terms, n_dom), derham) + stab_shapes, dom_means = separable(_block_weights(self._stab_terms, n_dom), derham) else: - stab_shapes, dom_means = [_ones_1d] * 3, _unit_means(self._stab_terms, n_dom) + stab_shapes, dom_means = [[_ones_1d] * 3] * n_dom, _unit_means(self._stab_terms, n_dom) # per operator L_i: its blocks, the means of the blocks of mid_i and their 1d mass matrices (with the essential # boundary conditions of mid_i as zero rows and columns) @@ -1074,9 +1085,9 @@ def _build_general_approximation(self, bc): mid_op = mid_terms[0][1] n_mid = len(mid_op.weights) if self._weights == "average": - mid_shapes, mid_means = _separable_weight(_block_weights(mid_terms, n_mid), derham) + mid_shapes, mid_means = separable(_block_weights(mid_terms, n_mid), derham) else: - mid_shapes, mid_means = [_ones_1d] * 3, _unit_means(mid_terms, n_mid) + mid_shapes, mid_means = [[_ones_1d] * 3] * n_mid, _unit_means(mid_terms, n_mid) mid_femspace = mid_op.domain_femspace mid_comps = (mid_femspace,) if isinstance(mid_femspace, TensorFemSpace) else mid_femspace.spaces mid_bc = _boundary_conditions(mid_op, bc is not None) @@ -1085,7 +1096,7 @@ def _build_general_approximation(self, bc): mats = [] for d in range(3): basis = comp.spaces[d].basis - M_d, _ = _dense_mass_1d(mid_op, basis, d, mid_shapes[d]) + M_d, _ = _dense_mass_1d(mid_op, basis, d, mid_shapes[r][d]) if mid_bc is not None and _has_essential_bc(mid_op, basis, r, d): M_d = _zero_bc_dense(M_d, mid_bc[d]) mats.append(M_d) @@ -1101,7 +1112,7 @@ def _build_general_approximation(self, bc): mass, local_mass, has_bc = [], [], [] for d in range(3): basis = comp.spaces[d].basis - M_d, domain_decomposition = _dense_mass_1d(self._mass_operator, basis, d, stab_shapes[d]) + M_d, domain_decomposition = _dense_mass_1d(self._mass_operator, basis, d, stab_shapes[c][d]) if self._stab_is_identity: M_d = np.eye(M_d.shape[0]) essential = bc is not None and _has_essential_bc(self._mass_operator, basis, c, d) @@ -1876,12 +1887,15 @@ class ShearAlfvenSchurPreconditioner: S = \mathbb M^{2,n} + \frac{\Delta t^2}{4} \mathbb T^\top \mathbb C^\top \mathbb M^2 \mathbb C \mathbb T\,, - for the velocity in :math:`H(\textnormal{div})`: the :class:`StiffnessPreconditioner` with the operator + for the velocity in :math:`H(\textnormal{div})` (or in :math:`(H^1)^3`, ``u_space="H1vec"``, with + :math:`\mathbb M^{v}[\rho_0]` instead of :math:`\mathbb M^{2,n}`, as in the variational models): the + :class:`StiffnessPreconditioner` with the operator :math:`\mathbb C \tilde{\mathbb T}` (:math:`\tilde{\mathbb T}` is the Kronecker approximation of the basis projection operator, :func:`~struphy.feec.kronecker_block.basis_projection_operator`), block Jacobi in the velocity components, plus the auxiliary space correction of the :math:`\mathbf E \times \mathbf B` flows :math:`\mathbf B_0 \times \nabla \phi` if ``exb_correction`` - (:class:`AuxiliarySpacePreconditioner` with :math:`\mathbb N = \hat \Pi^2[\mathbf B_0 \times \,\cdot\,] \, \mathbb G`). + (:class:`AuxiliarySpacePreconditioner` with :math:`\mathbb N = \hat \Pi^2[\mathbf B_0 \times \,\cdot\,] \, \mathbb G`, + or :math:`\hat \Pi^v[\mathbf B_0 \times \,\cdot\,] \, \mathbb G` for ``u_space="H1vec"``). Shared by :class:`~struphy.propagators.shear_alfven_propagator.ShearAlfvenPropagator` and :class:`~struphy.propagators.shear_alfven_current_coupling_5d.ShearAlfvenCurrentCoupling5D`. @@ -1907,15 +1921,19 @@ class ShearAlfvenSchurPreconditioner: exb_correction : bool Whether to add the auxiliary space correction of the :math:`\mathbf E \times \mathbf B` flows. + + u_space : str + Space of the velocity, ``"Hdiv"`` (2-form) or ``"H1vec"`` (vector field), the domain of ``T``. """ - def __init__(self, derham, mass_ops, basis_ops, T, id_M: str, exb_correction: bool = True): + def __init__(self, derham, mass_ops, basis_ops, T, id_M: str, exb_correction: bool = True, u_space: str = "Hdiv"): from struphy.feec.kronecker_block import basis_projection_operator, derivative_operator from struphy.feec.utilities import LocalRotationMatrix assert self.supported(derham), ( "The shear Alfven Kronecker preconditioner needs tensor-product splines and the global commuting projectors." ) + assert u_space in ("Hdiv", "H1vec"), f'u_space must be "Hdiv" or "H1vec", got {u_space!r}.' self._mass_ops = mass_ops self._id_M = id_M self._exb_correction = exb_correction @@ -1923,13 +1941,24 @@ def __init__(self, derham, mass_ops, basis_ops, T, id_M: str, exb_correction: bo # Kronecker approximation of the curl of the basis projection self._kron_curl_T = derivative_operator(derham, "curl") @ basis_projection_operator(T, derham) if exb_correction: - # E x B flows B0 x grad(phi): 2-form components eps_ijk b1_j d_k phi (b1: covariant components of B0) + # E x B flows B0 x grad(phi): 2-form components eps_ijk b1_j d_k phi (b1: covariant components of B0), + # or the vector-field (contravariant) components eps_ijk b1_j d_k phi / sqrt(g) eq = basis_ops.weights["eq_mhd"] rot_b1 = LocalRotationMatrix(eq.b1_1, eq.b1_2, eq.b1_3) - fun = [ - [(lambda e1, e2, e3, m=m, n=n: rot_b1(e1, e2, e3)[:, :, :, m, n]) for n in range(3)] for m in range(3) - ] - Pi_exb = basis_ops.create_basis_op(fun, "Hcurl", "Hdiv", assemble=True, name="Pi_ExB") + if u_space == "Hdiv": + fun = [ + [(lambda e1, e2, e3, m=m, n=n: rot_b1(e1, e2, e3)[:, :, :, m, n]) for n in range(3)] + for m in range(3) + ] + else: + fun = [ + [ + (lambda e1, e2, e3, m=m, n=n: rot_b1(e1, e2, e3)[:, :, :, m, n] / basis_ops.sqrt_g(e1, e2, e3)) + for n in range(3) + ] + for m in range(3) + ] + Pi_exb = basis_ops.create_basis_op(fun, "Hcurl", u_space, assemble=True, name="Pi_ExB") grad = derham.grad self._exb_N = Pi_exb @ grad self._exb_NT = grad.T @ Pi_exb.T @@ -2533,6 +2562,46 @@ def _separable_weight(weights: list, derham) -> tuple[list[xp.ndarray], list[flo Collective on ``derham.comm`` for array weights. """ + profiles, means = _weight_profiles(weights, derham) + nonzero = [(prof, m) for prof, m in zip(profiles, means) if m > 0.0] + shapes = [sum(prof[d] / m for prof, m in nonzero) / len(nonzero) for d in range(3)] + return shapes, means + + +def _separable_weights_per_block(weights: list, derham) -> tuple[list[list[xp.ndarray]], list[float]]: + r""" + Separable shape of each 3d weight, :math:`w_i \approx \bar w_i \prod_d \phi_{i,d}(\eta_d)`, as + :func:`_separable_weight` but with the shape of each weight itself instead of the common one (for the block Jacobi + approximation of general operators, whose diagonal blocks are inverted independently). A weight with vanishing + mean gets the common shape. + + Returns + ------- + shapes : list[list[xp.ndarray]] + Per weight, :math:`\phi_{i,d}` at the global 1d quadrature points (mean 1). + + means : list[float] + The means :math:`\bar w_i`, as in :func:`_separable_weight`. + + Collective on ``derham.comm`` for array weights. + """ + profiles, means = _weight_profiles(weights, derham) + nonzero = [(prof, m) for prof, m in zip(profiles, means) if m > 0.0] + common = [sum(prof[d] / m for prof, m in nonzero) / len(nonzero) for d in range(3)] + shapes = [[prof[d] / m for d in range(3)] if m > 0.0 else common for prof, m in zip(profiles, means)] + return shapes, means + + +def _common_shapes_per_block(weights: list, derham) -> tuple[list[list[xp.ndarray]], list[float]]: + """The common shape of :func:`_separable_weight`, repeated for each weight (the format of + :func:`_separable_weights_per_block`).""" + shapes, means = _separable_weight(weights, derham) + return [shapes] * len(weights), means + + +def _weight_profiles(weights: list, derham) -> tuple[list[list[xp.ndarray]], list[float]]: + """Per weight (see :func:`_separable_weight`): its 1d profiles (the mean over the other directions) at the global + 1d quadrature points, and its mean over the logical cube (0 if it vanishes relative to the largest one).""" grids = [derham.H1_1d_serial[d].get_assembly_grids(derham.nquads[d])[0] for d in range(3)] pts = [xp.asarray(np.ravel(g.points)) for g in grids] wts = [xp.asarray(np.ravel(g.weights)) for g in grids] @@ -2554,10 +2623,7 @@ def _separable_weight(weights: list, derham) -> tuple[list[xp.ndarray], list[flo assert largest > 0.0, "The weights must not all vanish." means = [0.0 if abs(m) <= 1e-12 * largest else m for m in means] assert all(m >= 0.0 for m in means), f"The weights must have non-negative means, got {means}." - - nonzero = [(prof, m) for prof, m in zip(profiles, means) if m > 0.0] - shapes = [sum(prof[d] / m for prof, m in nonzero) / len(nonzero) for d in range(3)] - return shapes, means + return profiles, means def _dense_mass_1d( diff --git a/src/struphy/feec/variational_utilities.py b/src/struphy/feec/variational_utilities.py index d2dc4b585..26b826fd1 100644 --- a/src/struphy/feec/variational_utilities.py +++ b/src/struphy/feec/variational_utilities.py @@ -1877,3 +1877,60 @@ def inv_norm2(inv: LinearOperator, v: Vector) -> float: err = inv.dot_inner(v, v) v *= norm return norm**2 * err + + +def alfven_schur_preconditioner(propagator, unmet: list[str]): + r"""The Kronecker preconditioner of the shear-Alfven Schur operator of a linear variational propagator, + + .. math:: + + \mathcal S_A = \mathbb M^v[\rho_0] + \frac{\Delta t^2}{4}\, \mathbb T_0^\top \mathbb M^2\, \mathbb T_0 \,, + \qquad \mathbb T_0 = \mathbb C\, \hat\Pi^1[\hat{\mathbf B}_0 \times \vec\Lambda^v] \,, + + as :class:`~struphy.feec.preconditioner.ShearAlfvenSchurPreconditioner` with ``u_space="H1vec"``: the + :class:`~struphy.feec.preconditioner.StiffnessPreconditioner` of the Kronecker approximation of + :math:`\mathbb C\, \hat\Pi^1[\mathbf B_0 \times \cdot\,]` (``basis_ops.Tv``, analytic equilibrium) with the + stabilization ``Mvn`` (the velocity mass matrix with the analytic density), plus the auxiliary space correction of + the :math:`\mathbf E \times \mathbf B` flows if ``options.exb_correction``. Its :meth:`build` gives the + preconditioner for a time step. + + Parameters + ---------- + propagator : Propagator + The propagator, with the options ``schur_precond``, ``solver`` and ``exb_correction``. + unmet : list[str] + Requirements of the propagator that do not hold (e.g. ``'model="linear"'``). + + Returns + ------- + ShearAlfvenSchurPreconditioner | None + None if ``options.schur_precond`` is not ``"StiffnessPreconditioner"`` or if a requirement does not hold + (then with a warning: the mass-matrix preconditioner is used). + """ + from struphy.feec.preconditioner import ShearAlfvenSchurPreconditioner + + options = propagator.options + if options.schur_precond != "StiffnessPreconditioner": + return None + unmet = list(unmet) + if options.solver != "pcg": + unmet += ['solver="pcg"'] + if propagator.derham.polar_splines: + unmet += ["tensor-product splines"] + if propagator.derham.with_local_projectors: + unmet += ["the global commuting projectors"] + if unmet: + logger.warning( + f'{propagator.__class__.__name__}: schur_precond="StiffnessPreconditioner" needs {", ".join(unmet)}; ' + "falling back to the mass-matrix preconditioner." + ) + return None + return ShearAlfvenSchurPreconditioner( + propagator.derham, + propagator.mass_ops, + propagator.basis_ops, + propagator.basis_ops.Tv, + "Mvn", + exb_correction=options.exb_correction, + u_space="H1vec", + ) diff --git a/src/struphy/linear_algebra/schur_solver.py b/src/struphy/linear_algebra/schur_solver.py index a46b98fe8..6c1d890a5 100644 --- a/src/struphy/linear_algebra/schur_solver.py +++ b/src/struphy/linear_algebra/schur_solver.py @@ -228,11 +228,25 @@ def __init__(self, M, solver_name, **solver_params): self._S = self._A - self._B @ self._C + self._solver_params = solver_params self._solver = inverse(self._S, solver_name, **solver_params) # right-hand side vector (avoids temporary memory allocation!) self._rhs = self._A.codomain.zeros() + @property + def precond(self): + """Preconditioner of the Schur system (None without preconditioner).""" + return self._solver_params.get("pc") + + @precond.setter + def precond(self, pc): + """Replace the preconditioner, e.g. by one of the Schur system for the current dt.""" + self._solver_params.pop("pc", None) + if pc is not None: + self._solver_params["pc"] = pc + self._solver = inverse(self._S, self._solver_name, **self._solver_params) + @profile @ProfileManager.profile("solve: SchurSolverFull") def dot(self, v, out=None): @@ -339,12 +353,26 @@ def __init__(self, M, solver_name, **solver_params): self._S = self._A - self._B @ self._C - self._D @ self._E + self._solver_params = solver_params self._solver = inverse(self._S, solver_name, **solver_params) # right-hand side vector (avoids temporary memory allocation!) self._rhs = self._A.codomain.zeros() self._rhs2 = self._A.codomain.zeros() + @property + def precond(self): + """Preconditioner of the Schur system (None without preconditioner).""" + return self._solver_params.get("pc") + + @precond.setter + def precond(self, pc): + """Replace the preconditioner, e.g. by one of the Schur system for the current dt.""" + self._solver_params.pop("pc", None) + if pc is not None: + self._solver_params["pc"] = pc + self._solver = inverse(self._S, self._solver_name, **self._solver_params) + @profile @ProfileManager.profile("solve: SchurSolverFull3") def dot(self, v, out=None): diff --git a/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py b/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py index bce1cf4b2..e37bcd345 100644 --- a/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py +++ b/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py @@ -22,7 +22,9 @@ 7. it is the linear limit of the perturbative model, also around the sheared slab; 8. it gives the analytic frequencies of the shear Alfvén, the fast magnetosonic and the sound wave in a homogeneous slab, with second-order convergence in the time step; -9. the p-based linear model agrees up to its first-order time discretization error. +9. the p-based linear model agrees up to its first-order time discretization error; +10. with ``split_sound=True`` (the shear-Alfven/fast and the sound waves in two substeps) the energy is conserved; +11. and the split model agrees with the unsplit one up to a splitting error of first order in the time step. """ import logging @@ -54,6 +56,7 @@ logger = logging.getLogger("struphy") LINEAR = linear(VariationalPerturbativeMHD_with_q) +LINEAR_SPLIT = linear(VariationalPerturbativeMHD_with_q, split_sound=True) def test_energy_conservation(tmp_path): @@ -292,3 +295,35 @@ def test_linear_same_as_p_model(tmp_path): for i, name in enumerate(("u", "B")): assert rel_diff[0.05][i] < 0.6, (name, rel_diff) assert rel_diff[0.025][i] < 0.6 * rel_diff[0.05][i], (name, rel_diff) + + +def test_linear_split_energy_conservation(tmp_path): + """With split_sound=True, each substep (u, B) and (u, q) conserves H, so the split step does too (sheared + slab, where the terms with grad q0 and curl B0 are active).""" + model = allocate_ideal( + tmp_path, "split_energy", LINEAR_SPLIT, sheared_slab(), dt=0.05, perturb=perturb_velocity(0.05) + ) + energies = energy_history(model, dt=0.05, n_steps=10) + + assert abs(model.scalars.host_value("en_thermo_2")[0]) > 1e-5 + assert abs(model.scalars.host_value("en_mag_2")[0]) > 1e-5 + drift = max(abs(en - energies[0]) for en in energies) + logger.info(f"split_sound: relative energy drift {drift / abs(energies[0]):.2e}") + assert drift < 2e-8 * abs(energies[0]), energies + + +def test_linear_split_same_as_unsplit(tmp_path): + """The split step (Lie-Trotter splitting of the (u, B) and (u, q) substeps) agrees with the unsplit step up to a + splitting error of first order in the time step (sheared slab, where the two couplings do not commute).""" + setup = sheared_slab() + t_end = 0.5 + rel_diff = {} + for dt in (0.05, 0.025): + n_steps = int(round(t_end / dt)) + split = _final_state(tmp_path, LINEAR_SPLIT, setup, 0.05, dt, n_steps) + unsplit = _final_state(tmp_path, LINEAR, setup, 0.05, dt, n_steps) + rel_diff[dt] = [norm(split[i] - unsplit[i]) / norm(unsplit[i]) for i in (1, 2, 3)] + logger.info(f"relative difference split / unsplit (u, q, B) per dt: {rel_diff}") + for i, name in enumerate(("u", "q", "B")): + assert rel_diff[0.05][i] < 0.3, (name, rel_diff) + assert rel_diff[0.025][i] < 0.6 * rel_diff[0.05][i], (name, rel_diff) diff --git a/src/struphy/models/tests/verification/variational_mhd_utils.py b/src/struphy/models/tests/verification/variational_mhd_utils.py index d1d4d1936..6b8223184 100644 --- a/src/struphy/models/tests/verification/variational_mhd_utils.py +++ b/src/struphy/models/tests/verification/variational_mhd_utils.py @@ -30,10 +30,11 @@ NONLIN = NonlinearSolverParameters() -def linear(model_class): - """The linear version of a perturbative model class (a constructor with ``linear=True``).""" - constructor = functools.partial(model_class, linear=True) - constructor.__name__ = f"{model_class.__name__}_linear" +def linear(model_class, **kwargs): + """The linear version of a perturbative model class (a constructor with ``linear=True`` and ``kwargs``).""" + constructor = functools.partial(model_class, linear=True, **kwargs) + suffix = "".join(f"_{key}" for key, value in kwargs.items() if value) + constructor.__name__ = f"{model_class.__name__}_linear{suffix}" return constructor @@ -59,11 +60,21 @@ def homogeneous_slab(B0=(0.0, 1.0, 1.0), Lz=10.0, num_elements=16): return equil, domain, (1, 1, num_elements), (1, 1, 3) -def allocate_ideal(tmp_path, name, model_class, setup, dt, perturb=None, implicit_weights=None): +def allocate_ideal( + tmp_path, + name, + model_class, + setup, + dt, + perturb=None, + implicit_weights=None, + schur_precond="MassMatrixPreconditioner", +): """Allocate an ideal model of `model_class` for the `setup` (equil, domain, num_elements, degree). `perturb(model)` adds the perturbations of the perturbative variables. The full-f model gets the equilibrium - as background in addition. `implicit_weights` is passed to VariationalQBEvolve (None: its default). Models must be + as background in addition. `implicit_weights` is passed to VariationalQBEvolve (None: its default), + `schur_precond` to VariationalPBEvolve and VariationalQBEvolve. Models must be allocated and run one after the other, since they share the class attributes of Propagator (derham, ...).""" equil, domain, num_elements, degree = setup model = model_class(with_viscosity=False, with_resistivity=False) @@ -78,10 +89,19 @@ def allocate_ideal(tmp_path, name, model_class, setup, dt, perturb=None, implici ) if hasattr(model.propagators, "variat_pb"): model.propagators.variat_pb.options = model.propagators.variat_pb.Options( - model=variant, solver_params=SOLVER, nonlin_solver=NONLIN + model=variant, solver_params=SOLVER, nonlin_solver=NONLIN, schur_precond=schur_precond ) else: model.propagators.variat_qb.options = model.propagators.variat_qb.Options( + model=variant, + solver_params=SOLVER, + nonlin_solver=NONLIN, + implicit_weights=implicit_weights, + schur_precond=schur_precond, + ) + # the sound substep of split_sound=True + if hasattr(model.propagators, "variat_sound"): + model.propagators.variat_sound.options = model.propagators.variat_sound.Options( model=variant, solver_params=SOLVER, nonlin_solver=NONLIN, implicit_weights=implicit_weights ) diff --git a/src/struphy/models/variational_perturbative_mhd_with_q.py b/src/struphy/models/variational_perturbative_mhd_with_q.py index 701ddd70f..afe5bfb07 100644 --- a/src/struphy/models/variational_perturbative_mhd_with_q.py +++ b/src/struphy/models/variational_perturbative_mhd_with_q.py @@ -38,6 +38,11 @@ class VariationalPerturbativeMHD_with_q(StruphyModel): Whether to include viscous dissipation (default: True) with_resistivity: bool Whether to include resistive dissipation (default: True) + split_sound: bool + Whether to advance the coupling of the velocity to the magnetic field and to q in two substeps + (``variat_qb`` and ``variat_sound``, both energy conserving), instead of together in ``variat_qb``. The + Newton systems are then smaller (shear-Alfven/fast and sound waves separately), at the cost of a splitting + error (default: False) """ @classmethod @@ -79,11 +84,17 @@ def __init__( linear: bool = False, with_viscosity: bool = True, with_resistivity: bool = True, + split_sound: bool = False, ): self.variat_dens = VariationalDensityEvolve() if not linear: self.variat_mom = VariationalMomentumAdvection() - self.variat_qb = VariationalQBEvolve(div_u=div_u, u2=u2, qt3=qt3, bt2=bt2) + if split_sound: + # the shear-Alfven/fast waves (u, b) and the sound waves (u, q) in two substeps + self.variat_qb = VariationalQBEvolve(div_u=div_u, u2=u2, qt3=qt3, bt2=bt2, coupling="b") + self.variat_sound = VariationalQBEvolve(div_u=div_u, u2=u2, qt3=qt3, bt2=bt2, coupling="q") + else: + self.variat_qb = VariationalQBEvolve(div_u=div_u, u2=u2, qt3=qt3, bt2=bt2) if with_viscosity: self.variat_viscous = VariationalViscosity(rho=rho, pt3=qt3) if with_resistivity: @@ -98,6 +109,7 @@ def __init__( linear: bool = False, with_viscosity: bool = True, with_resistivity: bool = True, + split_sound: bool = False, ): # 0. store input parameters @@ -121,6 +133,7 @@ def __init__( linear=linear, with_viscosity=with_viscosity, with_resistivity=with_resistivity, + split_sound=split_sound, ) # 4. assign variables to propagators @@ -131,6 +144,10 @@ def __init__( self.propagators.variat_qb.variables.u = self.mhd.velocity self.propagators.variat_qb.variables.q = self.mhd.sqrt_p self.propagators.variat_qb.variables.b = self.em_fields.b_field + if split_sound: + self.propagators.variat_sound.variables.u = self.mhd.velocity + self.propagators.variat_sound.variables.q = self.mhd.sqrt_p + self.propagators.variat_sound.variables.b = self.em_fields.b_field if with_viscosity: self.propagators.variat_viscous.variables.s = self.mhd.sqrt_p self.propagators.variat_viscous.variables.u = self.mhd.velocity @@ -165,7 +182,8 @@ def velocity_scale(self): def propagator_variants(self) -> dict: """The values of the option ``model`` of the variational propagators that match ``linear``.""" variant = "linear_q" if self.params["linear"] else "perturbative_q" - return {name: variant for name in ("variat_dens", "variat_qb", "variat_viscous", "variat_resist")} + names = ("variat_dens", "variat_qb", "variat_sound", "variat_viscous", "variat_resist") + return {name: variant for name in names} def post_allocate(self): # the propagator options must match the flag `linear` @@ -176,6 +194,13 @@ def post_allocate(self): f"{self.__class__.__name__}(linear={self.params['linear']}) needs " f"{name}.Options(model='{variant}'), got model='{prop.options.model}'." ) + # both substeps of the split step and the energy diagnostics must use the same adiabatic index + sound = getattr(self.propagators, "variat_sound", None) + if sound is not None and sound.options.gamma != self.propagators.variat_qb.options.gamma: + raise ValueError( + f"{self.__class__.__name__}(split_sound=True) needs the same gamma in variat_qb and variat_sound, " + f"got {self.propagators.variat_qb.options.gamma} and {sound.options.gamma}." + ) projV3 = L2Projector("L2", Propagator.mass_ops) @@ -221,6 +246,10 @@ def generate_default_parameter_file(self, path=None, prompt=True): new_file += [ f"model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='{variant}')\n", ] + elif "variat_sound.Options" in line: + new_file += [ + f"model.propagators.variat_sound.options = model.propagators.variat_sound.Options(model='{variant}')\n", + ] elif "variat_viscous.Options" in line: new_file += [ f"model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='{variant}')\n", @@ -345,7 +374,7 @@ def doc_scalar_quantities(cls): @classmethod def doc_discretization(cls): - """Time integration is performed by the following propagators (in sequence): + r"""Time integration is performed by the following propagators (in sequence): 1. :class:`~struphy.propagators.variational_density_evolve.VariationalDensityEvolve` 2. :class:`~struphy.propagators.variational_momentum_advection.VariationalMomentumAdvection` (if :attr:`linear` is False) @@ -354,7 +383,10 @@ def doc_discretization(cls): 5. :class:`~struphy.propagators.variational_resistivity.VariationalResistivity` (if :attr:`with_resistivity` is True) The option ``model`` of the variational propagators is ``"perturbative_q"``, or ``"linear_q"`` if - :attr:`linear` is True. + :attr:`linear` is True. With :attr:`split_sound` the step 3 is done in two substeps of + :class:`~struphy.propagators.variational_qb_evolve.VariationalQBEvolve`: ``variat_qb`` with + ``coupling="b"`` (shear-Alfven/fast waves, :math:`q` fixed) and ``variat_sound`` with ``coupling="q"`` (sound + waves, :math:`\mathbf B` fixed), both energy conserving. """ doc = rf"""**1. VariationalDensityEvolve:** diff --git a/src/struphy/propagators/tests/test_variational_schur_precond.py b/src/struphy/propagators/tests/test_variational_schur_precond.py new file mode 100644 index 000000000..729d51d88 --- /dev/null +++ b/src/struphy/propagators/tests/test_variational_schur_precond.py @@ -0,0 +1,64 @@ +"""The Kronecker preconditioner of the shear-Alfven Schur operator of the linear variational propagators +(``schur_precond="StiffnessPreconditioner"`` of VariationalPBEvolve and VariationalQBEvolve): it changes the +iteration, not the solution, and it falls back to the mass-matrix preconditioner where it does not apply.""" + +import pytest + +from struphy.models import VariationalPerturbativeMHD, VariationalPerturbativeMHD_with_q +from struphy.models.tests.verification.variational_mhd_utils import ( + allocate_ideal, + homogeneous_slab, + linear, + norm, + perturb_velocity, + state, +) + +LZ = 10.0 +DT = 0.2 +N_STEPS = 3 + +LINEAR_P = linear(VariationalPerturbativeMHD) +LINEAR_Q_SPLIT = linear(VariationalPerturbativeMHD_with_q, split_sound=True) +LINEAR_Q = linear(VariationalPerturbativeMHD_with_q) + + +def _run(tmp_path, model_class, schur_precond): + model = allocate_ideal( + tmp_path, + f"{model_class.__name__}_{schur_precond}", + model_class, + homogeneous_slab(Lz=LZ), + DT, + perturb_velocity(0.1, direction="z", L=LZ), + schur_precond=schur_precond, + ) + for prop in model.prop_list: + prop.prepare_time_step(DT) + for _ in range(N_STEPS): + model.integrate(DT) + return model + + +@pytest.mark.parametrize( + "model_class, name", + [(LINEAR_P, "variat_pb"), (LINEAR_Q_SPLIT, "variat_qb")], + ids=["linear_p", "linear_q_split"], +) +def test_same_solution(tmp_path, model_class, name): + """The stiffness preconditioner is active and gives the same solution as the mass-matrix preconditioner (up to the + solver tolerance).""" + stiff = _run(tmp_path, model_class, "StiffnessPreconditioner") + assert getattr(stiff.propagators, name)._schur_pc is not None + stiff_state = state(stiff) + mass_state = state(_run(tmp_path, model_class, "MassMatrixPreconditioner")) + for i, var in ((1, "u"), (2, "p/q"), (3, "B")): + rel = norm(stiff_state[i] - mass_state[i]) / norm(mass_state[i]) + assert rel < 1e-6, (var, rel) + + +def test_fallback(tmp_path): + """Where the Schur operator is not the shear-Alfven one (the unsplit q model), the mass-matrix preconditioner is + used.""" + model = _run(tmp_path, LINEAR_Q, "StiffnessPreconditioner") + assert model.propagators.variat_qb._schur_pc is None diff --git a/src/struphy/propagators/variational_pb_evolve.py b/src/struphy/propagators/variational_pb_evolve.py index 5e6cf2682..81d1fad27 100644 --- a/src/struphy/propagators/variational_pb_evolve.py +++ b/src/struphy/propagators/variational_pb_evolve.py @@ -17,6 +17,7 @@ Hdiv0_transport_operator, NewtonConvergence, Pressure_transport_operator, + alfven_schur_preconditioner, dot_normalized, inv_norm2, ) @@ -170,7 +171,7 @@ def __init__( @dataclass(repr=False) class Options(OptionsBase): - """Configuration options for :class:`VariationalPBEvolve`. + r"""Configuration options for :class:`VariationalPBEvolve`. Parameters ---------- @@ -189,9 +190,22 @@ class Options(OptionsBase): Nonlinear iteration controls. The iteration stops when the norm of the residual is below ``tol`` times the norm of the state at :math:`t^n` if that is smaller than one (absolute otherwise), or when it stagnates (with a warning). + schur_precond : {"MassMatrixPreconditioner", "StiffnessPreconditioner"}, default="MassMatrixPreconditioner" + Preconditioner of the Schur system of the Newton step: the mass-matrix preconditioner of the velocity + block, or the Kronecker (fast diagonalization) preconditioner of the whole shear-Alfven Schur operator + :math:`\mathbb M^v[\rho_0] + \frac{\Delta t^2}{4} \mathbb T_0^\top \mathbb M^2 \mathbb T_0` + (:func:`~struphy.feec.variational_utilities.alfven_schur_preconditioner`, built from the equilibrium for + each time step in :meth:`prepare_time_step`). The latter needs ``model="linear"``, ``solver="pcg"``, + tensor-product splines and the global commuting projectors; otherwise the mass-matrix preconditioner is + used (with a warning). + exb_correction : bool, default=True + For ``schur_precond="StiffnessPreconditioner"``: add the auxiliary space correction of the + :math:`\mathbf E \times \mathbf B` flows :math:`\mathbf B_0 \times \nabla \phi` (the incompressible shear + Alfven modes with small parallel wave number, which block Jacobi does not reduce). """ # specific literals + OptsSchurPrecond = Literal["MassMatrixPreconditioner", "StiffnessPreconditioner"] OptsModel = Literal["full_p", "linear", "perturbative"] # propagator options model: OptsModel = "full_p" @@ -200,12 +214,15 @@ class Options(OptionsBase): precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None nonlin_solver: NonlinearSolverParameters = None + schur_precond: OptsSchurPrecond = "MassMatrixPreconditioner" + exb_correction: bool = True def __post_init__(self): # checks check_option(self.model, self.OptsModel) check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsMassPrecond) + check_option(self.schur_precond, self.OptsSchurPrecond) # defaults if self.solver_params is None: @@ -263,6 +280,10 @@ def allocate(self): # Projector self._initialize_projectors_and_mass() + # Kronecker preconditioner of the Schur operator, built for the time step in prepare_time_step + self._schur_pc = alfven_schur_preconditioner(self, [] if self._model == "linear" else ['model="linear"']) + self._schur_pc_dt = None + # bunch of temporaries to avoid allocating in the loop u = self.variables.u.spline.vector p = self.variables.p.spline.vector @@ -300,7 +321,16 @@ def allocate(self): if self._linearize: self._extracted_b2 = self.derham.extraction_ops["2"].dot(self.projected_equil.b2) + def prepare_time_step(self, dt: float): + """Kronecker preconditioner of the Schur operator for the time step ``dt`` (with + ``schur_precond="StiffnessPreconditioner"``): built in the setup, rebuilt when dt changes.""" + if self._schur_pc is not None and dt != self._schur_pc_dt: + self._inv_Jacobian.precond = self._schur_pc.build(dt) + self._schur_pc_dt = dt + def __call__(self, dt): + # preconditioner for the current time step (no-op unless dt changed since the setup) + self.prepare_time_step(dt) # WMMnew is re-assembled in VariationalDensityEvolve: update the diagonal scaling of its preconditioner self._Mrho_pc.refresh() diff --git a/src/struphy/propagators/variational_qb_evolve.py b/src/struphy/propagators/variational_qb_evolve.py index 13ff2caec..0e3e22f1d 100644 --- a/src/struphy/propagators/variational_qb_evolve.py +++ b/src/struphy/propagators/variational_qb_evolve.py @@ -16,6 +16,7 @@ Hdiv0_transport_operator, NewtonConvergence, Pressure_transport_operator, + alfven_schur_preconditioner, dot_normalized, inv_norm2, ) @@ -28,6 +29,10 @@ logger = logging.getLogger("struphy") +# couplings to the velocity advanced by VariationalQBEvolve +OptsCoupling = Literal["bq", "b", "q"] + + class VariationalQBEvolve(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf u \in (H^1)^3`, :math:`q \in L^2` and :math:`\mathbf B \in H(\textnormal{div})` such that, @@ -84,6 +89,11 @@ class VariationalQBEvolve(Propagator): For ``"full_q"`` the mass density :math:`\rho` is the variable of :class:`~struphy.propagators.variational_density_evolve.VariationalDensityEvolve`, for ``"perturbative_q"`` it is :math:`\rho_0` plus that variable. + With the constructor argument ``coupling="b"`` (``"q"``) only the terms of the magnetic field (of :math:`q`) are + advanced, with :math:`q` (the magnetic field) and its auxiliary variable fixed; the Newton system is then a + :math:`2\times 2` block system in :math:`(\mathbf u, \mathbf B)` (:math:`(\mathbf u, q)`). Two such substeps split + the shear-Alfven/fast waves from the sound waves, and each substep conserves the energy. + The scheme below is written for ``"full_q"``; for the other models the operators and linear forms are replaced by the sums above. Below, hats denote pullbacks to the logical domain (not the auxiliary variables). The weights of the transport operators are taken at the same time in the momentum, induction and pressure @@ -193,6 +203,7 @@ def __init__( u2: FEECVariable = None, qt3: FEECVariable = None, bt2: FEECVariable = None, + coupling: OptsCoupling = "bq", ): r""" Parameters @@ -209,12 +220,22 @@ def __init__( Auxiliary second-order variable :math:`\hat{\mathbf B}` (``"Hdiv"`` space) of the linear and perturbative models, advanced with :math:`\partial_t \hat{\mathbf B} + \mathcal T_{\mathbf B}(\mathbf u) = 0`. It enters only the energy, not the momentum equation. + coupling : {"bq", "b", "q"}, default="bq" + The couplings to the velocity advanced (implicitly) by this propagator: of the magnetic field and + :math:`q` together (``"bq"``), or only of the magnetic field (``"b"``) or only of :math:`q` (``"q"``), the + other variable (and its auxiliary variable) being kept fixed. Each choice conserves the energy, so two + instances with ``"b"`` and ``"q"`` give an energy-conserving splitting of the shear-Alfven/fast and the + sound waves, with a smaller linear system per substep. """ + check_option(coupling, OptsCoupling) self.variables = self.Variables() self.div_u = div_u self.u2 = u2 self.qt3 = qt3 self.bt2 = bt2 + self._coupling = coupling + self._with_b = coupling in ("bq", "b") + self._with_q = coupling in ("bq", "q") @dataclass(repr=False) class Options(OptionsBase): @@ -246,9 +267,22 @@ class Options(OptionsBase): does not contain the dependence of the operators on their weights, so a step needs several iterations (typically 2 to 10, more for large time steps or amplitudes) instead of one, each with a Schur solve; ``False`` is the cheaper first-order choice. + schur_precond : {"MassMatrixPreconditioner", "StiffnessPreconditioner"}, default="MassMatrixPreconditioner" + Preconditioner of the Schur system of the Newton step: the mass-matrix preconditioner of the velocity + block, or the Kronecker (fast diagonalization) preconditioner of the whole shear-Alfven Schur operator + :math:`\mathbb M^v[\rho_0] + \frac{\Delta t^2}{4} \mathbb T_0^\top \mathbb M^2 \mathbb T_0` + (:func:`~struphy.feec.variational_utilities.alfven_schur_preconditioner`, built from the equilibrium for + each time step in :meth:`prepare_time_step`). The latter needs ``model="linear_q"`` with ``coupling="b"`` (the Alfven substep of a split step), ``solver="pcg"``, + tensor-product splines and the global commuting projectors; otherwise the mass-matrix preconditioner is + used (with a warning). + exb_correction : bool, default=True + For ``schur_precond="StiffnessPreconditioner"``: add the auxiliary space correction of the + :math:`\mathbf E \times \mathbf B` flows :math:`\mathbf B_0 \times \nabla \phi` (the incompressible shear + Alfven modes with small parallel wave number, which block Jacobi does not reduce). """ # specific literals + OptsSchurPrecond = Literal["MassMatrixPreconditioner", "StiffnessPreconditioner"] OptsModel = Literal["full_q", "linear_q", "perturbative_q"] # propagator options model: OptsModel = "full_q" @@ -258,12 +292,15 @@ class Options(OptionsBase): solver_params: SolverParameters = None nonlin_solver: NonlinearSolverParameters = None implicit_weights: bool | None = None + schur_precond: OptsSchurPrecond = "MassMatrixPreconditioner" + exb_correction: bool = True def __post_init__(self): # checks check_option(self.model, self.OptsModel) check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsMassPrecond) + check_option(self.schur_precond, self.OptsSchurPrecond) # defaults if self.solver_params is None: @@ -326,6 +363,14 @@ def allocate(self): # Projector self._initialize_projectors_and_mass() + # Kronecker preconditioner of the Schur operator, built for the time step in prepare_time_step + self._schur_pc = alfven_schur_preconditioner( + self, + ([] if self._model == "linear_q" else ['model="linear_q"']) + + ([] if self._coupling == "b" else ['coupling="b"']), + ) + self._schur_pc_dt = None + # bunch of temporaries to avoid allocating in the loop u = self.variables.u.spline.vector q = self.variables.q.spline.vector @@ -361,7 +406,16 @@ def allocate(self): self._extracted_b2 = self.derham.extraction_ops["2"].dot(self.projected_equil.b2) self._extracted_q3 = self.derham.extraction_ops["3"].dot(self.projected_equil.q3) + def prepare_time_step(self, dt: float): + """Kronecker preconditioner of the Schur operator for the time step ``dt`` (with + ``schur_precond="StiffnessPreconditioner"``): built in the setup, rebuilt when dt changes.""" + if self._schur_pc is not None and dt != self._schur_pc_dt: + self._inv_Jacobian.precond = self._schur_pc.build(dt) + self._schur_pc_dt = dt + def __call__(self, dt): + # preconditioner for the current time step (no-op unless dt changed since the setup) + self.prepare_time_step(dt) # WMMnew is re-assembled in VariationalDensityEvolve: update the diagonal scaling of its preconditioner self._Mrho_pc.refresh() @@ -416,88 +470,45 @@ def __call_picard(self, dt): self._update_Pib(self._tmp_bn12) self._update_Projq(self._tmp_qn12) - # Compute the advection terms - if self._model == "linear_q": - advection = self.curlPibT0.dot( - self._linear_form_dl_db, - out=self._tmp_advection, - ) - - advection += self.curlPibT.dot( - self._linear_form_dl_db0, - out=self._tmp_advection2, - ) - - advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) + # Compute the advection terms of the couplings advanced here (the other variable is kept fixed) + advection = self._tmp_advection + advection *= 0.0 + b_advection = self._tmp_b_advection + b_advection *= 0.0 + q_advection = self._tmp_q_advection + q_advection *= 0.0 - b_advection = self.curlPib0.dot( - un12, - out=self._tmp_b_advection, - ) - - q_advection = self._transop_q0.dot( - un12, - out=self._tmp_q_advection, - ) + if self._model == "linear_q": + if self._with_b: + advection += self.curlPibT0.dot(self._linear_form_dl_db, out=self._tmp_advection2) + advection += self.curlPibT.dot(self._linear_form_dl_db0, out=self._tmp_advection2) + b_advection += self.curlPib0.dot(un12, out=self._tmp_b_advection2) + if self._with_q: + advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) + q_advection += self._transop_q0.dot(un12, out=self._tmp_q_advection2) elif self._model == "perturbative_q": - advection = self.curlPibT0.dot( - self._linear_form_dl_db, - out=self._tmp_advection, - ) - - advection += self.curlPibT.dot( - self._linear_form_dl_db0, - out=self._tmp_advection2, - ) - - advection += self.curlPibT.dot( - self._linear_form_dl_db, - out=self._tmp_advection2, - ) - - advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) - - b_advection = self.curlPib.dot( - un12, - out=self._tmp_b_advection, - ) - - b_advection += self.curlPib0.dot( - un12, - out=self._tmp_b_advection2, - ) - - q_advection = self._transop_q0.dot( - un12, - out=self._tmp_q_advection, - ) - - q_advection += self._transop_q.dot( - un12, - out=self._tmp_q_advection2, - ) + if self._with_b: + advection += self.curlPibT0.dot(self._linear_form_dl_db, out=self._tmp_advection2) + advection += self.curlPibT.dot(self._linear_form_dl_db0, out=self._tmp_advection2) + advection += self.curlPibT.dot(self._linear_form_dl_db, out=self._tmp_advection2) + b_advection += self.curlPib.dot(un12, out=self._tmp_b_advection2) + b_advection += self.curlPib0.dot(un12, out=self._tmp_b_advection2) + if self._with_q: + advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) + q_advection += self._transop_q0.dot(un12, out=self._tmp_q_advection2) + q_advection += self._transop_q.dot(un12, out=self._tmp_q_advection2) else: - advection = self.curlPibT.dot( - self._linear_form_dl_db, - out=self._tmp_advection, - ) - - advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - - b_advection = self.curlPib.dot( - un12, - out=self._tmp_b_advection, - ) - - q_advection = self._transop_q.dot( - un12, - out=self._tmp_q_advection, - ) + if self._with_b: + advection += self.curlPibT.dot(self._linear_form_dl_db, out=self._tmp_advection2) + b_advection += self.curlPib.dot(un12, out=self._tmp_b_advection2) + if self._with_q: + advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + q_advection += self._transop_q.dot(un12, out=self._tmp_q_advection2) advection *= dt b_advection *= dt @@ -530,16 +541,24 @@ def __call_picard(self, dt): # Newton step self._tmp_f[0] = mn_diff - self._tmp_f[1] = bn_diff - self._tmp_f[2] = qn_diff + if self._coupling == "bq": + self._tmp_f[1] = bn_diff + self._tmp_f[2] = qn_diff + else: + self._tmp_f[1] = bn_diff if self._with_b else qn_diff # normalized right-hand side: the tolerance of the linear solver is relative to the residual incr = dot_normalized(self._inv_Jacobian, self._tmp_f, self._tmp_incr) if self._info: logger.info(f"information on the linear solver : {self._inv_Jacobian._solver._info}") un1 -= incr[0] - bn1 -= incr[1] - qn1 -= incr[2] + if self._coupling == "bq": + bn1 -= incr[1] + qn1 -= incr[2] + elif self._with_b: + bn1 -= incr[1] + else: + qn1 -= incr[1] # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.dot(un1, out=self._tmp_mn1) @@ -568,7 +587,7 @@ def __call_picard(self, dt): # Update the 2nd order variables - if self._qt3 is not None: + if self._qt3 is not None and self._with_q: q_advection = self._transop_q.dot( un12, out=self._tmp_q_advection, @@ -576,7 +595,7 @@ def __call_picard(self, dt): q_advection *= dt self._qt3 -= q_advection - if self._bt2 is not None: + if self._bt2 is not None and self._with_b: b_advection = self.curlPib.dot( un12, out=self._tmp_b_advection, @@ -621,11 +640,13 @@ def _initialize_projectors_and_mass(self): self._I2 = IdentityOperator(self.derham.V2pol) self._I3 = IdentityOperator(self.derham.V3pol) - Jacs = BlockVectorSpace( - self.derham.Vvpol, - self.derham.V2pol, - self.derham.V3pol, - ) + # unknowns of the Newton system: u and the coupled variables + if self._coupling == "bq": + Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V2pol, self.derham.V3pol) + elif self._with_b: + Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V2pol) + else: + Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V3pol) self._tmp_f = Jacs.zeros() self._tmp_incr = Jacs.zeros() @@ -682,17 +703,29 @@ def _initialize_projectors_and_mass(self): # local version to avoid creating new version of LinearOperator every time - self._Jacobian[0, 0] = self._Mrho - self._Jacobian[0, 1] = self._mdt2_pc_curlPibT_M - self._Jacobian[0, 2] = self._mdt2_pc_transopT_M - self._Jacobian[1, 0] = self._dt2_curlPib - self._Jacobian[1, 1] = self._I2 - self._Jacobian[2, 0] = self._dt2_transop - self._Jacobian[2, 2] = self._I3 + from struphy.linear_algebra.schur_solver import SchurSolverFull, SchurSolverFull3 - from struphy.linear_algebra.schur_solver import SchurSolverFull3 + self._Jacobian[0, 0] = self._Mrho + if self._coupling == "bq": + self._Jacobian[0, 1] = self._mdt2_pc_curlPibT_M + self._Jacobian[0, 2] = self._mdt2_pc_transopT_M + self._Jacobian[1, 0] = self._dt2_curlPib + self._Jacobian[1, 1] = self._I2 + self._Jacobian[2, 0] = self._dt2_transop + self._Jacobian[2, 2] = self._I3 + schur_class = SchurSolverFull3 + elif self._with_b: + self._Jacobian[0, 1] = self._mdt2_pc_curlPibT_M + self._Jacobian[1, 0] = self._dt2_curlPib + self._Jacobian[1, 1] = self._I2 + schur_class = SchurSolverFull + else: + self._Jacobian[0, 1] = self._mdt2_pc_transopT_M + self._Jacobian[1, 0] = self._dt2_transop + self._Jacobian[1, 1] = self._I3 + schur_class = SchurSolverFull - self._inv_Jacobian = SchurSolverFull3( + self._inv_Jacobian = schur_class( self._Jacobian, self.options.solver, # fixed (linear) preconditioner: an inner iterative solve is expensive and not a linear operator From de9f5dcb1e46768cd442606c804b3a2aeebabf9f Mon Sep 17 00:00:00 2001 From: Nicolo' Minichelli Date: Mon, 26 Oct 2026 09:03:18 +0100 Subject: [PATCH 2/4] Fix the model tests of VariationalPerturbativeMHD_with_q with split_sound - propagator_variants lists only the propagators of the model (variat_sound only with split_sound, the dissipative ones only if present): test_variational_perturbative_linear failed with "'Propagators' object has no attribute 'variat_sound'". - generate_default_parameter_file writes all non-default constructor arguments (linear, with_viscosity, with_resistivity, split_sound), so that the generated params file builds the same propagators; test_variational_perturbative_linear checks split_sound=True. - The polar StiffnessPreconditioner accepts (and ignores) block_shapes. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 1 + .../models/tests/default_params/test_models.py | 9 +++++++++ .../models/variational_perturbative_mhd_with_q.py | 12 ++++++++---- 3 files changed, 18 insertions(+), 4 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 5fd7078ff..04a07fefe 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -1209,6 +1209,7 @@ def _polar_stiffness_preconditioner( mid: str | LinearOperator | None = None, stab: str | LinearOperator | None = None, stab_average: str | None = None, + block_shapes: bool = True, **polar_kwargs, ) -> PolarStiffnessPreconditioner: """The PolarStiffnessPreconditioner for the arguments of StiffnessPreconditioner (plus its own options).""" diff --git a/src/struphy/models/tests/default_params/test_models.py b/src/struphy/models/tests/default_params/test_models.py index 81b0e3bc0..0de579674 100644 --- a/src/struphy/models/tests/default_params/test_models.py +++ b/src/struphy/models/tests/default_params/test_models.py @@ -115,6 +115,15 @@ def test_variational_perturbative_linear(perturbative_model: str, tmp_path, test for name, variant in model.propagator_variants.items(): assert getattr(params_model.propagators, name).options.model == variant, name + # the q model with split_sound: the flag is kept, and the sound substep gets the linear variant + if perturbative_model == "VariationalPerturbativeMHD_with_q": + split = model_class(linear=True, split_sound=True) + path_split = os.path.join(tmp_path, f"params_{perturbative_model}_split.py") + split.generate_default_parameter_file(path=path_split, prompt=False) + split_model = import_parameters_py(path_split).model + assert split_model.params["split_sound"] and split_model.params["linear"] + assert split_model.propagators.variat_sound.options.model == "linear_q" + ut.call_test(model=model, test_profiling=test_profiling) diff --git a/src/struphy/models/variational_perturbative_mhd_with_q.py b/src/struphy/models/variational_perturbative_mhd_with_q.py index afe5bfb07..b6492d8f6 100644 --- a/src/struphy/models/variational_perturbative_mhd_with_q.py +++ b/src/struphy/models/variational_perturbative_mhd_with_q.py @@ -180,10 +180,11 @@ def velocity_scale(self): @property def propagator_variants(self) -> dict: - """The values of the option ``model`` of the variational propagators that match ``linear``.""" + """The values of the option ``model`` of the variational propagators of this model that match ``linear`` + (``variat_sound`` only with ``split_sound``, the dissipative propagators only if present).""" variant = "linear_q" if self.params["linear"] else "perturbative_q" names = ("variat_dens", "variat_qb", "variat_sound", "variat_viscous", "variat_resist") - return {name: variant for name in names} + return {name: variant for name in names if hasattr(self.propagators, name)} def post_allocate(self): # the propagator options must match the flag `linear` @@ -233,11 +234,14 @@ def generate_default_parameter_file(self, path=None, prompt=True): params_path = super().generate_default_parameter_file(path=path, prompt=prompt) variant = self.propagator_variants["variat_qb"] linear = self.params["linear"] + # the constructor arguments that differ from their defaults (they determine the propagators) + defaults = {"linear": False, "with_viscosity": True, "with_resistivity": True, "split_sound": False} + args = "".join(f", {k}={self.params[k]}" for k, v in defaults.items() if self.params[k] != v) new_file = [] with open(params_path, "r") as f: for line in f: - if line.startswith("model = ") and linear: - new_file += [line.replace("(base_units=base_units)", "(base_units=base_units, linear=True)")] + if line.startswith("model = ") and args: + new_file += [line.replace("(base_units=base_units)", f"(base_units=base_units{args})")] elif "variat_dens.Options" in line: new_file += [ f"model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='{variant}')\n", From 834b38a0d85bf4d5076804b737f48cf6f76158f1 Mon Sep 17 00:00:00 2001 From: Nicolo' Minichelli Date: Mon, 26 Oct 2026 09:19:13 +0100 Subject: [PATCH 3/4] Per-component weight shapes only in the shear-Alfven preconditioner; projected density - StiffnessPreconditioner(block_shapes=False) by default again, so the other propagators are unchanged; ShearAlfvenSchurPreconditioner enables it (main and E x B preconditioner). - The variational Schur preconditioner uses the projected density of the Schur operator (WMMnew's density spline) instead of the analytic equilibrium density (Mvn). The density is passed as a function weight, since WMMnew stores only G as weight and evaluates the spline during the assembly. Lanczos kappa on the TAE grid: 411 -> 391 (block Jacobi), 412 -> 301 (with E x B). - ShearAlfvenSchurPreconditioner accepts the velocity mass operator itself as id_M. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 28 +++++++++++++++++------ src/struphy/feec/variational_utilities.py | 16 +++++++++++-- 2 files changed, 35 insertions(+), 9 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 04a07fefe..cfef0c0e3 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -742,7 +742,8 @@ class StiffnessPreconditioner(KroneckerPreconditioner): block_shapes : bool General operators only: approximate the weights of each diagonal block of ``mid`` and ``stab`` by their own - separable shape (default), instead of one shape common to all blocks. The blocks are inverted independently + separable shape, instead of one shape common to all blocks (default False; used by + :class:`ShearAlfvenSchurPreconditioner`). The blocks are inverted independently (block Jacobi), so this costs nothing and is exact for separable weights with different profiles per component, as the metric coefficients of curvilinear maps (e.g. :math:`G_{22} \sqrt g \propto r^3` for the poloidal component of a vector field in a torus); in a Cartesian geometry all blocks have the same shape. @@ -776,7 +777,7 @@ def __init__( mid: str | LinearOperator | None = None, stab: str | LinearOperator | None = None, stab_average: str | None = None, - block_shapes: bool = True, + block_shapes: bool = False, ): derham = mass_ops.derham self._block_shapes = block_shapes @@ -1209,7 +1210,7 @@ def _polar_stiffness_preconditioner( mid: str | LinearOperator | None = None, stab: str | LinearOperator | None = None, stab_average: str | None = None, - block_shapes: bool = True, + block_shapes: bool = False, **polar_kwargs, ) -> PolarStiffnessPreconditioner: """The PolarStiffnessPreconditioner for the arguments of StiffnessPreconditioner (plus its own options).""" @@ -1917,8 +1918,8 @@ class ShearAlfvenSchurPreconditioner: T : BasisProjectionOperator The basis projection operator :math:`\mathbb T` of :math:`\mathbf B_0 \times` (on 2-forms). - id_M : str - Name of the weighted mass matrix of the velocity in ``mass_ops`` (e.g. ``"M2n"``). + id_M : str | WeightedMassOperator + The weighted mass matrix of the velocity: its name in ``mass_ops`` (e.g. ``"M2n"``) or the operator itself. exb_correction : bool Whether to add the auxiliary space correction of the :math:`\mathbf E \times \mathbf B` flows. @@ -1927,7 +1928,16 @@ class ShearAlfvenSchurPreconditioner: Space of the velocity, ``"Hdiv"`` (2-form) or ``"H1vec"`` (vector field), the domain of ``T``. """ - def __init__(self, derham, mass_ops, basis_ops, T, id_M: str, exb_correction: bool = True, u_space: str = "Hdiv"): + def __init__( + self, + derham, + mass_ops, + basis_ops, + T, + id_M: str | WeightedMassOperator, + exb_correction: bool = True, + u_space: str = "Hdiv", + ): from struphy.feec.kronecker_block import basis_projection_operator, derivative_operator from struphy.feec.utilities import LocalRotationMatrix @@ -2002,7 +2012,10 @@ def resolve(precond: str, derham, algo: str, solver: str, u_space: str) -> str: def build(self, dt: float) -> LinearOperator: """The preconditioner of the Schur operator for the time step ``dt``.""" mid = (dt**2 / 4) * self._mass_ops.M2 - pc = StiffnessPreconditioner(self._mass_ops, self._kron_curl_T, sigma=1.0, mid=mid, stab=self._id_M) + # separable weights per velocity component (their metric weights differ strongly in curvilinear geometry) + pc = StiffnessPreconditioner( + self._mass_ops, self._kron_curl_T, sigma=1.0, mid=mid, stab=self._id_M, block_shapes=True + ) if self._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 = StiffnessPreconditioner( @@ -2010,6 +2023,7 @@ def build(self, dt: float) -> LinearOperator: [self._kron_exb, self._kron_curl_T_exb], sigma=0.0, mid=[self._id_M, mid], + block_shapes=True, ) pc = AuxiliarySpacePreconditioner(pc, [(self._exb_N, self._exb_NT, pc_phi)]) return pc diff --git a/src/struphy/feec/variational_utilities.py b/src/struphy/feec/variational_utilities.py index 26b826fd1..aebd1b1b3 100644 --- a/src/struphy/feec/variational_utilities.py +++ b/src/struphy/feec/variational_utilities.py @@ -1890,7 +1890,8 @@ def alfven_schur_preconditioner(propagator, unmet: list[str]): as :class:`~struphy.feec.preconditioner.ShearAlfvenSchurPreconditioner` with ``u_space="H1vec"``: the :class:`~struphy.feec.preconditioner.StiffnessPreconditioner` of the Kronecker approximation of :math:`\mathbb C\, \hat\Pi^1[\mathbf B_0 \times \cdot\,]` (``basis_ops.Tv``, analytic equilibrium) with the - stabilization ``Mvn`` (the velocity mass matrix with the analytic density), plus the auxiliary space correction of + stabilization :math:`\mathbb M^v[\rho_0]` with the projected density of the Schur operator (``WMMnew``), plus the + auxiliary space correction of the :math:`\mathbf E \times \mathbf B` flows if ``options.exb_correction``. Its :meth:`build` gives the preconditioner for a time step. @@ -1925,12 +1926,23 @@ def alfven_schur_preconditioner(propagator, unmet: list[str]): "falling back to the mass-matrix preconditioner." ) return None + # the velocity mass matrix of the Schur operator, G weighted with the projected density: WMMnew stores only G as + # weight and evaluates its density spline during the assembly, so the density is passed here as a function, which + # is evaluated into the weights seen by the separable approximation of the preconditioner + rho = propagator.mass_ops.WMMnew.spline_functions["l2_field"] + stab = propagator.mass_ops.create_weighted_mass( + "H1vec", + "H1vec", + weights=("G", lambda e1, e2, e3: rho(e1, e2, e3)), + name="WMM_alfven_precond", + assemble=False, + ) return ShearAlfvenSchurPreconditioner( propagator.derham, propagator.mass_ops, propagator.basis_ops, propagator.basis_ops.Tv, - "Mvn", + stab, exb_correction=options.exb_correction, u_space="H1vec", ) From bcc1a707a563d2fd714a4d82a38b783ee4de9215 Mon Sep 17 00:00:00 2001 From: Nicolo' Minichelli Date: Mon, 26 Oct 2026 10:16:25 +0100 Subject: [PATCH 4/4] Variational q model: separate Alfven and sound propagators, split step by default Review of #829: the coupling argument of VariationalQBEvolve is replaced by two propagators. - VariationalAlfvenEvolve (u, b): shear-Alfven/fast waves, variants full/linear/perturbative (independent of the thermodynamic variable), auxiliary b^, implicit_weights, relative Newton criteria; its Schur complement is the shear-Alfven operator, so it has the options schur_precond / exb_correction. - VariationalSoundEvolve (u, q): sound waves, variants full_q/linear_q/perturbative_q, auxiliary q^, diagnostics div_u and u2. - VariationalPerturbativeMHD_with_q(split_sound=True) is now the default (variat_alfven and variat_sound); split_sound=False keeps the coupled VariationalQBEvolve, which is back to its version on devel. The p model is unchanged. - Tests: the verification tests 1-9 use the unsplit step (same scheme as the full-f model, second order in time); the perturbative energy test runs both steps; one-step tests of the three propagators; the q slab examples keep split_sound=False (their documented results). - Low-mode TAE (q, 4 ranks): Alfven 163-171 / sound 8-9 PCG iterations as before, 1.9 s per step (was 2.5 s: no work for the variable that is kept fixed). Co-Authored-By: Claude Opus 5.5 --- .../subsections/propagators-avail.rst | 10 + .../with_q/lowmode_tae/params_lowmode_tae.py | 17 +- .../slab_dispersion/params_slab_dispersion.py | 6 +- .../params_slab_dispersion_linear.py | 4 +- .../tests/default_params/test_models.py | 16 +- .../tests/test_variational_dissipation.py | 3 +- ...verif_VariationalPerturbativeMHD_with_q.py | 32 +- .../verification/variational_mhd_utils.py | 25 +- .../variational_perturbative_mhd_with_q.py | 110 ++--- src/struphy/propagators/__init__.py | 6 + .../tests/test_variational_qb_evolve.py | 21 +- .../tests/test_variational_schur_precond.py | 17 +- .../propagators/variational_alfven_evolve.py | 450 ++++++++++++++++++ .../propagators/variational_qb_evolve.py | 227 ++++----- .../propagators/variational_sound_evolve.py | 441 +++++++++++++++++ 15 files changed, 1146 insertions(+), 239 deletions(-) create mode 100644 src/struphy/propagators/variational_alfven_evolve.py create mode 100644 src/struphy/propagators/variational_sound_evolve.py diff --git a/doc/sections/subsections/propagators-avail.rst b/doc/sections/subsections/propagators-avail.rst index ac971567a..9dd572831 100644 --- a/doc/sections/subsections/propagators-avail.rst +++ b/doc/sections/subsections/propagators-avail.rst @@ -72,6 +72,11 @@ Linear MHD solvers Nonlinear MHD solvers ^^^^^^^^^^^^^^^^^^^^^ +.. automodule:: struphy.propagators.variational_alfven_evolve + :members: + :exclude-members: options, allocate + :show-inheritance: + .. automodule:: struphy.propagators.variational_density_evolve :members: :exclude-members: options, allocate @@ -107,6 +112,11 @@ Nonlinear MHD solvers :exclude-members: options, allocate :show-inheritance: +.. automodule:: struphy.propagators.variational_sound_evolve + :members: + :exclude-members: options, allocate + :show-inheritance: + .. automodule:: struphy.propagators.variational_viscosity :members: :exclude-members: options, allocate diff --git a/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py b/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py index bdfc1e62b..d284fe814 100644 --- a/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py +++ b/examples/VariationalPerturbativeLinearMHD/with_q/lowmode_tae/params_lowmode_tae.py @@ -49,9 +49,9 @@ The variational propagators run with model="linear_q" and their default (Newton/Picard) solver options, except implicit_weights=False: the weights of the transport operators at t^n, i.e. first order in time like the p model (no midpoint option) and LinearMHD (Crank-Nicolson substeps, Lie-Trotter -splitting). With split_sound=True the coupling of the velocity to B (shear-Alfven/fast waves, variat_qb) and to q -(sound waves, variat_sound) are advanced in two implicit substeps, like the shear-Alfven and magnetosonic steps of -LinearMHD: each substep is energy conserving and needs a single Picard iteration, with a smaller linear system. +splitting). The coupling of the velocity to B (shear-Alfven/fast waves, VariationalAlfvenEvolve) and to q (sound +waves, VariationalSoundEvolve) are advanced in two implicit substeps (split_sound=True, the default), like the +shear-Alfven and magnetosonic steps of LinearMHD: each substep is energy conserving and needs a single Newton step. Expected result. With n = 1 and m = 1 the gap is at q(r_c) = 3/2, i.e. r_c = 0.5, where n = 1 and |B| = 1.0 (flux-surface average), so v_A(r_c) = 1.0 v_unit and @@ -137,9 +137,9 @@ base_units = BaseUnits() # Model instance: linear and ideal (no viscosity, no resistivity), as LinearMHD -# split_sound: shear-Alfven/fast and sound waves in two substeps, as in LinearMHD +# split_sound=True (default): shear-Alfven/fast and sound waves in two substeps, as in LinearMHD model = VariationalPerturbativeMHD_with_q( - base_units=base_units, linear=True, with_viscosity=False, with_resistivity=False, split_sound=True + base_units=base_units, linear=True, with_viscosity=False, with_resistivity=False ) # List all variables and decide whether to save their data @@ -197,12 +197,13 @@ # ------------------ # Propagator options # ------------------ -# linear=True needs the "linear_q" variants of the variational propagators (checked by the model) +# linear=True needs the "linear_q" variants of the variational propagators, and "linear" for the Alfven step +# (checked by the model) model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model="linear_q", gamma=gamma) # implicit_weights=False: first order in time, as the p model and LinearMHD (Lie-Trotter splitting) -model.propagators.variat_qb.options = model.propagators.variat_qb.Options( - model="linear_q", gamma=gamma, implicit_weights=False +model.propagators.variat_alfven.options = model.propagators.variat_alfven.Options( + model="linear", implicit_weights=False ) model.propagators.variat_sound.options = model.propagators.variat_sound.Options( model="linear_q", gamma=gamma, implicit_weights=False diff --git a/examples/VariationalPerturbativeMHD_with_q/slab_dispersion/params_slab_dispersion.py b/examples/VariationalPerturbativeMHD_with_q/slab_dispersion/params_slab_dispersion.py index 5a2c0438d..4cb9f1017 100644 --- a/examples/VariationalPerturbativeMHD_with_q/slab_dispersion/params_slab_dispersion.py +++ b/examples/VariationalPerturbativeMHD_with_q/slab_dispersion/params_slab_dispersion.py @@ -95,7 +95,11 @@ # --------------------- base_units = BaseUnits() -model = VariationalPerturbativeMHD_with_q(base_units=base_units, with_viscosity=False, with_resistivity=False) +# split_sound=False: u, B and q together in VariationalQBEvolve (second order in time with the default implicit +# weights, the results below) +model = VariationalPerturbativeMHD_with_q( + base_units=base_units, with_viscosity=False, with_resistivity=False, split_sound=False +) model.em_fields.b_field.save_data = True model.mhd.density.save_data = True diff --git a/examples/VariationalPerturbativeMHD_with_q/slab_dispersion_linear/params_slab_dispersion_linear.py b/examples/VariationalPerturbativeMHD_with_q/slab_dispersion_linear/params_slab_dispersion_linear.py index 9466116b3..73adb103c 100644 --- a/examples/VariationalPerturbativeMHD_with_q/slab_dispersion_linear/params_slab_dispersion_linear.py +++ b/examples/VariationalPerturbativeMHD_with_q/slab_dispersion_linear/params_slab_dispersion_linear.py @@ -95,8 +95,10 @@ # --------------------- base_units = BaseUnits() +# split_sound=False: u, B and q together in VariationalQBEvolve (second order in time with the default implicit +# weights, the results below) model = VariationalPerturbativeMHD_with_q( - base_units=base_units, linear=True, with_viscosity=False, with_resistivity=False + base_units=base_units, linear=True, with_viscosity=False, with_resistivity=False, split_sound=False ) model.em_fields.b_field.save_data = True diff --git a/src/struphy/models/tests/default_params/test_models.py b/src/struphy/models/tests/default_params/test_models.py index 0de579674..667fc48cf 100644 --- a/src/struphy/models/tests/default_params/test_models.py +++ b/src/struphy/models/tests/default_params/test_models.py @@ -115,14 +115,16 @@ def test_variational_perturbative_linear(perturbative_model: str, tmp_path, test for name, variant in model.propagator_variants.items(): assert getattr(params_model.propagators, name).options.model == variant, name - # the q model with split_sound: the flag is kept, and the sound substep gets the linear variant + # the q model with the unsplit step (split_sound=False): the flag is kept, and VariationalQBEvolve gets the + # linear variant if perturbative_model == "VariationalPerturbativeMHD_with_q": - split = model_class(linear=True, split_sound=True) - path_split = os.path.join(tmp_path, f"params_{perturbative_model}_split.py") - split.generate_default_parameter_file(path=path_split, prompt=False) - split_model = import_parameters_py(path_split).model - assert split_model.params["split_sound"] and split_model.params["linear"] - assert split_model.propagators.variat_sound.options.model == "linear_q" + assert params_model.propagators.variat_alfven.options.model == "linear" + unsplit = model_class(linear=True, split_sound=False) + path_unsplit = os.path.join(tmp_path, f"params_{perturbative_model}_unsplit.py") + unsplit.generate_default_parameter_file(path=path_unsplit, prompt=False) + unsplit_model = import_parameters_py(path_unsplit).model + assert not unsplit_model.params["split_sound"] and unsplit_model.params["linear"] + assert unsplit_model.propagators.variat_qb.options.model == "linear_q" ut.call_test(model=model, test_profiling=test_profiling) diff --git a/src/struphy/models/tests/test_variational_dissipation.py b/src/struphy/models/tests/test_variational_dissipation.py index 6868b8646..64533a913 100644 --- a/src/struphy/models/tests/test_variational_dissipation.py +++ b/src/struphy/models/tests/test_variational_dissipation.py @@ -26,7 +26,8 @@ def test_linear_q_viscosity_resistivity(fast, tmp_path): model = VariationalPerturbativeMHD_with_q(linear=True) props = model.propagators props.variat_dens.options = props.variat_dens.Options(model="linear_q") - props.variat_qb.options = props.variat_qb.Options(model="linear_q") + props.variat_alfven.options = props.variat_alfven.Options(model="linear") + props.variat_sound.options = props.variat_sound.Options(model="linear_q") props.variat_viscous.options = props.variat_viscous.Options(model="linear_q", mu=0.1, fast=fast) props.variat_resist.options = props.variat_resist.Options(model="linear_q", eta=0.1, fast=fast) model.mhd.density.add_background(FieldsBackground()) diff --git a/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py b/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py index e37bcd345..103902375 100644 --- a/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py +++ b/src/struphy/models/tests/verification/test_verif_VariationalPerturbativeMHD_with_q.py @@ -23,10 +23,16 @@ 8. it gives the analytic frequencies of the shear Alfvén, the fast magnetosonic and the sound wave in a homogeneous slab, with second-order convergence in the time step; 9. the p-based linear model agrees up to its first-order time discretization error; -10. with ``split_sound=True`` (the shear-Alfven/fast and the sound waves in two substeps) the energy is conserved; +10. with ``split_sound=True`` (the default: the shear-Alfven/fast and the sound waves in two substeps, + VariationalAlfvenEvolve and VariationalSoundEvolve) the energy is conserved, for the perturbative (1.) and the + linear model; 11. and the split model agrees with the unsplit one up to a splitting error of first order in the time step. + +Tests 1.-9. use the unsplit step (``split_sound=False``, VariationalQBEvolve), which gives the same scheme as the +full-f model (3.) and second-order convergence in the time step (8.). """ +import functools import logging import cunumpy as xp @@ -55,17 +61,21 @@ set_logging_level() logger = logging.getLogger("struphy") -LINEAR = linear(VariationalPerturbativeMHD_with_q) +# the unsplit step (VariationalQBEvolve) and the split step (the default) +PERT = functools.partial(VariationalPerturbativeMHD_with_q, split_sound=False) +PERT.__name__ = "VariationalPerturbativeMHD_with_q_unsplit" +LINEAR = linear(VariationalPerturbativeMHD_with_q, split_sound=False) LINEAR_SPLIT = linear(VariationalPerturbativeMHD_with_q, split_sound=True) -def test_energy_conservation(tmp_path): +@pytest.mark.parametrize("split_sound", [False, True]) +def test_energy_conservation(tmp_path, split_sound): """The energy H = int rho |u|^2/2 + (q~^2 + 2 q0 q^)/(gamma-1) + |B~|^2/2 + B0.B^ is conserved, around a - non-uniform equilibrium and for a nonlinear amplitude.""" + non-uniform equilibrium and for a nonlinear amplitude, with the unsplit and the split step.""" model = allocate_ideal( tmp_path, - "energy", - VariationalPerturbativeMHD_with_q, + f"energy_{split_sound}", + functools.partial(VariationalPerturbativeMHD_with_q, split_sound=split_sound), sheared_slab(), dt=0.05, perturb=perturb_velocity(0.05), @@ -89,7 +99,7 @@ def test_equilibrium_preservation(tmp_path): setup = sheared_slab() dt, n_steps = 0.05, 5 - pert = allocate_ideal(tmp_path, "pert", VariationalPerturbativeMHD_with_q, setup, dt) + pert = allocate_ideal(tmp_path, "pert", PERT, setup, dt) for _ in range(n_steps): pert.integrate(dt) for vec in state(pert): @@ -124,7 +134,7 @@ def run(model_class, perturb): model.integrate(dt) return initial, state(model) - pert_0, pert_n = run(VariationalPerturbativeMHD_with_q, perturb_all_along_z(amp, Lz)) + pert_0, pert_n = run(PERT, perturb_all_along_z(amp, Lz)) full_0, full_n = run(VariationalMHD_with_q, perturb_all_along_z(amp, Lz)) # both runs are converged up to the nonlinear solver tolerance (1e-8, relative to the perturbation in the @@ -154,7 +164,7 @@ def finalstate(model_class, dt): rel_diff = {} for dt in (0.05, 0.025): - q_n = finalstate(VariationalPerturbativeMHD_with_q, dt) + q_n = finalstate(PERT, dt) p_n = finalstate(VariationalPerturbativeMHD, dt) rel_diff[dt] = [norm(q_n[i] - p_n[i]) / norm(q_n[i]) for i in (1, 3)] @@ -187,7 +197,7 @@ def run(model_class, amp): u_lin, omega_lin = run(LINEAR, 1.0) rel_diff = [] for amp in (1e-2, 1e-3): - u_pert, omega_pert = run(VariationalPerturbativeMHD_with_q, amp) + u_pert, omega_pert = run(PERT, amp) rel_diff.append(norm(u_pert / amp - u_lin) / norm(u_lin)) logger.info(f"{wave}: amp {amp}, omega {omega_pert:.5f}") assert abs(omega_pert - omega_exact) < 2e-3 * omega_exact, (omega_pert, omega_exact) @@ -246,7 +256,7 @@ def test_limit_of_perturbative_model(tmp_path): u_lin = _final_state(tmp_path, LINEAR, setup, 1.0, 0.05, 10)[1] rel_diff = [] for amp in (1e-2, 1e-3): - u_pert = _final_state(tmp_path, VariationalPerturbativeMHD_with_q, setup, amp, 0.05, 10)[1] + u_pert = _final_state(tmp_path, PERT, setup, amp, 0.05, 10)[1] rel_diff.append(norm(u_pert / amp - u_lin) / norm(u_lin)) logger.info(f"relative difference perturbative / linear for amplitudes 1e-2, 1e-3: {rel_diff}") assert rel_diff[0] < 5e-2, rel_diff diff --git a/src/struphy/models/tests/verification/variational_mhd_utils.py b/src/struphy/models/tests/verification/variational_mhd_utils.py index 6b8223184..5e10e220c 100644 --- a/src/struphy/models/tests/verification/variational_mhd_utils.py +++ b/src/struphy/models/tests/verification/variational_mhd_utils.py @@ -73,8 +73,8 @@ def allocate_ideal( """Allocate an ideal model of `model_class` for the `setup` (equil, domain, num_elements, degree). `perturb(model)` adds the perturbations of the perturbative variables. The full-f model gets the equilibrium - as background in addition. `implicit_weights` is passed to VariationalQBEvolve (None: its default), - `schur_precond` to VariationalPBEvolve and VariationalQBEvolve. Models must be + as background in addition. `implicit_weights` is passed to VariationalQBEvolve, VariationalAlfvenEvolve and + VariationalSoundEvolve (None: their default), `schur_precond` to VariationalPBEvolve and VariationalAlfvenEvolve. Models must be allocated and run one after the other, since they share the class attributes of Propagator (derham, ...).""" equil, domain, num_elements, degree = setup model = model_class(with_viscosity=False, with_resistivity=False) @@ -87,21 +87,26 @@ def allocate_ideal( model.propagators.variat_mom.options = model.propagators.variat_mom.Options( solver_params=SOLVER, nonlin_solver=NONLIN ) - if hasattr(model.propagators, "variat_pb"): - model.propagators.variat_pb.options = model.propagators.variat_pb.Options( + props = model.propagators + if hasattr(props, "variat_pb"): + props.variat_pb.options = props.variat_pb.Options( model=variant, solver_params=SOLVER, nonlin_solver=NONLIN, schur_precond=schur_precond ) - else: - model.propagators.variat_qb.options = model.propagators.variat_qb.Options( - model=variant, + if hasattr(props, "variat_qb"): + props.variat_qb.options = props.variat_qb.Options( + model=variant, solver_params=SOLVER, nonlin_solver=NONLIN, implicit_weights=implicit_weights + ) + # the two substeps of split_sound=True (the Alfven step does not depend on the thermodynamic variable) + if hasattr(props, "variat_alfven"): + props.variat_alfven.options = props.variat_alfven.Options( + model=variant.removesuffix("_q"), solver_params=SOLVER, nonlin_solver=NONLIN, implicit_weights=implicit_weights, schur_precond=schur_precond, ) - # the sound substep of split_sound=True - if hasattr(model.propagators, "variat_sound"): - model.propagators.variat_sound.options = model.propagators.variat_sound.Options( + if hasattr(props, "variat_sound"): + props.variat_sound.options = props.variat_sound.Options( model=variant, solver_params=SOLVER, nonlin_solver=NONLIN, implicit_weights=implicit_weights ) diff --git a/src/struphy/models/variational_perturbative_mhd_with_q.py b/src/struphy/models/variational_perturbative_mhd_with_q.py index b6492d8f6..b2d390996 100644 --- a/src/struphy/models/variational_perturbative_mhd_with_q.py +++ b/src/struphy/models/variational_perturbative_mhd_with_q.py @@ -14,10 +14,12 @@ from struphy.models.variables import FEECVariable from struphy.polar.basic import PolarVector from struphy.propagators.base import Propagator +from struphy.propagators.variational_alfven_evolve import VariationalAlfvenEvolve from struphy.propagators.variational_density_evolve import VariationalDensityEvolve from struphy.propagators.variational_momentum_advection import VariationalMomentumAdvection from struphy.propagators.variational_qb_evolve import VariationalQBEvolve from struphy.propagators.variational_resistivity import VariationalResistivity +from struphy.propagators.variational_sound_evolve import VariationalSoundEvolve from struphy.propagators.variational_viscosity import VariationalViscosity @@ -39,10 +41,10 @@ class VariationalPerturbativeMHD_with_q(StruphyModel): with_resistivity: bool Whether to include resistive dissipation (default: True) split_sound: bool - Whether to advance the coupling of the velocity to the magnetic field and to q in two substeps - (``variat_qb`` and ``variat_sound``, both energy conserving), instead of together in ``variat_qb``. The - Newton systems are then smaller (shear-Alfven/fast and sound waves separately), at the cost of a splitting - error (default: False) + Whether to advance the coupling of the velocity to the magnetic field (shear-Alfvén/fast waves, + ``variat_alfven``) and to q (sound waves, ``variat_sound``) in two substeps, both energy conserving (default: + True). With False they are advanced together in ``variat_qb``, without the splitting error but with a larger + and stiffer Newton system. """ @classmethod @@ -84,15 +86,15 @@ def __init__( linear: bool = False, with_viscosity: bool = True, with_resistivity: bool = True, - split_sound: bool = False, + split_sound: bool = True, ): self.variat_dens = VariationalDensityEvolve() if not linear: self.variat_mom = VariationalMomentumAdvection() if split_sound: # the shear-Alfven/fast waves (u, b) and the sound waves (u, q) in two substeps - self.variat_qb = VariationalQBEvolve(div_u=div_u, u2=u2, qt3=qt3, bt2=bt2, coupling="b") - self.variat_sound = VariationalQBEvolve(div_u=div_u, u2=u2, qt3=qt3, bt2=bt2, coupling="q") + self.variat_alfven = VariationalAlfvenEvolve(bt2=bt2) + self.variat_sound = VariationalSoundEvolve(div_u=div_u, u2=u2, qt3=qt3) else: self.variat_qb = VariationalQBEvolve(div_u=div_u, u2=u2, qt3=qt3, bt2=bt2) if with_viscosity: @@ -109,7 +111,7 @@ def __init__( linear: bool = False, with_viscosity: bool = True, with_resistivity: bool = True, - split_sound: bool = False, + split_sound: bool = True, ): # 0. store input parameters @@ -141,13 +143,15 @@ def __init__( self.propagators.variat_dens.variables.u = self.mhd.velocity if not linear: self.propagators.variat_mom.variables.u = self.mhd.velocity - self.propagators.variat_qb.variables.u = self.mhd.velocity - self.propagators.variat_qb.variables.q = self.mhd.sqrt_p - self.propagators.variat_qb.variables.b = self.em_fields.b_field if split_sound: + self.propagators.variat_alfven.variables.u = self.mhd.velocity + self.propagators.variat_alfven.variables.b = self.em_fields.b_field self.propagators.variat_sound.variables.u = self.mhd.velocity self.propagators.variat_sound.variables.q = self.mhd.sqrt_p - self.propagators.variat_sound.variables.b = self.em_fields.b_field + else: + self.propagators.variat_qb.variables.u = self.mhd.velocity + self.propagators.variat_qb.variables.q = self.mhd.sqrt_p + self.propagators.variat_qb.variables.b = self.em_fields.b_field if with_viscosity: self.propagators.variat_viscous.variables.s = self.mhd.sqrt_p self.propagators.variat_viscous.variables.u = self.mhd.velocity @@ -181,10 +185,20 @@ def velocity_scale(self): @property def propagator_variants(self) -> dict: """The values of the option ``model`` of the variational propagators of this model that match ``linear`` - (``variat_sound`` only with ``split_sound``, the dissipative propagators only if present).""" + (``variat_alfven`` and ``variat_sound`` with ``split_sound``, else ``variat_qb``; the dissipative propagators + only if present). ``variat_alfven`` does not depend on the thermodynamic variable.""" variant = "linear_q" if self.params["linear"] else "perturbative_q" - names = ("variat_dens", "variat_qb", "variat_sound", "variat_viscous", "variat_resist") - return {name: variant for name in names if hasattr(self.propagators, name)} + alfven = "linear" if self.params["linear"] else "perturbative" + names = ("variat_dens", "variat_alfven", "variat_sound", "variat_qb", "variat_viscous", "variat_resist") + return { + name: alfven if name == "variat_alfven" else variant for name in names if hasattr(self.propagators, name) + } + + @property + def _thermo_propagator(self): + """The propagator of q (with the adiabatic index and the auxiliary variable qt3).""" + sound = getattr(self.propagators, "variat_sound", None) + return sound if sound is not None else self.propagators.variat_qb def post_allocate(self): # the propagator options must match the flag `linear` @@ -195,13 +209,6 @@ def post_allocate(self): f"{self.__class__.__name__}(linear={self.params['linear']}) needs " f"{name}.Options(model='{variant}'), got model='{prop.options.model}'." ) - # both substeps of the split step and the energy diagnostics must use the same adiabatic index - sound = getattr(self.propagators, "variat_sound", None) - if sound is not None and sound.options.gamma != self.propagators.variat_qb.options.gamma: - raise ValueError( - f"{self.__class__.__name__}(split_sound=True) needs the same gamma in variat_qb and variat_sound, " - f"got {self.propagators.variat_qb.options.gamma} and {sound.options.gamma}." - ) projV3 = L2Projector("L2", Propagator.mass_ops) @@ -221,46 +228,31 @@ def f(e1, e2, e3): def _compute_en_thermo_1(self): q = self.mhd.sqrt_p.spline.vector - gamma = self.propagators.variat_qb.options.gamma + gamma = self._thermo_propagator.options.gamma return 1.0 / (gamma - 1.0) * Propagator.mass_ops.M3.dot_inner(q, q) def _compute_en_thermo_2(self): - qt3 = self.propagators.variat_qb.qt3.spline.vector - gamma = self.propagators.variat_qb.options.gamma + qt3 = self._thermo_propagator.qt3.spline.vector + gamma = self._thermo_propagator.options.gamma return 2.0 / (gamma - 1.0) * Propagator.mass_ops.M3.dot_inner(qt3, Propagator.projected_equil.q3) # default parameters def generate_default_parameter_file(self, path=None, prompt=True): params_path = super().generate_default_parameter_file(path=path, prompt=prompt) - variant = self.propagator_variants["variat_qb"] + variants = self.propagator_variants linear = self.params["linear"] # the constructor arguments that differ from their defaults (they determine the propagators) - defaults = {"linear": False, "with_viscosity": True, "with_resistivity": True, "split_sound": False} + defaults = {"linear": False, "with_viscosity": True, "with_resistivity": True, "split_sound": True} args = "".join(f", {k}={self.params[k]}" for k, v in defaults.items() if self.params[k] != v) new_file = [] with open(params_path, "r") as f: for line in f: + name = next((n for n in variants if f"{n}.Options" in line), None) if line.startswith("model = ") and args: new_file += [line.replace("(base_units=base_units)", f"(base_units=base_units{args})")] - elif "variat_dens.Options" in line: - new_file += [ - f"model.propagators.variat_dens.options = model.propagators.variat_dens.Options(model='{variant}')\n", - ] - elif "variat_qb.Options" in line: - new_file += [ - f"model.propagators.variat_qb.options = model.propagators.variat_qb.Options(model='{variant}')\n", - ] - elif "variat_sound.Options" in line: + elif name is not None: new_file += [ - f"model.propagators.variat_sound.options = model.propagators.variat_sound.Options(model='{variant}')\n", - ] - elif "variat_viscous.Options" in line: - new_file += [ - f"model.propagators.variat_viscous.options = model.propagators.variat_viscous.Options(model='{variant}')\n", - ] - elif "variat_resist.Options" in line: - new_file += [ - f"model.propagators.variat_resist.options = model.propagators.variat_resist.Options(model='{variant}')\n", + f"model.propagators.{name}.options = model.propagators.{name}.Options(model='{variants[name]}')\n", ] elif "sqrt_p.add_background" in line and linear: new_file += ["model.mhd.density.add_background(FieldsBackground())\n"] @@ -382,15 +374,15 @@ def doc_discretization(cls): 1. :class:`~struphy.propagators.variational_density_evolve.VariationalDensityEvolve` 2. :class:`~struphy.propagators.variational_momentum_advection.VariationalMomentumAdvection` (if :attr:`linear` is False) - 3. :class:`~struphy.propagators.variational_qb_evolve.VariationalQBEvolve` - 4. :class:`~struphy.propagators.variational_viscosity.VariationalViscosity` (if :attr:`with_viscosity` is True) - 5. :class:`~struphy.propagators.variational_resistivity.VariationalResistivity` (if :attr:`with_resistivity` is True) + 3. :class:`~struphy.propagators.variational_alfven_evolve.VariationalAlfvenEvolve` (shear-Alfvén/fast waves) and + 4. :class:`~struphy.propagators.variational_sound_evolve.VariationalSoundEvolve` (sound waves), if + :attr:`split_sound` is True (default); otherwise + :class:`~struphy.propagators.variational_qb_evolve.VariationalQBEvolve` (both together) + 5. :class:`~struphy.propagators.variational_viscosity.VariationalViscosity` (if :attr:`with_viscosity` is True) + 6. :class:`~struphy.propagators.variational_resistivity.VariationalResistivity` (if :attr:`with_resistivity` is True) The option ``model`` of the variational propagators is ``"perturbative_q"``, or ``"linear_q"`` if - :attr:`linear` is True. With :attr:`split_sound` the step 3 is done in two substeps of - :class:`~struphy.propagators.variational_qb_evolve.VariationalQBEvolve`: ``variat_qb`` with - ``coupling="b"`` (shear-Alfven/fast waves, :math:`q` fixed) and ``variat_sound`` with ``coupling="q"`` (sound - waves, :math:`\mathbf B` fixed), both energy conserving. + :attr:`linear` is True (``"perturbative"`` or ``"linear"`` for VariationalAlfvenEvolve). """ doc = rf"""**1. VariationalDensityEvolve:** @@ -400,15 +392,23 @@ def doc_discretization(cls): {VariationalMomentumAdvection.__doc__} -**3. VariationalQBEvolve:** +**3. VariationalAlfvenEvolve:** + +{VariationalAlfvenEvolve.__doc__} + +**4. VariationalSoundEvolve:** + +{VariationalSoundEvolve.__doc__} + +**3 and 4 together (split_sound=False). VariationalQBEvolve:** {VariationalQBEvolve.__doc__} -**4. VariationalViscosity:** +**5. VariationalViscosity:** {VariationalViscosity.__doc__} -**5. VariationalResistivity:** +**6. VariationalResistivity:** {VariationalResistivity.__doc__} """ diff --git a/src/struphy/propagators/__init__.py b/src/struphy/propagators/__init__.py index 770e51e8b..5bacdb28a 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -41,6 +41,7 @@ "ShearAlfvenPropagator": "struphy.propagators.shear_alfven_propagator", "TimeDependentSource": "struphy.propagators.time_dependent_source", "TwoFluidQuasiNeutralFull": "struphy.propagators.two_fluid_quasi_neutral_full", + "VariationalAlfvenEvolve": "struphy.propagators.variational_alfven_evolve", "VariationalDensityEvolve": "struphy.propagators.variational_density_evolve", "VariationalEntropyEvolve": "struphy.propagators.variational_entropy_evolve", "VariationalMagFieldEvolve": "struphy.propagators.variational_mag_field_evolve", @@ -48,6 +49,7 @@ "VariationalPBEvolve": "struphy.propagators.variational_pb_evolve", "VariationalQBEvolve": "struphy.propagators.variational_qb_evolve", "VariationalResistivity": "struphy.propagators.variational_resistivity", + "VariationalSoundEvolve": "struphy.propagators.variational_sound_evolve", "VariationalViscosity": "struphy.propagators.variational_viscosity", "VlasovAmpereCoupling": "struphy.propagators.vlasov_ampere_coupling", "VlasovMaxwellCoupling": "struphy.propagators.vlasov_maxwell_coupling", @@ -94,6 +96,7 @@ from struphy.propagators.shear_alfven_propagator import ShearAlfvenPropagator from struphy.propagators.time_dependent_source import TimeDependentSource from struphy.propagators.two_fluid_quasi_neutral_full import TwoFluidQuasiNeutralFull + from struphy.propagators.variational_alfven_evolve import VariationalAlfvenEvolve from struphy.propagators.variational_density_evolve import VariationalDensityEvolve from struphy.propagators.variational_entropy_evolve import VariationalEntropyEvolve from struphy.propagators.variational_mag_field_evolve import VariationalMagFieldEvolve @@ -101,6 +104,7 @@ from struphy.propagators.variational_pb_evolve import VariationalPBEvolve from struphy.propagators.variational_qb_evolve import VariationalQBEvolve from struphy.propagators.variational_resistivity import VariationalResistivity + from struphy.propagators.variational_sound_evolve import VariationalSoundEvolve from struphy.propagators.variational_viscosity import VariationalViscosity from struphy.propagators.vlasov_ampere_coupling import VlasovAmpereCoupling from struphy.propagators.vlasov_maxwell_coupling import VlasovMaxwellCoupling @@ -144,6 +148,7 @@ "ShearAlfvenPropagator", "TimeDependentSource", "TwoFluidQuasiNeutralFull", + "VariationalAlfvenEvolve", "VariationalDensityEvolve", "VariationalEntropyEvolve", "VariationalMagFieldEvolve", @@ -151,6 +156,7 @@ "VariationalPBEvolve", "VariationalQBEvolve", "VariationalResistivity", + "VariationalSoundEvolve", "VariationalViscosity", "VlasovAmpereCoupling", "VlasovMaxwellCoupling", diff --git a/src/struphy/propagators/tests/test_variational_qb_evolve.py b/src/struphy/propagators/tests/test_variational_qb_evolve.py index 84a70d3f8..199e2d325 100644 --- a/src/struphy/propagators/tests/test_variational_qb_evolve.py +++ b/src/struphy/propagators/tests/test_variational_qb_evolve.py @@ -10,14 +10,16 @@ @pytest.mark.mpi_skip @pytest.mark.parametrize("linearize", [False, True]) @pytest.mark.parametrize("info", [False, True]) -def test_variational_qb_evolve_step(linearize: bool, info: bool): - """One step of VariationalQBEvolve on a tiny grid, with and without linearization and solver info.""" +@pytest.mark.parametrize("name", ["variat_qb", "variat_alfven", "variat_sound"]) +def test_variational_qb_evolve_step(linearize: bool, info: bool, name: str): + """One step of VariationalQBEvolve (unsplit model), VariationalAlfvenEvolve or VariationalSoundEvolve (split model) + on a tiny grid, with and without linearization and solver info.""" from struphy import DerhamOptions, FieldsBackground, Simulation, domains, equils, grids, perturbations from struphy.linear_algebra.solver import NonlinearSolverParameters from struphy.models import VariationalPerturbativeMHD_with_q - model = VariationalPerturbativeMHD_with_q(linear=True) + model = VariationalPerturbativeMHD_with_q(linear=True, split_sound=name != "variat_qb") sim = Simulation( model=model, domain=domains.Cuboid(), @@ -26,12 +28,15 @@ def test_variational_qb_evolve_step(linearize: bool, info: bool): derham_opts=DerhamOptions(degree=(1, 1, 1)), ) - for name in ("variat_dens", "variat_viscous", "variat_resist"): - prop = getattr(model.propagators, name) - prop.options = prop.Options(model="linear_q") - prop = model.propagators.variat_qb + for other in ("variat_dens", "variat_viscous", "variat_resist", "variat_sound"): + prop = getattr(model.propagators, other, None) + if prop is not None: + prop.options = prop.Options(model="linear_q") + if hasattr(model.propagators, "variat_alfven"): + model.propagators.variat_alfven.options = model.propagators.variat_alfven.Options(model="linear") + prop = getattr(model.propagators, name) prop.options = prop.Options( - model="linear_q", + model="linear" if name == "variat_alfven" else "linear_q", nonlin_solver=NonlinearSolverParameters(linearize=linearize, maxiter=5, info=info), ) diff --git a/src/struphy/propagators/tests/test_variational_schur_precond.py b/src/struphy/propagators/tests/test_variational_schur_precond.py index 729d51d88..61a6ce253 100644 --- a/src/struphy/propagators/tests/test_variational_schur_precond.py +++ b/src/struphy/propagators/tests/test_variational_schur_precond.py @@ -1,5 +1,5 @@ """The Kronecker preconditioner of the shear-Alfven Schur operator of the linear variational propagators -(``schur_precond="StiffnessPreconditioner"`` of VariationalPBEvolve and VariationalQBEvolve): it changes the +(``schur_precond="StiffnessPreconditioner"`` of VariationalPBEvolve and VariationalAlfvenEvolve): it changes the iteration, not the solution, and it falls back to the mass-matrix preconditioner where it does not apply.""" import pytest @@ -20,7 +20,6 @@ LINEAR_P = linear(VariationalPerturbativeMHD) LINEAR_Q_SPLIT = linear(VariationalPerturbativeMHD_with_q, split_sound=True) -LINEAR_Q = linear(VariationalPerturbativeMHD_with_q) def _run(tmp_path, model_class, schur_precond): @@ -42,7 +41,7 @@ def _run(tmp_path, model_class, schur_precond): @pytest.mark.parametrize( "model_class, name", - [(LINEAR_P, "variat_pb"), (LINEAR_Q_SPLIT, "variat_qb")], + [(LINEAR_P, "variat_pb"), (LINEAR_Q_SPLIT, "variat_alfven")], ids=["linear_p", "linear_q_split"], ) def test_same_solution(tmp_path, model_class, name): @@ -58,7 +57,11 @@ def test_same_solution(tmp_path, model_class, name): def test_fallback(tmp_path): - """Where the Schur operator is not the shear-Alfven one (the unsplit q model), the mass-matrix preconditioner is - used.""" - model = _run(tmp_path, LINEAR_Q, "StiffnessPreconditioner") - assert model.propagators.variat_qb._schur_pc is None + """Where the Schur operator is not the linear shear-Alfven one (the perturbative models), the mass-matrix + preconditioner is used.""" + for model_class, name in ( + (VariationalPerturbativeMHD, "variat_pb"), + (VariationalPerturbativeMHD_with_q, "variat_alfven"), + ): + model = _run(tmp_path, model_class, "StiffnessPreconditioner") + assert getattr(model.propagators, name)._schur_pc is None, name diff --git a/src/struphy/propagators/variational_alfven_evolve.py b/src/struphy/propagators/variational_alfven_evolve.py new file mode 100644 index 000000000..9040474d6 --- /dev/null +++ b/src/struphy/propagators/variational_alfven_evolve.py @@ -0,0 +1,450 @@ +import logging +from dataclasses import dataclass +from typing import Literal + +from feectools.linalg.basic import IdentityOperator +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace +from feectools.linalg.solvers import inverse +from line_profiler import profile + +from struphy.feec import preconditioner +from struphy.feec.preconditioner import MassMatrixPreconditioner +from struphy.feec.variational_utilities import ( + Hdiv0_transport_operator, + NewtonConvergence, + alfven_schur_preconditioner, + dot_normalized, + inv_norm2, +) +from struphy.io.options import LiteralOptions, OptionsBase +from struphy.linear_algebra.schur_solver import SchurSolverFull +from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.utils.utils import check_option + +logger = logging.getLogger("struphy") + + +class VariationalAlfvenEvolve(Propagator): + r""":ref:`FEEC ` discretization of the coupling of the velocity and the magnetic field (shear-Alfvén and + fast magnetosonic waves) of the variational MHD models: find :math:`\mathbf u \in (H^1)^3` and + :math:`\mathbf B \in H(\textnormal{div})` such that, with the transport operator + :math:`\mathcal T_{\bar{\mathbf B}}(\mathbf v) = \nabla \times (\bar{\mathbf B} \times \mathbf v)`, depending on the + option ``model``: + + * ``"full"`` (full MHD): + + .. math:: + + \int_\Omega \partial_t (\rho \mathbf u) \cdot \mathbf v\,\textrm d \mathbf x + - \int_\Omega \mathbf B \cdot \mathcal T_{\mathbf B}(\mathbf v)\,\textrm d \mathbf x = 0 + \qquad \forall \, \mathbf v \in (H^1)^3\,, + \qquad + \partial_t \mathbf B + \mathcal T_{\mathbf B}(\mathbf u) = 0 \,. + + * ``"linear"`` (:math:`\mathbf B` is the perturbation of the equilibrium :math:`\mathbf B_0`): + + .. math:: + + \int_\Omega \partial_t (\rho_0 \mathbf u) \cdot \mathbf v\,\textrm d \mathbf x + - \int_\Omega \left(\mathbf B \cdot \mathcal T_{\mathbf B_0}(\mathbf v) + + \mathbf B_0 \cdot \mathcal T_{\mathbf B}(\mathbf v)\right)\textrm d \mathbf x = 0\,, + \qquad + \partial_t \mathbf B + \mathcal T_{\mathbf B_0}(\mathbf u) = 0 \,. + + * ``"perturbative"`` (:math:`\mathbf B` is the perturbation, all nonlinear terms are kept): + + .. math:: + + \int_\Omega \partial_t (\rho \mathbf u) \cdot \mathbf v\,\textrm d \mathbf x + - \int_\Omega \left(\mathbf B \cdot \mathcal T_{\mathbf B_0 + \mathbf B}(\mathbf v) + + \mathbf B_0 \cdot \mathcal T_{\mathbf B}(\mathbf v)\right)\textrm d \mathbf x = 0\,, + \qquad + \partial_t \mathbf B + \mathcal T_{\mathbf B_0 + \mathbf B}(\mathbf u) = 0 \,. + + The thermodynamic variable is not changed, so this propagator serves the models in :math:`p` and in + :math:`q = \sqrt p` alike; together with :class:`~struphy.propagators.variational_sound_evolve.VariationalSoundEvolve` + it splits the coupled step of :class:`~struphy.propagators.variational_qb_evolve.VariationalQBEvolve` into the + shear-Alfvén/fast and the sound waves. The second-order term of the linear and perturbative models comes from the + auxiliary variable :math:`\hat{\mathbf B}` (``bt2``), which, if given, is advanced with + :math:`\partial_t \hat{\mathbf B} + \mathcal T_{\mathbf B}(\mathbf u) = 0`; it enters the conserved energy as + :math:`\int \mathbf B_0 \cdot \hat{\mathbf B}`. For ``"perturbative"``, :math:`\rho` is :math:`\rho_0` plus the + variable of :class:`~struphy.propagators.variational_density_evolve.VariationalDensityEvolve`. + + Discretization: with :math:`\mathbb T[\bar{\mathbf B}] = \mathbb C \hat\Pi^1[\hat{\bar{\mathbf B}} \times + \vec{\boldsymbol\Lambda}^v]` (:class:`~struphy.feec.variational_utilities.Hdiv0_transport_operator`) and, for + ``"linear"``, + + .. math:: + + \begin{align} + &\mathbb M^v[\hat\rho_0] \frac{\mathbf u^{n+1} - \mathbf u^n}{\Delta t} + - \mathbb T[\mathbf B_0]^\top \mathbb M^2 \mathbf b^{n+\frac12} + - \mathbb T[\bar{\mathbf b}]^\top \mathbb M^2 \mathbf b_0 = 0 \,, + \\[2mm] + &\frac{\mathbf b^{n+1} - \mathbf b^n}{\Delta t} + \mathbb T[\mathbf B_0]\, \mathbf u^{n+\frac12} = 0 \,, + \qquad + \frac{\hat{\mathbf b}^{n+1} - \hat{\mathbf b}^n}{\Delta t} + \mathbb T[\bar{\mathbf b}]\, \mathbf u^{n+\frac12} = 0 \,, + \end{align} + + and correspondingly for the other models. The weight :math:`\bar{\mathbf b}` of the transport operators is taken + at the same time in all equations, which makes the step energy conserving: at :math:`t^n` (first order in time, + one Newton step solves the linear system), or at the midpoint :math:`t^{n+1/2}` with the option + ``implicit_weights`` (second order in time, a Picard iteration). The Newton system is solved with + :class:`~struphy.linear_algebra.schur_solver.SchurSolverFull`; its Schur complement for ``"linear"`` is the + shear-Alfvén operator :math:`\mathbb M^v[\hat\rho_0] + \frac{\Delta t^2}{4} \mathbb T[\mathbf B_0]^\top \mathbb M^2 + \mathbb T[\mathbf B_0]`, for which the option ``schur_precond`` provides a Kronecker preconditioner. + """ + + class Variables: + """Container for variables advanced by :class:`VariationalAlfvenEvolve`. + + Attributes + ---------- + u : FEECVariable + Velocity variable in ``"H1vec"`` space. + b : FEECVariable + Magnetic-field variable in ``"Hdiv"`` space. + """ + + def __init__(self): + self._u: FEECVariable = None + self._b: FEECVariable = None + + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new + + @property + def b(self) -> FEECVariable: + return self._b + + @b.setter + def b(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "Hdiv" + self._b = new + + def __init__(self, bt2: FEECVariable = None): + r""" + Parameters + ---------- + bt2 : FEECVariable, default=None + Auxiliary second-order variable :math:`\hat{\mathbf B}` (``"Hdiv"`` space) of the linear and perturbative + models, advanced with :math:`\partial_t \hat{\mathbf B} + \mathcal T_{\mathbf B}(\mathbf u) = 0`. It enters + only the energy, not the momentum equation. + """ + self.variables = self.Variables() + self.bt2 = bt2 + + @dataclass(repr=False) + class Options(OptionsBase): + r"""Configuration options for :class:`VariationalAlfvenEvolve`. + + Parameters + ---------- + model : {"full", "linear", "perturbative"}, default="full" + Model variant. + solver : LiteralOptions.OptsSymmSolver, default="pcg" + Linear solver of the Schur system. + precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner" + Preconditioner used in linear solves. + solver_params : SolverParameters, default=None + Linear-solver controls. The Newton systems are solved for a normalized right-hand side, so ``tol`` is + relative to the norm of the residual. + nonlin_solver : NonlinearSolverParameters, default=None + Nonlinear/Picard iteration controls. The iteration stops when the norm of the residual is below ``tol`` + times the norm of the state at :math:`t^n` if that is smaller than one (absolute otherwise), or when it + stagnates (with a warning). + implicit_weights : bool, default=None + If ``True``, the weight of the transport operators is taken at the midpoint :math:`t^{n+1/2}` and updated + in every Picard iteration (second order in time, several iterations per step); if ``False``, it is taken + at :math:`t^n` (first order in time, one linear solve). Both conserve the energy. ``None`` means ``True`` + for ``"linear"`` and ``"perturbative"`` and ``False`` for ``"full"``. + schur_precond : {"MassMatrixPreconditioner", "StiffnessPreconditioner"}, default="MassMatrixPreconditioner" + Preconditioner of the Schur system: the mass-matrix preconditioner of the velocity block, or the Kronecker + (fast diagonalization) preconditioner of the whole shear-Alfvén Schur operator + (:func:`~struphy.feec.variational_utilities.alfven_schur_preconditioner`, built for each time step in + :meth:`prepare_time_step`). The latter needs ``model="linear"``, ``solver="pcg"``, tensor-product splines + and the global commuting projectors; otherwise the mass-matrix preconditioner is used (with a warning). + exb_correction : bool, default=True + For ``schur_precond="StiffnessPreconditioner"``: add the auxiliary space correction of the + :math:`\mathbf E \times \mathbf B` flows :math:`\mathbf B_0 \times \nabla \phi`. + """ + + # specific literals + OptsModel = Literal["full", "linear", "perturbative"] + OptsSchurPrecond = Literal["MassMatrixPreconditioner", "StiffnessPreconditioner"] + # propagator options + model: OptsModel = "full" + solver: LiteralOptions.OptsSymmSolver = "pcg" + precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + implicit_weights: bool | None = None + schur_precond: OptsSchurPrecond = "MassMatrixPreconditioner" + exb_correction: bool = True + + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, LiteralOptions.OptsSymmSolver) + check_option(self.precond, LiteralOptions.OptsMassPrecond) + check_option(self.schur_precond, self.OptsSchurPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters() + if self.implicit_weights is None: + self.implicit_weights = self.model in ("linear", "perturbative") + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + self._options = new + logger.info(f"\nNew options for propagator '{self.__class__.__name__}':\n{self._options}") + + @profile + def allocate(self): + self._model = self.options.model + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize = self.options.nonlin_solver.linearize + self._implicit_weights = self.options.implicit_weights + self._bt2 = None if self.bt2 is None else self.bt2.spline.vector + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._info = self._nonlin_solver.info and (self.rank == 0) + + # assembly of WMMnew happens in VariationalDensityEvolve + self._Mrho = self.mass_ops.WMMnew + self._Mrho_pc = MassMatrixPreconditioner(self._Mrho, dim_reduce=None) + + self._initialize_operators() + + # Kronecker preconditioner of the Schur operator, built for the time step in prepare_time_step + self._schur_pc = alfven_schur_preconditioner(self, [] if self._model == "linear" else ['model="linear"']) + self._schur_pc_dt = None + + # temporaries + u = self.variables.u.spline.vector + b = self.variables.b.spline.vector + self._tmp_un1 = u.space.zeros() + self._tmp_un12 = u.space.zeros() + self._tmp_un_diff = u.space.zeros() + self._tmp_mn = u.space.zeros() + self._tmp_mn1 = u.space.zeros() + self._tmp_mn_diff = u.space.zeros() + self._tmp_advection = u.space.zeros() + self._tmp_advection2 = u.space.zeros() + self._tmp_bn1 = b.space.zeros() + self._tmp_bn12 = b.space.zeros() + self._tmp_bn_diff = b.space.zeros() + self._tmp_b_advection = b.space.zeros() + self._tmp_b_advection2 = b.space.zeros() + self._linear_form_dl_db = b.space.zeros() + + if self._linearize: + self._extracted_b2 = self.derham.extraction_ops["2"].dot(self.projected_equil.b2) + + def prepare_time_step(self, dt: float): + """Kronecker preconditioner of the Schur operator for the time step ``dt`` (with + ``schur_precond="StiffnessPreconditioner"``): built in the setup, rebuilt when dt changes.""" + if self._schur_pc is not None and dt != self._schur_pc_dt: + self._inv_Jacobian.precond = self._schur_pc.build(dt) + self._schur_pc_dt = dt + + @profile + def __call__(self, dt): + # preconditioner for the current time step (no-op unless dt changed since the setup) + self.prepare_time_step(dt) + # WMMnew is re-assembled in VariationalDensityEvolve: update the diagonal scaling of its preconditioner + self._Mrho_pc.refresh() + + if self._nonlin_solver.type != "Picard": + raise ValueError("Only the Picard solver is implemented for VariationalAlfvenEvolve") + + if self._info: + logger.info("") + logger.info("Newton iteration in VariationalAlfvenEvolve") + + un = self.variables.u.spline.vector + bn = self.variables.b.spline.vector + + # with implicit weights the operators are updated at the midpoint inside the iteration + if not self._implicit_weights: + self._update_weights(bn) + + mn = self._Mrho.dot(un, out=self._tmp_mn) + bn1 = bn.copy(out=self._tmp_bn1) + bn1 += self._tmp_bn_diff + un1 = un.copy(out=self._tmp_un1) + un1 += self._tmp_un_diff + mn1 = self._Mrho.dot(un1, out=self._tmp_mn1) + tol = self._nonlin_solver.tol + # stopping criterion relative to the state at t^n, with a check for stagnation + conv = NewtonConvergence("VariationalAlfvenEvolve", tol, self._get_error(mn, bn)) + + for it in range(self._nonlin_solver.maxiter): + un12 = un.copy(out=self._tmp_un12) + un12 += un1 + un12 *= 0.5 + self._update_linear_form_dl_db(bn, bn1) + if self._implicit_weights: + self._update_weights(self._tmp_bn12) + + advection, b_advection = self._advection(un12) + advection *= dt + b_advection *= dt + + bn_diff = bn1.copy(out=self._tmp_bn_diff) + bn_diff -= bn + bn_diff += b_advection + mn_diff = mn1.copy(out=self._tmp_mn_diff) + mn_diff -= mn + mn_diff += advection + + stop = conv.stop(self._get_error(mn_diff, bn_diff), it) + if self._info: + logger.info(f"iteration : {it} relative error : {conv.err}") + if stop: + break + + # Newton step for the normalized residual (the tolerance of the linear solver is relative to it) + self._get_jacobian(dt) + self._tmp_f[0] = mn_diff + self._tmp_f[1] = bn_diff + incr = dot_normalized(self._inv_Jacobian, self._tmp_f, self._tmp_incr) + if self._info: + logger.info(f"information on the linear solver : {self._inv_Jacobian._solver._info}") + un1 -= incr[0] + bn1 -= incr[1] + mn1 = self._Mrho.dot(un1, out=self._tmp_mn1) + + conv.warn_if_not_converged(it, self._nonlin_solver.maxiter) + + if not conv.converged: + # the last iteration applied an update after the residual: midpoint velocity (and weights) of the final + # iterate, such that the update of the auxiliary variable below is consistent with it + un12 = un.copy(out=self._tmp_un12) + un12 += un1 + un12 *= 0.5 + if self._implicit_weights: + self._update_linear_form_dl_db(bn, bn1) + self._update_weights(self._tmp_bn12) + + self._tmp_un_diff = un1 - un + self._tmp_bn_diff = bn1 - bn + self.update_feec_variables(b=bn1, u=un1) + + # auxiliary second-order variable, with the same weight + if self._bt2 is not None: + b_advection = self.curlPib.dot(un12, out=self._tmp_b_advection) + b_advection *= dt + self._bt2 -= b_advection + + def _advection(self, un12): + """The force of the magnetic field on the velocity (without dt) and the induction term.""" + advection = self._tmp_advection + advection *= 0.0 + b_advection = self._tmp_b_advection + b_advection *= 0.0 + if self._model == "linear": + advection += self.curlPibT0.dot(self._linear_form_dl_db, out=self._tmp_advection2) + advection += self.curlPibT.dot(self._linear_form_dl_db0, out=self._tmp_advection2) + b_advection += self.curlPib0.dot(un12, out=self._tmp_b_advection2) + elif self._model == "perturbative": + advection += self.curlPibT0.dot(self._linear_form_dl_db, out=self._tmp_advection2) + advection += self.curlPibT.dot(self._linear_form_dl_db0, out=self._tmp_advection2) + advection += self.curlPibT.dot(self._linear_form_dl_db, out=self._tmp_advection2) + b_advection += self.curlPib.dot(un12, out=self._tmp_b_advection2) + b_advection += self.curlPib0.dot(un12, out=self._tmp_b_advection2) + else: + advection += self.curlPibT.dot(self._linear_form_dl_db, out=self._tmp_advection2) + b_advection += self.curlPib.dot(un12, out=self._tmp_b_advection2) + return advection, b_advection + + def _initialize_operators(self): + """Transport operators, the Jacobian of the Newton step and its Schur solver.""" + self.curlPib = Hdiv0_transport_operator(self.derham) + self.curlPibT = self.curlPib.T + + # inverse mass matrix for the error norm + self.pc_Mv = preconditioner.MassMatrixPreconditioner(self.mass_ops.Mv, dim_reduce=None) + self._inv_Mv = inverse(self.mass_ops.Mv, "pcg", pc=self.pc_Mv, tol=1e-10, maxiter=1000, verbose=False) + + if self._model in ("linear", "perturbative"): + self.curlPib0 = Hdiv0_transport_operator(self.derham) + self.curlPibT0 = self.curlPib0.T + self.curlPib0.update_coeffs(self.projected_equil.b2) + self.curlPibT0.update_coeffs(self.projected_equil.b2) + self._linear_form_dl_db0 = -self.mass_ops.M2.dot(self.projected_equil.b2) + + if self._model == "linear": + T, TT = self.curlPib0, self.curlPibT0 + elif self._model == "perturbative": + T, TT = self.curlPib0 + self.curlPib, self.curlPibT0 + self.curlPibT + else: + T, TT = self.curlPib, self.curlPibT + # dt-dependent scalars set in _get_jacobian + self._mdt2_TT_M2 = 2 * (TT @ self.mass_ops.M2) + self._dt2_T = 2 * T + + Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V2pol) + self._tmp_f = Jacs.zeros() + self._tmp_incr = Jacs.zeros() + self._Jacobian = BlockLinearOperator(Jacs, Jacs) + self._Jacobian[0, 0] = self._Mrho + self._Jacobian[0, 1] = self._mdt2_TT_M2 + self._Jacobian[1, 0] = self._dt2_T + self._Jacobian[1, 1] = IdentityOperator(self.derham.V2pol) + self._inv_Jacobian = SchurSolverFull( + self._Jacobian, + self.options.solver, + # fixed (linear) preconditioner: an inner iterative solve is expensive and not a linear operator + pc=self._Mrho_pc, + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, + recycle=True, + ) + + def _update_weights(self, b): + """Weights of the transport operator with the (perturbed) magnetic field.""" + self.curlPib.update_coeffs(b) + self.curlPibT.update_coeffs(b) + + def _update_linear_form_dl_db(self, bn, bn1): + """The linear form -M2 b^{n+1/2} (derivative of the magnetic energy).""" + bn12 = bn.copy(out=self._tmp_bn12) + bn12 += bn1 + bn12 *= 0.5 + if self._linearize: + wb = self.mass_ops.M2.dot(bn12 - self._extracted_b2, out=self._linear_form_dl_db) + else: + wb = self.mass_ops.M2.dot(bn12, out=self._linear_form_dl_db) + wb *= -1 + + def _get_error(self, mn_diff, bn_diff): + # the inverse mass matrix is applied to the normalized residual (relative solver tolerance) + err_u = inv_norm2(self._inv_Mv, self.derham.boundary_ops["v"].dot(mn_diff)) + err_b = self.mass_ops.M2.dot_inner(bn_diff, bn_diff) + return max(err_b, err_u) + + def _get_jacobian(self, dt): + self._mdt2_TT_M2._scalar = -dt / 2 + self._dt2_T._scalar = dt / 2 diff --git a/src/struphy/propagators/variational_qb_evolve.py b/src/struphy/propagators/variational_qb_evolve.py index 0e3e22f1d..13ff2caec 100644 --- a/src/struphy/propagators/variational_qb_evolve.py +++ b/src/struphy/propagators/variational_qb_evolve.py @@ -16,7 +16,6 @@ Hdiv0_transport_operator, NewtonConvergence, Pressure_transport_operator, - alfven_schur_preconditioner, dot_normalized, inv_norm2, ) @@ -29,10 +28,6 @@ logger = logging.getLogger("struphy") -# couplings to the velocity advanced by VariationalQBEvolve -OptsCoupling = Literal["bq", "b", "q"] - - class VariationalQBEvolve(Propagator): r""":ref:`FEEC ` discretization of the following equations: find :math:`\mathbf u \in (H^1)^3`, :math:`q \in L^2` and :math:`\mathbf B \in H(\textnormal{div})` such that, @@ -89,11 +84,6 @@ class VariationalQBEvolve(Propagator): For ``"full_q"`` the mass density :math:`\rho` is the variable of :class:`~struphy.propagators.variational_density_evolve.VariationalDensityEvolve`, for ``"perturbative_q"`` it is :math:`\rho_0` plus that variable. - With the constructor argument ``coupling="b"`` (``"q"``) only the terms of the magnetic field (of :math:`q`) are - advanced, with :math:`q` (the magnetic field) and its auxiliary variable fixed; the Newton system is then a - :math:`2\times 2` block system in :math:`(\mathbf u, \mathbf B)` (:math:`(\mathbf u, q)`). Two such substeps split - the shear-Alfven/fast waves from the sound waves, and each substep conserves the energy. - The scheme below is written for ``"full_q"``; for the other models the operators and linear forms are replaced by the sums above. Below, hats denote pullbacks to the logical domain (not the auxiliary variables). The weights of the transport operators are taken at the same time in the momentum, induction and pressure @@ -203,7 +193,6 @@ def __init__( u2: FEECVariable = None, qt3: FEECVariable = None, bt2: FEECVariable = None, - coupling: OptsCoupling = "bq", ): r""" Parameters @@ -220,22 +209,12 @@ def __init__( Auxiliary second-order variable :math:`\hat{\mathbf B}` (``"Hdiv"`` space) of the linear and perturbative models, advanced with :math:`\partial_t \hat{\mathbf B} + \mathcal T_{\mathbf B}(\mathbf u) = 0`. It enters only the energy, not the momentum equation. - coupling : {"bq", "b", "q"}, default="bq" - The couplings to the velocity advanced (implicitly) by this propagator: of the magnetic field and - :math:`q` together (``"bq"``), or only of the magnetic field (``"b"``) or only of :math:`q` (``"q"``), the - other variable (and its auxiliary variable) being kept fixed. Each choice conserves the energy, so two - instances with ``"b"`` and ``"q"`` give an energy-conserving splitting of the shear-Alfven/fast and the - sound waves, with a smaller linear system per substep. """ - check_option(coupling, OptsCoupling) self.variables = self.Variables() self.div_u = div_u self.u2 = u2 self.qt3 = qt3 self.bt2 = bt2 - self._coupling = coupling - self._with_b = coupling in ("bq", "b") - self._with_q = coupling in ("bq", "q") @dataclass(repr=False) class Options(OptionsBase): @@ -267,22 +246,9 @@ class Options(OptionsBase): does not contain the dependence of the operators on their weights, so a step needs several iterations (typically 2 to 10, more for large time steps or amplitudes) instead of one, each with a Schur solve; ``False`` is the cheaper first-order choice. - schur_precond : {"MassMatrixPreconditioner", "StiffnessPreconditioner"}, default="MassMatrixPreconditioner" - Preconditioner of the Schur system of the Newton step: the mass-matrix preconditioner of the velocity - block, or the Kronecker (fast diagonalization) preconditioner of the whole shear-Alfven Schur operator - :math:`\mathbb M^v[\rho_0] + \frac{\Delta t^2}{4} \mathbb T_0^\top \mathbb M^2 \mathbb T_0` - (:func:`~struphy.feec.variational_utilities.alfven_schur_preconditioner`, built from the equilibrium for - each time step in :meth:`prepare_time_step`). The latter needs ``model="linear_q"`` with ``coupling="b"`` (the Alfven substep of a split step), ``solver="pcg"``, - tensor-product splines and the global commuting projectors; otherwise the mass-matrix preconditioner is - used (with a warning). - exb_correction : bool, default=True - For ``schur_precond="StiffnessPreconditioner"``: add the auxiliary space correction of the - :math:`\mathbf E \times \mathbf B` flows :math:`\mathbf B_0 \times \nabla \phi` (the incompressible shear - Alfven modes with small parallel wave number, which block Jacobi does not reduce). """ # specific literals - OptsSchurPrecond = Literal["MassMatrixPreconditioner", "StiffnessPreconditioner"] OptsModel = Literal["full_q", "linear_q", "perturbative_q"] # propagator options model: OptsModel = "full_q" @@ -292,15 +258,12 @@ class Options(OptionsBase): solver_params: SolverParameters = None nonlin_solver: NonlinearSolverParameters = None implicit_weights: bool | None = None - schur_precond: OptsSchurPrecond = "MassMatrixPreconditioner" - exb_correction: bool = True def __post_init__(self): # checks check_option(self.model, self.OptsModel) check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsMassPrecond) - check_option(self.schur_precond, self.OptsSchurPrecond) # defaults if self.solver_params is None: @@ -363,14 +326,6 @@ def allocate(self): # Projector self._initialize_projectors_and_mass() - # Kronecker preconditioner of the Schur operator, built for the time step in prepare_time_step - self._schur_pc = alfven_schur_preconditioner( - self, - ([] if self._model == "linear_q" else ['model="linear_q"']) - + ([] if self._coupling == "b" else ['coupling="b"']), - ) - self._schur_pc_dt = None - # bunch of temporaries to avoid allocating in the loop u = self.variables.u.spline.vector q = self.variables.q.spline.vector @@ -406,16 +361,7 @@ def allocate(self): self._extracted_b2 = self.derham.extraction_ops["2"].dot(self.projected_equil.b2) self._extracted_q3 = self.derham.extraction_ops["3"].dot(self.projected_equil.q3) - def prepare_time_step(self, dt: float): - """Kronecker preconditioner of the Schur operator for the time step ``dt`` (with - ``schur_precond="StiffnessPreconditioner"``): built in the setup, rebuilt when dt changes.""" - if self._schur_pc is not None and dt != self._schur_pc_dt: - self._inv_Jacobian.precond = self._schur_pc.build(dt) - self._schur_pc_dt = dt - def __call__(self, dt): - # preconditioner for the current time step (no-op unless dt changed since the setup) - self.prepare_time_step(dt) # WMMnew is re-assembled in VariationalDensityEvolve: update the diagonal scaling of its preconditioner self._Mrho_pc.refresh() @@ -470,45 +416,88 @@ def __call_picard(self, dt): self._update_Pib(self._tmp_bn12) self._update_Projq(self._tmp_qn12) - # Compute the advection terms of the couplings advanced here (the other variable is kept fixed) - advection = self._tmp_advection - advection *= 0.0 - b_advection = self._tmp_b_advection - b_advection *= 0.0 - q_advection = self._tmp_q_advection - q_advection *= 0.0 - + # Compute the advection terms if self._model == "linear_q": - if self._with_b: - advection += self.curlPibT0.dot(self._linear_form_dl_db, out=self._tmp_advection2) - advection += self.curlPibT.dot(self._linear_form_dl_db0, out=self._tmp_advection2) - b_advection += self.curlPib0.dot(un12, out=self._tmp_b_advection2) - if self._with_q: - advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) - q_advection += self._transop_q0.dot(un12, out=self._tmp_q_advection2) + advection = self.curlPibT0.dot( + self._linear_form_dl_db, + out=self._tmp_advection, + ) + + advection += self.curlPibT.dot( + self._linear_form_dl_db0, + out=self._tmp_advection2, + ) + + advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) + + b_advection = self.curlPib0.dot( + un12, + out=self._tmp_b_advection, + ) + + q_advection = self._transop_q0.dot( + un12, + out=self._tmp_q_advection, + ) elif self._model == "perturbative_q": - if self._with_b: - advection += self.curlPibT0.dot(self._linear_form_dl_db, out=self._tmp_advection2) - advection += self.curlPibT.dot(self._linear_form_dl_db0, out=self._tmp_advection2) - advection += self.curlPibT.dot(self._linear_form_dl_db, out=self._tmp_advection2) - b_advection += self.curlPib.dot(un12, out=self._tmp_b_advection2) - b_advection += self.curlPib0.dot(un12, out=self._tmp_b_advection2) - if self._with_q: - advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) - q_advection += self._transop_q0.dot(un12, out=self._tmp_q_advection2) - q_advection += self._transop_q.dot(un12, out=self._tmp_q_advection2) + advection = self.curlPibT0.dot( + self._linear_form_dl_db, + out=self._tmp_advection, + ) + + advection += self.curlPibT.dot( + self._linear_form_dl_db0, + out=self._tmp_advection2, + ) + + advection += self.curlPibT.dot( + self._linear_form_dl_db, + out=self._tmp_advection2, + ) + + advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) + + b_advection = self.curlPib.dot( + un12, + out=self._tmp_b_advection, + ) + + b_advection += self.curlPib0.dot( + un12, + out=self._tmp_b_advection2, + ) + + q_advection = self._transop_q0.dot( + un12, + out=self._tmp_q_advection, + ) + + q_advection += self._transop_q.dot( + un12, + out=self._tmp_q_advection2, + ) else: - if self._with_b: - advection += self.curlPibT.dot(self._linear_form_dl_db, out=self._tmp_advection2) - b_advection += self.curlPib.dot(un12, out=self._tmp_b_advection2) - if self._with_q: - advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) - q_advection += self._transop_q.dot(un12, out=self._tmp_q_advection2) + advection = self.curlPibT.dot( + self._linear_form_dl_db, + out=self._tmp_advection, + ) + + advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + + b_advection = self.curlPib.dot( + un12, + out=self._tmp_b_advection, + ) + + q_advection = self._transop_q.dot( + un12, + out=self._tmp_q_advection, + ) advection *= dt b_advection *= dt @@ -541,24 +530,16 @@ def __call_picard(self, dt): # Newton step self._tmp_f[0] = mn_diff - if self._coupling == "bq": - self._tmp_f[1] = bn_diff - self._tmp_f[2] = qn_diff - else: - self._tmp_f[1] = bn_diff if self._with_b else qn_diff + self._tmp_f[1] = bn_diff + self._tmp_f[2] = qn_diff # normalized right-hand side: the tolerance of the linear solver is relative to the residual incr = dot_normalized(self._inv_Jacobian, self._tmp_f, self._tmp_incr) if self._info: logger.info(f"information on the linear solver : {self._inv_Jacobian._solver._info}") un1 -= incr[0] - if self._coupling == "bq": - bn1 -= incr[1] - qn1 -= incr[2] - elif self._with_b: - bn1 -= incr[1] - else: - qn1 -= incr[1] + bn1 -= incr[1] + qn1 -= incr[2] # Multiply by the mass matrix to get the momentum mn1 = self._Mrho.dot(un1, out=self._tmp_mn1) @@ -587,7 +568,7 @@ def __call_picard(self, dt): # Update the 2nd order variables - if self._qt3 is not None and self._with_q: + if self._qt3 is not None: q_advection = self._transop_q.dot( un12, out=self._tmp_q_advection, @@ -595,7 +576,7 @@ def __call_picard(self, dt): q_advection *= dt self._qt3 -= q_advection - if self._bt2 is not None and self._with_b: + if self._bt2 is not None: b_advection = self.curlPib.dot( un12, out=self._tmp_b_advection, @@ -640,13 +621,11 @@ def _initialize_projectors_and_mass(self): self._I2 = IdentityOperator(self.derham.V2pol) self._I3 = IdentityOperator(self.derham.V3pol) - # unknowns of the Newton system: u and the coupled variables - if self._coupling == "bq": - Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V2pol, self.derham.V3pol) - elif self._with_b: - Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V2pol) - else: - Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V3pol) + Jacs = BlockVectorSpace( + self.derham.Vvpol, + self.derham.V2pol, + self.derham.V3pol, + ) self._tmp_f = Jacs.zeros() self._tmp_incr = Jacs.zeros() @@ -703,29 +682,17 @@ def _initialize_projectors_and_mass(self): # local version to avoid creating new version of LinearOperator every time - from struphy.linear_algebra.schur_solver import SchurSolverFull, SchurSolverFull3 - self._Jacobian[0, 0] = self._Mrho - if self._coupling == "bq": - self._Jacobian[0, 1] = self._mdt2_pc_curlPibT_M - self._Jacobian[0, 2] = self._mdt2_pc_transopT_M - self._Jacobian[1, 0] = self._dt2_curlPib - self._Jacobian[1, 1] = self._I2 - self._Jacobian[2, 0] = self._dt2_transop - self._Jacobian[2, 2] = self._I3 - schur_class = SchurSolverFull3 - elif self._with_b: - self._Jacobian[0, 1] = self._mdt2_pc_curlPibT_M - self._Jacobian[1, 0] = self._dt2_curlPib - self._Jacobian[1, 1] = self._I2 - schur_class = SchurSolverFull - else: - self._Jacobian[0, 1] = self._mdt2_pc_transopT_M - self._Jacobian[1, 0] = self._dt2_transop - self._Jacobian[1, 1] = self._I3 - schur_class = SchurSolverFull + self._Jacobian[0, 1] = self._mdt2_pc_curlPibT_M + self._Jacobian[0, 2] = self._mdt2_pc_transopT_M + self._Jacobian[1, 0] = self._dt2_curlPib + self._Jacobian[1, 1] = self._I2 + self._Jacobian[2, 0] = self._dt2_transop + self._Jacobian[2, 2] = self._I3 + + from struphy.linear_algebra.schur_solver import SchurSolverFull3 - self._inv_Jacobian = schur_class( + self._inv_Jacobian = SchurSolverFull3( self._Jacobian, self.options.solver, # fixed (linear) preconditioner: an inner iterative solve is expensive and not a linear operator diff --git a/src/struphy/propagators/variational_sound_evolve.py b/src/struphy/propagators/variational_sound_evolve.py new file mode 100644 index 000000000..1359d1e36 --- /dev/null +++ b/src/struphy/propagators/variational_sound_evolve.py @@ -0,0 +1,441 @@ +import logging +from dataclasses import dataclass +from typing import Literal + +from feectools.linalg.basic import IdentityOperator +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace +from feectools.linalg.solvers import inverse +from line_profiler import profile + +from struphy.feec import preconditioner +from struphy.feec.preconditioner import MassMatrixPreconditioner +from struphy.feec.variational_utilities import ( + NewtonConvergence, + Pressure_transport_operator, + dot_normalized, + inv_norm2, +) +from struphy.io.options import LiteralOptions, OptionsBase +from struphy.linear_algebra.schur_solver import SchurSolverFull +from struphy.linear_algebra.solver import NonlinearSolverParameters, SolverParameters +from struphy.models.variables import FEECVariable +from struphy.propagators.base import Propagator +from struphy.utils.utils import check_option + +logger = logging.getLogger("struphy") + + +class VariationalSoundEvolve(Propagator): + r""":ref:`FEEC ` discretization of the coupling of the velocity and :math:`q = \sqrt p` (sound waves) of the + variational MHD models in :math:`q`: find :math:`\mathbf u \in (H^1)^3` and :math:`q \in L^2` such that, with the + transport operator :math:`\mathcal S_{\bar q}(\mathbf v) = \nabla \cdot (\bar q \mathbf v) + \left(\frac \gamma 2 - + 1\right) \bar q\, \nabla \cdot \mathbf v`, depending on the option ``model``: + + * ``"full_q"`` (full MHD): + + .. math:: + + \int_\Omega \partial_t (\rho \mathbf u) \cdot \mathbf v\,\textrm d \mathbf x + - \int_\Omega \frac{2 q}{\gamma -1} \mathcal S_q(\mathbf v)\,\textrm d \mathbf x = 0 + \qquad \forall \, \mathbf v \in (H^1)^3\,, + \qquad + \partial_t q + \mathcal S_q(\mathbf u) = 0 \,. + + * ``"linear_q"`` (:math:`q` is the perturbation of the equilibrium :math:`q_0`): + + .. math:: + + \int_\Omega \partial_t (\rho_0 \mathbf u) \cdot \mathbf v\,\textrm d \mathbf x + - \int_\Omega \frac{2}{\gamma -1} \left(q\, \mathcal S_{q_0}(\mathbf v) + + q_0\, \mathcal S_q(\mathbf v)\right)\textrm d \mathbf x = 0\,, + \qquad + \partial_t q + \mathcal S_{q_0}(\mathbf u) = 0 \,. + + * ``"perturbative_q"`` (:math:`q` is the perturbation, all nonlinear terms are kept): + + .. math:: + + \int_\Omega \partial_t (\rho \mathbf u) \cdot \mathbf v\,\textrm d \mathbf x + - \int_\Omega \frac{2}{\gamma -1} \left(q\, \mathcal S_{q_0 + q}(\mathbf v) + + q_0\, \mathcal S_q(\mathbf v)\right)\textrm d \mathbf x = 0\,, + \qquad + \partial_t q + \mathcal S_{q_0 + q}(\mathbf u) = 0 \,. + + The magnetic field is not changed; together with + :class:`~struphy.propagators.variational_alfven_evolve.VariationalAlfvenEvolve` this splits the coupled step of + :class:`~struphy.propagators.variational_qb_evolve.VariationalQBEvolve` into the shear-Alfvén/fast and the sound + waves. The second-order term of the linear and perturbative models comes from the auxiliary variable + :math:`\hat q` (``qt3``), which, if given, is advanced with :math:`\partial_t \hat q + \mathcal S_q(\mathbf u) = 0`; + it enters the conserved energy as :math:`\int 2 q_0 \hat q / (\gamma - 1)`. For ``"perturbative_q"``, :math:`\rho` + is :math:`\rho_0` plus the variable of + :class:`~struphy.propagators.variational_density_evolve.VariationalDensityEvolve`. + + Discretization: with :math:`\mathbb S[\bar q] = \mathbb D \hat\Pi^2[\hat{\bar q} \vec{\boldsymbol\Lambda}^v] + + (\gamma/2 - 1) \hat\Pi^3[\hat{\bar q}/\sqrt g\, \vec{\boldsymbol\Lambda}^3] \mathbb D \mathcal U^v` + (:class:`~struphy.feec.variational_utilities.Pressure_transport_operator`) and, for ``"linear_q"``, + + .. math:: + + \begin{align} + &\mathbb M^v[\hat\rho_0] \frac{\mathbf u^{n+1} - \mathbf u^n}{\Delta t} + - \frac{2}{\gamma - 1}\left(\mathbb S[q_0]^\top \mathbb M^3 \mathbf q^{n+\frac12} + + \mathbb S[\bar{\mathbf q}]^\top \mathbb M^3 \mathbf q_0\right) = 0 \,, + \\[2mm] + &\frac{\mathbf q^{n+1} - \mathbf q^n}{\Delta t} + \mathbb S[q_0]\, \mathbf u^{n+\frac12} = 0 \,, + \qquad + \frac{\hat{\mathbf q}^{n+1} - \hat{\mathbf q}^n}{\Delta t} + \mathbb S[\bar{\mathbf q}]\, \mathbf u^{n+\frac12} = 0 \,, + \end{align} + + and correspondingly for the other models. The weight :math:`\bar{\mathbf q}` is taken at the same time in all + equations, which makes the step energy conserving: at :math:`t^n` (first order in time, one Newton step solves the + linear system), or at the midpoint :math:`t^{n+1/2}` with the option ``implicit_weights`` (second order in time, a + Picard iteration). The Newton system is solved with + :class:`~struphy.linear_algebra.schur_solver.SchurSolverFull`; for ``"linear_q"`` its Schur complement is + :math:`\mathbb M^v[\hat\rho_0] + \frac{\Delta t^2}{2(\gamma - 1)} \mathbb S[q_0]^\top \mathbb M^3 \mathbb S[q_0]`. + """ + + class Variables: + """Container for variables advanced by :class:`VariationalSoundEvolve`. + + Attributes + ---------- + u : FEECVariable + Velocity variable in ``"H1vec"`` space. + q : FEECVariable + Square root of the pressure in ``"L2"`` space. + """ + + def __init__(self): + self._u: FEECVariable = None + self._q: FEECVariable = None + + @property + def u(self) -> FEECVariable: + return self._u + + @u.setter + def u(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "H1vec" + self._u = new + + @property + def q(self) -> FEECVariable: + return self._q + + @q.setter + def q(self, new): + assert isinstance(new, FEECVariable) + assert new.space == "L2" + self._q = new + + def __init__(self, div_u: FEECVariable = None, u2: FEECVariable = None, qt3: FEECVariable = None): + r""" + Parameters + ---------- + div_u : FEECVariable, default=None + Divergence diagnostic (``"L2"`` space) of the velocity at the midpoint. + u2 : FEECVariable, default=None + Velocity as 2-form, :math:`\mathcal U^v \mathbf u` (``"Hdiv"`` space), diagnostic. + qt3 : FEECVariable, default=None + Auxiliary second-order variable :math:`\hat q` (``"L2"`` space) of the linear and perturbative models, + advanced with :math:`\partial_t \hat q + \mathcal S_q(\mathbf u) = 0`. It enters only the energy, not the + momentum equation. + """ + self.variables = self.Variables() + self.div_u = div_u + self.u2 = u2 + self.qt3 = qt3 + + @dataclass(repr=False) + class Options(OptionsBase): + r"""Configuration options for :class:`VariationalSoundEvolve`. + + Parameters + ---------- + model : {"full_q", "linear_q", "perturbative_q"}, default="full_q" + Model variant. + gamma : float, default=5/3 + Adiabatic index. + solver : LiteralOptions.OptsSymmSolver, default="pcg" + Linear solver of the Schur system. + precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner" + Preconditioner used in linear solves. + solver_params : SolverParameters, default=None + Linear-solver controls. The Newton systems are solved for a normalized right-hand side, so ``tol`` is + relative to the norm of the residual. + nonlin_solver : NonlinearSolverParameters, default=None + Nonlinear/Picard iteration controls. The iteration stops when the norm of the residual is below ``tol`` + times the norm of the state at :math:`t^n` if that is smaller than one (absolute otherwise), or when it + stagnates (with a warning). + implicit_weights : bool, default=None + If ``True``, the weight of the transport operators is taken at the midpoint :math:`t^{n+1/2}` and updated + in every Picard iteration (second order in time, several iterations per step); if ``False``, it is taken + at :math:`t^n` (first order in time, one linear solve). Both conserve the energy. ``None`` means ``True`` + for ``"linear_q"`` and ``"perturbative_q"`` and ``False`` for ``"full_q"``. + """ + + # specific literals + OptsModel = Literal["full_q", "linear_q", "perturbative_q"] + # propagator options + model: OptsModel = "full_q" + gamma: float = 5.0 / 3.0 + solver: LiteralOptions.OptsSymmSolver = "pcg" + precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" + solver_params: SolverParameters = None + nonlin_solver: NonlinearSolverParameters = None + implicit_weights: bool | None = None + + def __post_init__(self): + # checks + check_option(self.model, self.OptsModel) + check_option(self.solver, LiteralOptions.OptsSymmSolver) + check_option(self.precond, LiteralOptions.OptsMassPrecond) + + # defaults + if self.solver_params is None: + self.solver_params = SolverParameters() + if self.nonlin_solver is None: + self.nonlin_solver = NonlinearSolverParameters() + if self.implicit_weights is None: + self.implicit_weights = self.model in ("linear_q", "perturbative_q") + + @property + def options(self) -> Options: + if not hasattr(self, "_options"): + self._options = self.Options() + return self._options + + @options.setter + def options(self, new): + assert isinstance(new, self.Options) + self._options = new + logger.info(f"\nNew options for propagator '{self.__class__.__name__}':\n{self._options}") + + @profile + def allocate(self): + self._model = self.options.model + self._lin_solver = self.options.solver_params + self._nonlin_solver = self.options.nonlin_solver + self._linearize = self.options.nonlin_solver.linearize + self._gamma = self.options.gamma + self._implicit_weights = self.options.implicit_weights + self._divu = None if self.div_u is None else self.div_u.spline.vector + self._u2 = None if self.u2 is None else self.u2.spline.vector + self._qt3 = None if self.qt3 is None else self.qt3.spline.vector + self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 + self._info = self._nonlin_solver.info and (self.rank == 0) + + # assembly of WMMnew happens in VariationalDensityEvolve + self._Mrho = self.mass_ops.WMMnew + self._Mrho_pc = MassMatrixPreconditioner(self._Mrho, dim_reduce=None) + + self._initialize_operators() + + # temporaries + u = self.variables.u.spline.vector + q = self.variables.q.spline.vector + self._tmp_un1 = u.space.zeros() + self._tmp_un12 = u.space.zeros() + self._tmp_un_diff = u.space.zeros() + self._tmp_mn = u.space.zeros() + self._tmp_mn1 = u.space.zeros() + self._tmp_mn_diff = u.space.zeros() + self._tmp_advection = u.space.zeros() + self._tmp_advection2 = u.space.zeros() + self._tmp_qn1 = q.space.zeros() + self._tmp_qn12 = q.space.zeros() + self._tmp_qn_diff = q.space.zeros() + self._tmp_q_advection = q.space.zeros() + self._tmp_q_advection2 = q.space.zeros() + self._linear_form_dl_dq = q.space.zeros() + + if self._linearize: + self._extracted_q3 = self.derham.extraction_ops["3"].dot(self.projected_equil.q3) + + @profile + def __call__(self, dt): + # WMMnew is re-assembled in VariationalDensityEvolve: update the diagonal scaling of its preconditioner + self._Mrho_pc.refresh() + + if self._nonlin_solver.type != "Picard": + raise ValueError("Only the Picard solver is implemented for VariationalSoundEvolve") + + if self._info: + logger.info("") + logger.info("Newton iteration in VariationalSoundEvolve") + + un = self.variables.u.spline.vector + qn = self.variables.q.spline.vector + + # with implicit weights the operators are updated at the midpoint inside the iteration + if not self._implicit_weights: + self._update_weights(qn) + + mn = self._Mrho.dot(un, out=self._tmp_mn) + qn1 = qn.copy(out=self._tmp_qn1) + qn1 += self._tmp_qn_diff + un1 = un.copy(out=self._tmp_un1) + un1 += self._tmp_un_diff + mn1 = self._Mrho.dot(un1, out=self._tmp_mn1) + tol = self._nonlin_solver.tol + # stopping criterion relative to the state at t^n, with a check for stagnation + conv = NewtonConvergence("VariationalSoundEvolve", tol, self._get_error(mn, qn)) + + for it in range(self._nonlin_solver.maxiter): + un12 = un.copy(out=self._tmp_un12) + un12 += un1 + un12 *= 0.5 + self._update_linear_form_dl_dq(qn, qn1) + if self._implicit_weights: + self._update_weights(self._tmp_qn12) + + advection, q_advection = self._advection(un12) + advection *= dt + q_advection *= dt + + qn_diff = qn1.copy(out=self._tmp_qn_diff) + qn_diff -= qn + qn_diff += q_advection + mn_diff = mn1.copy(out=self._tmp_mn_diff) + mn_diff -= mn + mn_diff += advection + + stop = conv.stop(self._get_error(mn_diff, qn_diff), it) + if self._info: + logger.info(f"iteration : {it} relative error : {conv.err}") + if stop: + break + + # Newton step for the normalized residual (the tolerance of the linear solver is relative to it) + self._get_jacobian(dt) + self._tmp_f[0] = mn_diff + self._tmp_f[1] = qn_diff + incr = dot_normalized(self._inv_Jacobian, self._tmp_f, self._tmp_incr) + if self._info: + logger.info(f"information on the linear solver : {self._inv_Jacobian._solver._info}") + un1 -= incr[0] + qn1 -= incr[1] + mn1 = self._Mrho.dot(un1, out=self._tmp_mn1) + + conv.warn_if_not_converged(it, self._nonlin_solver.maxiter) + + if not conv.converged: + # the last iteration applied an update after the residual: midpoint velocity (and weights) of the final + # iterate, such that the update of the auxiliary variable below is consistent with it + un12 = un.copy(out=self._tmp_un12) + un12 += un1 + un12 *= 0.5 + if self._implicit_weights: + self._update_linear_form_dl_dq(qn, qn1) + self._update_weights(self._tmp_qn12) + + self._tmp_un_diff = un1 - un + self._tmp_qn_diff = qn1 - qn + self.update_feec_variables(q=qn1, u=un1) + + # diagnostics + if self._divu is not None: + self._transop_q.div.dot(un12, out=self._divu) + if self._u2 is not None: + self._transop_q._Uv.dot(un1, out=self._u2) + + # auxiliary second-order variable, with the same weight + if self._qt3 is not None: + q_advection = self._transop_q.dot(un12, out=self._tmp_q_advection) + q_advection *= dt + self._qt3 -= q_advection + + def _advection(self, un12): + """The force of q on the velocity (without dt) and the transport of q.""" + advection = self._tmp_advection + advection *= 0.0 + q_advection = self._tmp_q_advection + q_advection *= 0.0 + if self._model == "linear_q": + advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) + q_advection += self._transop_q0.dot(un12, out=self._tmp_q_advection2) + elif self._model == "perturbative_q": + advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_q0T.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + advection += self._transop_qT.dot(self._linear_form_dl_dq0, out=self._tmp_advection2) + q_advection += self._transop_q0.dot(un12, out=self._tmp_q_advection2) + q_advection += self._transop_q.dot(un12, out=self._tmp_q_advection2) + else: + advection += self._transop_qT.dot(self._linear_form_dl_dq, out=self._tmp_advection2) + q_advection += self._transop_q.dot(un12, out=self._tmp_q_advection2) + return advection, q_advection + + def _initialize_operators(self): + """Transport operators, the Jacobian of the Newton step and its Schur solver.""" + self._transop_q = Pressure_transport_operator(self.derham, self.domain, self.basis_ops.Uv, self._gamma / 2.0) + self._transop_qT = self._transop_q.T + + # inverse mass matrix for the error norm + self.pc_Mv = preconditioner.MassMatrixPreconditioner(self.mass_ops.Mv, dim_reduce=None) + self._inv_Mv = inverse(self.mass_ops.Mv, "pcg", pc=self.pc_Mv, tol=1e-10, maxiter=1000, verbose=False) + + if self._model in ("linear_q", "perturbative_q"): + self._transop_q0 = Pressure_transport_operator( + self.derham, self.domain, self.basis_ops.Uv, self._gamma / 2.0 + ) + self._transop_q0T = self._transop_q0.T + self._transop_q0.update_coeffs(self.projected_equil.q3) + self._transop_q0T.update_coeffs(self.projected_equil.q3) + self._linear_form_dl_dq0 = -2 / (self._gamma - 1.0) * self.mass_ops.M3.dot(self.projected_equil.q3) + + if self._model == "linear_q": + S, ST = self._transop_q0, self._transop_q0T + elif self._model == "perturbative_q": + S, ST = self._transop_q0 + self._transop_q, self._transop_q0T + self._transop_qT + else: + S, ST = self._transop_q, self._transop_qT + # dt-dependent scalars set in _get_jacobian + self._mdt_ST_M3 = 2 * (ST @ self.mass_ops.M3) + self._dt2_S = 2 * S + + Jacs = BlockVectorSpace(self.derham.Vvpol, self.derham.V3pol) + self._tmp_f = Jacs.zeros() + self._tmp_incr = Jacs.zeros() + self._Jacobian = BlockLinearOperator(Jacs, Jacs) + self._Jacobian[0, 0] = self._Mrho + self._Jacobian[0, 1] = self._mdt_ST_M3 + self._Jacobian[1, 0] = self._dt2_S + self._Jacobian[1, 1] = IdentityOperator(self.derham.V3pol) + self._inv_Jacobian = SchurSolverFull( + self._Jacobian, + self.options.solver, + # fixed (linear) preconditioner: an inner iterative solve is expensive and not a linear operator + pc=self._Mrho_pc, + tol=self._lin_solver.tol, + maxiter=self._lin_solver.maxiter, + verbose=self._lin_solver.verbose, + recycle=True, + ) + + def _update_weights(self, q): + """Weights of the transport operator with the (perturbed) q.""" + self._transop_q.update_coeffs(q) + self._transop_qT.update_coeffs(q) + + def _update_linear_form_dl_dq(self, qn, qn1): + """The linear form -2/(gamma-1) M3 q^{n+1/2} (derivative of the thermal energy).""" + qn12 = qn.copy(out=self._tmp_qn12) + qn12 += qn1 + qn12 *= 0.5 + if self._linearize: + wq = self.mass_ops.M3.dot(qn12 - self._extracted_q3, out=self._linear_form_dl_dq) + else: + wq = self.mass_ops.M3.dot(qn12, out=self._linear_form_dl_dq) + wq *= -2 / (self._gamma - 1) + + def _get_error(self, mn_diff, qn_diff): + # the inverse mass matrix is applied to the normalized residual (relative solver tolerance) + err_u = inv_norm2(self._inv_Mv, self.derham.boundary_ops["v"].dot(mn_diff)) + err_q = self.mass_ops.M3.dot_inner(qn_diff, qn_diff) + return max(err_q, err_u) + + def _get_jacobian(self, dt): + self._mdt_ST_M3._scalar = -dt / (self._gamma - 1.0) + self._dt2_S._scalar = dt / 2