import logging
from dataclasses import dataclass
from typing import Callable, Literal
import cunumpy as xp
from feectools.linalg.basic import IdentityOperator, Vector
from feectools.linalg.solvers import inverse
from feectools.linalg.stencil import StencilVector
from line_profiler import profile
from maybempi import MPI
from scope_profiler import ProfileManager
from struphy.feec import preconditioner
from struphy.feec.mass import L2Projector, WeightedMassOperator
from struphy.io.options import LiteralOptions
from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions, MultiGridPreconditioner
from struphy.linear_algebra.solver import SolverParameters
from struphy.models.variables import FEECVariable
from struphy.pic.accumulation.particles_to_grid import AccumulatorVector
from struphy.pic.base import Particles
from struphy.propagators.base import Propagator
from struphy.utils.utils import check_option
logger = logging.getLogger("struphy")
[docs]
class CurlCurlSolve(Propagator):
r"""
Weak discretization of the curl-curl problem.
Find :math:`\mathbf E \in H(\textnormal{curl})` such that
.. math::
\int_\Omega \nabla \times \mathbf F \cdot \nabla \times \mathbf E\,\textrm d \mathbf x + \sigma \int_\Omega \mathbf F \cdot \mathbf E\,\textrm d \mathbf x = \sum_i \int_\Omega \mathbf F \cdot \mathbf J _i\,\textrm d \mathbf x \qquad \forall \,\mathbf F \in H(\textnormal{curl})\,,
where :math:`\mathbf J _i:\Omega \to \mathbb R^3` are real-valued and
:math:`\sigma > 0` is a scalar (the operator is symmetric positive definite).
Boundary terms from integration by parts are assumed to vanish.
The equation is discretized as
.. math::
\left( \mathbb C^\top \cdot \mathbb M^2 \cdot \mathbb C + \sigma \mathbb M^1 \right)\, \boldsymbol e^{n+1} =\sum_i \boldsymbol j_i\,,
where :math:`\mathbb M^1` and :math:`\mathbb M^2` are :class:`WeightedMassOperators <struphy.feec.mass.WeightedMassOperators>` and :math:`\boldsymbol j_i`
is the vector of coefficients of the projection of :math:`\mathbf J_i` into the discrete space :math:`V^1_h\subset H(\textnormal{curl})`.
"""
[docs]
class Variables:
"""
Attributes
----------
e : FEECVariable
Vector-valued solution in ``"Hcurl"`` space.
"""
def __init__(self):
self._e: FEECVariable = None
@property
def e(self) -> FEECVariable:
return self._e
@e.setter
def e(self, new):
assert isinstance(new, FEECVariable)
assert new.space == "Hcurl"
self._e = new
def __init__(
self,
j: FEECVariable | tuple[Callable, Callable, Callable] | tuple[AccumulatorVector, Particles] | list = None,
j_coeffs: float | list = None,
):
"""
Parameters
----------
j : FEECVariable or tuple of Callables or tuple or list, default=None
Source term(s) on the right-hand side.
Accepted entries are:
- ``None``: zero source.
- ``FEECVariable`` in ``"Hcurl"``.
- ``tuple`` (or ``list``) of three ``Callable``s to be projected to ``"Hcurl"`` via
``L2Projector``.
- ``AccumulatorVector``.
- a ``list`` containing any mix of the entries above.
The tuple form is accepted by typing for compatibility with other
propagator interfaces that pair particle data with accumulators.
j_coeffs : float or list, default=None
Multiplicative coefficient(s) for ``j`` sources.
If a scalar is provided, it is applied to a single source.
If a sequence is provided, its length must match the number of
collected sources.
If ``None``, all coefficients default to ``1.0``.
"""
self.variables = self.Variables()
self.j = j
self.j_coeffs = j_coeffs
[docs]
@dataclass
class Options:
"""Configuration options for :class:`CurlCurlSolve`.
Parameters
----------
sigma : float, default=1.0
Coefficient of the stabilization/mass term on the left-hand side,
``curl.T @ diffusion_mat @ curl + sigma * stab_mat``. Must be positive for the
operator to be positive definite (values below ``1e-14`` are set to ``1e-14``).
stab_mat : {"M1", "Id"}, default="M1"
Stabilization operator used in the term weighted by ``sigma``.
- ``"M1"``: standard weighted 1-form mass operator.
- ``"Id"``: identity operator.
diffusion_mat : "M2" or WeightedMassOperator, default="M2"
Diffusion metric in the bilinear form
``curl.T @ diffusion_mat @ curl``.
You can pass the name of a pre-built operator in ``mass_ops`` or a
custom ``WeightedMassOperator`` compatible with the codomain of
``curl``.
x0 : StencilVector, default=None
Initial guess for the iterative linear solver.
solver : LiteralOptions.OptsSymmSolver, default="pcg"
Name of the symmetric iterative solver passed to
:func:`psydac.linalg.solvers.inverse`.
precond : LiteralOptions.OptsCurlCurlPrecond, default=None
Name of the preconditioner (requires ``solver="pcg"``):
- ``None``: no preconditioner.
- ``"MassMatrixPreconditioner"``: Kronecker approximation of the inverse of
the 1-form mass matrix ``M1`` (see :mod:`struphy.feec.preconditioner`).
- ``"PolarMassMatrixPreconditioner"``: 2D x 1D Kronecker approximation of the inverse
of ``M1`` on polar splines.
- ``"StiffnessPreconditioner"``: Kronecker approximation of the inverse of
``curl.T @ diffusion_mat @ curl + sigma * stab_mat``
(:class:`~struphy.feec.preconditioner.StiffnessPreconditioner` with ``derivative="curl"``,
``mid=diffusion_mat`` and ``stab=stab_mat``): block Jacobi with fast diagonalization,
plus a correction on the gradients (``kernel_correction``, default True). With polar
splines, this is a :class:`~struphy.feec.preconditioner.PolarStiffnessPreconditioner`.
- ``"HiptmairXu"``: auxiliary space preconditioner of Hiptmair and Xu
(:class:`~struphy.feec.preconditioner.HiptmairXuPreconditioner` with ``mid=diffusion_mat`` and
``stab=stab_mat``): smoother plus corrections on the gradients and on vector-valued
:math:`H^1` splines.
- ``"MultiGrid"``: geometric multigrid V-cycle
(:class:`~struphy.linear_algebra.multigrid.preconditioner.MultiGridPreconditioner`)
with the hybrid smoother of Hiptmair (see ``multigrid``).
precond_params : dict, default=None
Keyword arguments passed to the constructor of the preconditioner,
e.g. ``{"kernel_correction": False}`` for ``"StiffnessPreconditioner"``
or ``{"smoother": "jacobi"}`` for ``"HiptmairXu"``.
multigrid : MultiGridOptions, default=None
Options of the multigrid preconditioner (if ``precond="MultiGrid"``).
If ``None``, defaults to ``MultiGridOptions(hiptmair=True, smoother_precond="jacobi")``:
without the Hiptmair smoother, multigrid is not robust for the curl-curl operator.
solver_params : SolverParameters, default=None
Iterative-solver controls (for example ``tol``, ``maxiter``,
``verbose``, ``info``, ``recycle``).
If ``None``, defaults to ``SolverParameters()``.
"""
# specific literals
OptsStabMat = Literal["M1", "Id"]
OptsDiffusionMat = Literal["M2"]
# propagator options
sigma: float = 1.0
stab_mat: OptsStabMat = "M1"
diffusion_mat: OptsDiffusionMat = "M2"
x0: StencilVector = None
solver: LiteralOptions.OptsSymmSolver = "pcg"
precond: LiteralOptions.OptsCurlCurlPrecond = None
precond_params: dict = None
multigrid: MultiGridOptions = None
solver_params: SolverParameters = None
def __post_init__(self):
# checks
check_option(self.stab_mat, self.OptsStabMat)
check_option(self.diffusion_mat, self.OptsDiffusionMat)
check_option(self.solver, LiteralOptions.OptsSymmSolver)
check_option(self.precond, LiteralOptions.OptsCurlCurlPrecond)
if self.precond is not None:
assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'."
# defaults
if self.precond_params is None:
self.precond_params = {}
if self.multigrid is None:
self.multigrid = MultiGridOptions(hiptmair=True, smoother_precond="jacobi")
if self.solver_params is None:
self.solver_params = SolverParameters()
@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
@profile
def allocate(self, verbose: bool = False):
if self.options.sigma < 0.0:
raise ValueError(f"CurlCurlSolve needs sigma >= 0 (SPD operator), got {self.options.sigma = }.")
# always stabilize
if self.options.sigma < 1e-14:
self.options.sigma = 1e-14
if MPI.COMM_WORLD.Get_rank() == 0:
logger.info(f"Running Curl-Curl solve with {self.options.sigma =}")
# model parameters
self._sigma = self.options.sigma
e = self.variables.e.spline.vector
# collect rhs
def is_callable_triple(j) -> bool:
"""True if j is a tuple/list of three callables (components of a 1-form)."""
return isinstance(j, (tuple, list)) and len(j) == 3 and all(callable(ji) for ji in j)
def verify_rhs(j) -> Vector | FEECVariable | AccumulatorVector:
"""Perform preliminary operations on j to compute the rhs and return the result."""
if j is None:
rhs = e.space.zeros()
elif isinstance(j, FEECVariable):
assert j.space == "Hcurl"
rhs = j
elif isinstance(j, AccumulatorVector):
rhs = j
elif is_callable_triple(j):
rhs = L2Projector("Hcurl", self.mass_ops).get_dofs(j, apply_bc=True)
else:
raise TypeError(f"{type(j) =} is not accepted.")
return rhs
j = self.j
if isinstance(j, list) and not is_callable_triple(j):
self._sources = []
for ji in j:
self._sources += [verify_rhs(ji)]
else:
self._sources = [verify_rhs(j)]
# coeffs of rhs
if self.j_coeffs is not None:
if isinstance(self.j_coeffs, (list, tuple)):
self._coeffs = self.j_coeffs
else:
self._coeffs = [self.j_coeffs]
assert len(self._coeffs) == len(self._sources)
else:
self._coeffs = [1.0 for src in self.sources]
# initial guess and solver params
self._x0 = self.options.x0
self._info = self.options.solver_params.info
if self.options.stab_mat == "Id":
stab_mat = IdentityOperator(e.space)
else:
stab_mat = getattr(self.mass_ops, self.options.stab_mat)
if isinstance(self.options.diffusion_mat, str):
diffusion_mat = getattr(self.mass_ops, self.options.diffusion_mat)
else:
diffusion_mat = self.options.diffusion_mat
assert isinstance(diffusion_mat, WeightedMassOperator)
assert diffusion_mat.domain == self.derham.curl.codomain
assert diffusion_mat.codomain == self.derham.curl.codomain
# Set lhs matrices (without dt)
self._stab_mat = stab_mat
self._diffusion_op = self.derham.curl.T @ diffusion_mat @ self.derham.curl
# preconditioner and solver for Ax=b
lhs = self._diffusion_op + self._sigma * self._stab_mat
precond = self.options.precond
if precond is None:
pc = None
elif precond == "StiffnessPreconditioner":
pc = preconditioner.StiffnessPreconditioner(
self.mass_ops,
"curl",
sigma=self._sigma,
mid=self.options.diffusion_mat,
stab=self.options.stab_mat,
**self.options.precond_params,
)
elif precond == "MultiGrid":
pc = MultiGridPreconditioner(
lhs,
self.derham,
self.domain,
self.options.multigrid,
mass_ops=self.mass_ops,
)
elif precond == "HiptmairXu":
pc = preconditioner.HiptmairXuPreconditioner(
lhs,
self.mass_ops,
self._sigma,
mid=self.options.diffusion_mat,
stab=self.options.stab_mat,
**self.options.precond_params,
)
else:
pc_class = getattr(preconditioner, precond)
pc = pc_class(self.mass_ops.M1, **self.options.precond_params)
# (only "pcg" takes a preconditioner)
pc_kwargs = {"pc": pc} if self.options.solver == "pcg" else {}
self._solver = inverse(
lhs,
self.options.solver,
**pc_kwargs,
x0=self.x0,
tol=self.options.solver_params.tol,
maxiter=self.options.solver_params.maxiter,
verbose=self.options.solver_params.verbose,
recycle=self.options.solver_params.recycle,
)
# allocate memory for solution
self._tmp = e.space.zeros()
self._rhs = e.space.zeros()
self._tmp_src = e.space.zeros()
@property
def sources(self) -> list[Vector | FEECVariable | AccumulatorVector]:
"""
Right-hand side of the equation (sources).
"""
return self._sources
@property
def coeffs(self) -> list[float]:
"""
Same length as self.sources. Coefficients multiplied with sources before solve (default is 1.0).
"""
return self._coeffs
@property
def x0(self):
"""
feectools.linalg.stencil.StencilVector or struphy.polar.basic.PolarVector. First guess of the iterative solver.
"""
return self.options.x0
@x0.setter
def x0(self, value: StencilVector):
"""In-place setter for StencilVector/PolarVector. First guess of the iterative solver."""
assert value.space == self.derham.V1
assert value.space.symbolic_space == "Hcurl", (
f"Right-hand side must be in Hcurl, but is in {value.space.symbolic_space}."
)
if self.options.x0 is None:
self.options.x0 = value
else:
self.options.x0[:] = value[:]
@profile
def __call__(self, dt):
# compute rhs
self._rhs *= 0.0
for src, coeff in zip(self.sources, self.coeffs):
if isinstance(src, Vector):
self._rhs += coeff * src
elif isinstance(src, FEECVariable):
v = src.spline.vector
self._rhs += coeff * self.mass_ops.M1.dot(v, out=self._tmp_src)
elif isinstance(src, AccumulatorVector):
src() # accumulate
self._rhs += coeff * src.vectors[0]
# solve
with ProfileManager.profile_region(self._solve_region, functions=[self._solver.solve]):
out = self._solver.solve(self._rhs, out=self._tmp)
info = self._solver._info
if self._info:
logger.info(info)
self.update_feec_variables(e=out)
