diff --git a/doc/sections/subsections/propagators-avail.rst b/doc/sections/subsections/propagators-avail.rst index 37c30958c..d4b982187 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 c6e695dc4..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,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). 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 @@ -135,6 +137,7 @@ base_units = BaseUnits() # Model instance: linear and ideal (no viscosity, no resistivity), as 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 ) @@ -194,11 +197,15 @@ # ------------------ # 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.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/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index ce27a6eb6..95e82329f 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,14 @@ 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, 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. + References ---------- R. E. Lynch, J. R. Rice and D. H. Thomas, Direct solution of partial difference equations by tensor product @@ -769,8 +777,10 @@ def __init__( mid: str | LinearOperator | None = None, stab: str | LinearOperator | None = None, stab_average: str | None = None, + block_shapes: bool = False, ): 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 +1069,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 +1086,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 +1097,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 +1113,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) @@ -1198,6 +1210,7 @@ def _polar_stiffness_preconditioner( mid: str | LinearOperator | None = None, stab: str | LinearOperator | None = None, stab_average: str | None = None, + block_shapes: bool = False, **polar_kwargs, ) -> PolarStiffnessPreconditioner: """The PolarStiffnessPreconditioner for the arguments of StiffnessPreconditioner (plus its own options).""" @@ -1876,12 +1889,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`. @@ -1902,20 +1918,33 @@ 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. + + 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 | 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 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 +1952,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 @@ -1972,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( @@ -1980,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 @@ -2533,6 +2577,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 +2638,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..aebd1b1b3 100644 --- a/src/struphy/feec/variational_utilities.py +++ b/src/struphy/feec/variational_utilities.py @@ -1877,3 +1877,72 @@ 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 :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. + + 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 + # 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, + stab, + 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/default_params/test_models.py b/src/struphy/models/tests/default_params/test_models.py index 81b0e3bc0..667fc48cf 100644 --- a/src/struphy/models/tests/default_params/test_models.py +++ b/src/struphy/models/tests/default_params/test_models.py @@ -115,6 +115,17 @@ 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 the unsplit step (split_sound=False): the flag is kept, and VariationalQBEvolve gets the + # linear variant + if perturbative_model == "VariationalPerturbativeMHD_with_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 bce1cf4b2..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 @@ -22,9 +22,17 @@ 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 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 @@ -53,16 +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), @@ -86,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): @@ -121,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 @@ -151,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)] @@ -184,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) @@ -243,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 @@ -292,3 +305,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..5e10e220c 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, 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) @@ -76,12 +87,26 @@ def allocate_ideal(tmp_path, name, model_class, setup, dt, perturb=None, implici 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( - model=variant, solver_params=SOLVER, nonlin_solver=NONLIN + 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( + 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, + ) + 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 701ddd70f..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 @@ -38,6 +40,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 (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 @@ -79,11 +86,17 @@ def __init__( linear: bool = False, with_viscosity: bool = True, with_resistivity: bool = True, + split_sound: bool = True, ): 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_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: self.variat_viscous = VariationalViscosity(rho=rho, pt3=qt3) if with_resistivity: @@ -98,6 +111,7 @@ def __init__( linear: bool = False, with_viscosity: bool = True, with_resistivity: bool = True, + split_sound: bool = True, ): # 0. store input parameters @@ -121,6 +135,7 @@ def __init__( linear=linear, with_viscosity=with_viscosity, with_resistivity=with_resistivity, + split_sound=split_sound, ) # 4. assign variables to propagators @@ -128,9 +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 + 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 @@ -163,9 +184,21 @@ 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_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" - return {name: variant for name in ("variat_dens", "variat_qb", "variat_viscous", "variat_resist")} + 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,39 +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": 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: - if line.startswith("model = ") and linear: - new_file += [line.replace("(base_units=base_units)", "(base_units=base_units, linear=True)")] - 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_viscous.Options" in line: + 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 name is not None: 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"] @@ -345,16 +370,19 @@ 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) - 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. + :attr:`linear` is True (``"perturbative"`` or ``"linear"`` for VariationalAlfvenEvolve). """ doc = rf"""**1. VariationalDensityEvolve:** @@ -364,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 c9ba267c1..618f35d6b 100644 --- a/src/struphy/propagators/__init__.py +++ b/src/struphy/propagators/__init__.py @@ -42,6 +42,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", @@ -49,6 +50,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", @@ -96,6 +98,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 @@ -103,6 +106,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 @@ -147,6 +151,7 @@ "ShearAlfvenPropagator", "TimeDependentSource", "TwoFluidQuasiNeutralFull", + "VariationalAlfvenEvolve", "VariationalDensityEvolve", "VariationalEntropyEvolve", "VariationalMagFieldEvolve", @@ -154,6 +159,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 new file mode 100644 index 000000000..61a6ce253 --- /dev/null +++ b/src/struphy/propagators/tests/test_variational_schur_precond.py @@ -0,0 +1,67 @@ +"""The Kronecker preconditioner of the shear-Alfven Schur operator of the linear variational propagators +(``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 + +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) + + +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_alfven")], + 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 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_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_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