[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("")