Source code for struphy.propagators.shear_alfven_propagator

import logging
from dataclasses import dataclass
from typing import Literal

from feectools.ddm.mpi import mpi as MPI
from feectools.linalg.block import BlockVector
from feectools.linalg.solvers import inverse
from line_profiler import profile

from struphy.feec import preconditioner
from struphy.io.options import LiteralOptions, OptionsBase
from struphy.linear_algebra.schur_solver import SchurSolver
from struphy.linear_algebra.solver import SolverParameters
from struphy.models.variables import FEECVariable
from struphy.ode.solvers import ODEsolverFEEC
from struphy.ode.utils import ButcherTableau
from struphy.propagators.base import Propagator
from struphy.utils.utils import check_option

logger = logging.getLogger("struphy")


[docs] class ShearAlfvenPropagator(Propagator): r"""FEEC discretization of the following equations: find :math:`\tilde{\mathbf{U}} \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\}` and :math:`\tilde{\mathbf{B}} \in H(\textnormal{div})` such that .. math:: \int_{\Omega} \rho_0\frac{\partial \tilde{\mathbf{U}}}{\partial t} \cdot \mathbf V\,\textnormal d \mathbf x &= \int_{\Omega} \tilde{\mathbf{B}} \cdot \nabla \times (\mathbf{B}_0 \times \mathbf V) \,\textnormal d \mathbf x \qquad \forall \,\mathbf V \in \{H(\textnormal{curl}), H(\textnormal{div}), (H^1)^3\} \\[2mm] \frac{\partial \tilde{\mathbf{B}}}{\partial t} &= \nabla\times(\tilde{\mathbf{U}} \times \mathbf{B}_0) Time discretization: - implicit: Crank-Nicolson (implicit mid-point) - explicit: explicit RK methods from ButcherTableau System size reduction via :class:`~struphy.linear_algebra.schur_solver.SchurSolver`. """
[docs] class Variables: """Container for variables advanced by :class:`ShearAlfvenPropagator`. Attributes ---------- u : FEECVariable Velocity variable in one of ``"Hcurl"``, ``"Hdiv"``, or ``"H1vec"``. 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 in ("Hcurl", "Hdiv", "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): self.variables = self.Variables()
[docs] @dataclass(repr=False) class Options(OptionsBase): """Configuration options for :class:`ShearAlfvenPropagator`. Parameters ---------- u_space : LiteralOptions.OptsVecSpace, default="Hdiv" FEEC space used for the velocity variable. algo : {"implicit", "explicit"}, default="implicit" Time stepping algorithm. solver : LiteralOptions.OptsSymmSolver, default="pcg" Symmetric iterative solver used by implicit or explicit operators. precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner" Preconditioner for the mass-matrix block. solver_params : SolverParameters, default=None Solver controls; defaults to ``SolverParameters()``. butcher : ButcherTableau, default=None Explicit Runge-Kutta tableau when ``algo="explicit"``. """ # specific literals OptsAlgo = Literal["implicit", "explicit"] # propagator options u_space: LiteralOptions.OptsVecSpace = "Hdiv" algo: OptsAlgo = "implicit" solver: LiteralOptions.OptsSymmSolver = "pcg" precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" solver_params: SolverParameters = None butcher: ButcherTableau = None def __post_init__(self): # checks check_option(self.u_space, LiteralOptions.OptsVecSpace) check_option(self.algo, self.OptsAlgo) 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.algo == "explicit" and self.butcher is None: self.butcher = ButcherTableau()
@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): u_space = self.options.u_space # define block matrix [[A B], [C I]] (without time step size dt in the diagonals) id_M = "M" + self.derham.space_to_form[u_space] + "n" id_T = "T" + self.derham.space_to_form[u_space] # call operators _A = getattr(self.mass_ops, id_M) _T = getattr(self.basis_ops, id_T) _M2 = self.mass_ops.M2 curl = self.derham.curl # Preconditioner if self.options.precond is None: pc = None else: pc_class = getattr(preconditioner, self.options.precond) pc = pc_class(getattr(self.mass_ops, id_M)) if self.options.algo == "implicit": self._info = self.options.solver_params.info self._B = -1 / 2 * _T.T @ curl.T @ _M2 self._C = 1 / 2 * curl @ _T # instantiate Schur solver (constant in this case) _BC = self._B @ self._C self._schur_solver = SchurSolver( _A, _BC, self.options.solver, precond=pc, solver_params=self.options.solver_params, ) # pre-allocate arrays self._byn = self._B.codomain.zeros() else: self._info = False # define vector field A_inv = inverse( _A, self.options.solver, pc=pc, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, verbose=self.options.solver_params.verbose, ) _f1 = A_inv @ _T.T @ curl.T @ _M2 _f2 = curl @ _T # allocate output of vector field out1 = self.variables.u.spline.vector.space.zeros() out2 = self.variables.b.spline.vector.space.zeros() def f1(t, y1, y2, out: BlockVector = out1): _f1.dot(y2, out=out) out.update_ghost_regions() return out def f2(t, y1, y2, out: BlockVector = out2): _f2.dot(y1, out=out) out *= -1.0 out.update_ghost_regions() return out vector_field = {self.variables.u.spline.vector: f1, self.variables.b.spline.vector: f2} self._ode_solver = ODEsolverFEEC(vector_field, butcher=self.options.butcher) # allocate dummy vectors to avoid temporary array allocations self._u_tmp1 = self.variables.u.spline.vector.space.zeros() self._u_tmp2 = self.variables.u.spline.vector.space.zeros() self._b_tmp1 = self.variables.b.spline.vector.space.zeros() @profile def __call__(self, dt): # current FE coeffs un = self.variables.u.spline.vector bn = self.variables.b.spline.vector if self.options.algo == "implicit": # solve for new u coeffs byn = self._B.dot(bn, out=self._byn) un1, info = self._schur_solver(un, byn, dt, out=self._u_tmp1) # new b coeffs _u = un.copy(out=self._u_tmp2) _u += self._u_tmp1 bn1 = self._C.dot(_u, out=self._b_tmp1) bn1 *= -dt bn1 += bn diffs = self.update_feec_variables(u=un1, b=bn1) else: self._ode_solver(0.0, dt) if self._info and MPI.COMM_WORLD.Get_rank() == 0: if self.options.algo == "implicit": logger.info(f"Status for ShearAlfven: {info['success']}") logger.info(f"Iterations for ShearAlfven: {info['niter']}") logger.info(f"Maxdiff up for ShearAlfven: {diffs['u']}") logger.info(f"Maxdiff b2 for ShearAlfven: {diffs['b']}") logger.info("")