import dataclasses
import logging
from dataclasses import dataclass
from typing import Callable, Literal
import cunumpy as xp
from feectools.linalg.basic import IdentityOperator
from feectools.linalg.solvers import inverse
from feectools.linalg.stencil import StencilVector
from line_profiler import profile
from scope_profiler import ProfileManager
from struphy.feec import preconditioner
from struphy.feec.fused_derivatives import fuse_boundary_gradient
from struphy.feec.mass import L2Projector, WeightedAverageProjection, WeightedMassOperator
from struphy.io.options import LiteralOptions, OptionsBase
from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions, MultiGridPreconditioner
from struphy.linear_algebra.solver import SolverParameters
from struphy.models.variables import FEECVariable, PICVariable, SPHVariable
from struphy.pic.accumulation.filter import FilterParameters
from struphy.pic.accumulation.particles_to_grid import AccumulatorVector, ParticlesToGrid
from struphy.pic.particles import DeltaFParticles6D
from struphy.propagators.base import Propagator
from struphy.utils.utils import check_option
logger = logging.getLogger("struphy")
opts_feec_space = LiteralOptions.OptsFEECSpace
[docs]
class ImplicitDiffusion(Propagator):
r"""
Weak, implicit discretization of the diffusion (or heat) equation (can be used as a Poisson solver too).
Find :math:`\phi \in H^1` such that
.. math::
\int_\Omega \psi\, n_0(\mathbf x)\frac{\partial \phi}{\partial t}\,\textrm d \mathbf x + \int_\Omega \nabla \psi^\top D_0(\mathbf x) \nabla \phi \,\textrm d \mathbf x = \sum_i \int_\Omega \psi\, \rho_i(\mathbf x)\,\textrm d \mathbf x \qquad \forall \ \psi \in H^1\,,
where :math:`n_0, \rho_i:\Omega \to \mathbb R` are real-valued functions and
:math:`D_0:\Omega \to \mathbb R^{3\times 3}`
is a positive diffusion matrix.
Boundary terms from integration by parts are assumed to vanish.
The equation is discretized as
.. math::
\left( \frac{\sigma_1}{\Delta t} \mathbb M^0_{n_0} + \mathbb G^\top \mathbb M^1_{D_0} \mathbb G \right)\, \boldsymbol\phi^{n+1} = \frac{\sigma_2}{\Delta t} \mathbb M^0_{n_0} \boldsymbol\phi^{n} + \frac{\sigma_3}{\Delta t} \sum_i(\Lambda^0, \rho_i )_{L^2}\,,
where :math:`M^0_{n_0}` and :math:`M^1_{D_0}` are :class:`WeightedMassOperators <struphy.feec.mass.WeightedMassOperators>`
and :math:`\sigma_1, \sigma_2, \sigma_3 \in \mathbb R` are artificial parameters that can be tuned to
change the model (see Notes).
Parameters
----------
rho : FEECVariable or Callable or ParticlesToGrid or list, default=None
Source term(s) :math:`\rho_i` on the right-hand side. Several entries can be combined
by passing a ``list``; each is weighted by the corresponding entry of ``rho_coeffs``.
Accepted entries are:
- ``None``: zero source.
- ``FEECVariable`` in ``H1``: its coefficients are mass-multiplied by :math:`M^0` on the fly.
- ``Callable`` :math:`\rho_i(\mathbf x)`, projected to ``H1`` via ``L2Projector``.
- :class:`~struphy.pic.accumulation.particles_to_grid.ParticlesToGrid`, describing a
particle-to-grid deposition. An
:class:`~struphy.pic.accumulation.particles_to_grid.AccumulatorVector` is built from
it at ``allocate()`` and re-accumulated at every call. Its ``accum_space`` selects
what is deposited:
- ``"H1"``: the kernel deposits a scalar density directly, used as-is.
- ``"Hcurl"``: the kernel deposits a 1-form (vector) quantity :math:`\mathbf w`,
e.g. :math:`n\mathbf u`; the propagator then applies the weak divergence
:math:`-\mathbb G^\top` to it before adding it to the right-hand side. This is
used in the Chorin-projection scheme for example.
- a ``list`` containing any mix of the entries above.
rho_coeffs : float or list, default=None
Multiplicative coefficient(s) for ``rho`` sources.
If a scalar is provided, it is applied to a single source.
If a sequence is provided, it has one entry per entry of ``rho``
(the neutralising background of a full-f ``"H1"`` accumulator gets the coefficient
of its accumulator), or one entry per collected source (no background source is collected
with ``delta_f_source=True``, the control variate or delta-f markers).
If ``None``, all coefficients default to ``1.0``.
diagnostic : FEECVariable, default=None
If not None, updates at each call to the propagator, takes the value of the right-hand side.
Otherwise does not provide diagnostic.
delta_f_source : bool, default=False
Deposit particle sources with the delta-f weights :math:`(f - f_0)/s` also for full-f markers (without
the control variate): the weights are updated before and the full-f weights restored after the deposit,
and no background source is added. The Monte-Carlo noise is then that of :math:`f - f_0` instead of
:math:`f`, which is much smaller for small perturbations; meant for the initial Poisson solve of full-f models.
Note
----
The deposited particle density is :math:`f - f_0` with the control variate (the weights are updated before
each deposit), with delta-f markers (:class:`~struphy.pic.particles.DeltaFParticles6D`) and with
``delta_f_source=True``. Otherwise (full-f), for ``"H1"`` accumulators the neutralizing background
:math:`-n_0` is added as a separate source with the coefficient of the particle source.
Notes
-----
* :math:`\sigma_1=\sigma_2=0` and :math:`\sigma_3 = \Delta t`: **Poisson solver** with a given charge density :math:`\sum_i\rho_i`.
* :math:`\sigma_2=0` and :math:`\sigma_1 = \sigma_3 = \Delta t` : Poisson with **adiabatic electrons**.
* :math:`\sigma_1=\sigma_2=1` and :math:`\sigma_3 = 0`: **Implicit heat equation**.
"""
[docs]
class Variables:
"""
Attributes
----------
phi : FEECVariable
Scalar solution field in H^1 space.
"""
def __init__(self):
self._phi: FEECVariable = None
@property
def phi(self) -> FEECVariable:
return self._phi
@phi.setter
def phi(self, new):
assert isinstance(new, FEECVariable)
assert new.space == "H1"
self._phi = new
def __init__(
self,
rho: FEECVariable | ParticlesToGrid | Callable | list = None,
rho_coeffs: float | list = None,
diagnostic: FEECVariable | None = None,
delta_f_source: bool = False,
):
if isinstance(rho, list):
for r in rho:
if isinstance(r, ParticlesToGrid):
assert r.accum_space is not None, "If rho contains a PICVariable, accum_space must be provided."
assert r.accum_kernel is not None, "If rho contains a PICVariable, accum_kernel must be provided."
elif isinstance(rho, ParticlesToGrid):
assert rho.accum_space is not None, "If rho is a PICVariable, accum_space must be provided."
assert rho.accum_kernel is not None, "If rho is a PICVariable, accum_kernel must be provided."
self.variables = self.Variables()
self.rho = rho
self.rho_coeffs = rho_coeffs
self.diagnostic = diagnostic
self.delta_f_source = delta_f_source
[docs]
@dataclass(repr=False)
class Options(OptionsBase):
r"""Configuration options for :class:`ImplicitDiffusion`.
Parameters
----------
sigma_1 : float, default=1.0
Coefficient multiplying the stabilization/mass contribution on the
left-hand side.
With ``divide_by_dt=True`` it is interpreted as ``sigma_1 / dt`` in
the assembled linear system.
sigma_2 : float, default=0.0
Coefficient multiplying the previous solution contribution
``stab_mat * phi^n`` on the right-hand side.
With ``divide_by_dt=True`` it is interpreted as ``sigma_2 / dt``.
sigma_3 : float, default=1.0
Coefficient multiplying the source terms ``rho`` on the
right-hand side.
With ``divide_by_dt=True`` it is interpreted as ``sigma_3 / dt``.
divide_by_dt : bool, default=False
If ``True``, divide ``sigma_1``, ``sigma_2`` and ``sigma_3`` by the
time step passed to ``__call__(dt)`` before assembling the system.
stab_mat : {"M0", "M0ad", "M0ad_withT", "Id"}, default="M0"
Stabilization operator used in the term weighted by ``sigma_1`` and
``sigma_2``.
- ``"M0"``: standard weighted 0-form mass operator.
- ``"M0ad"``: adiabatic-electron weighted 0-form mass operator (weight :math:`n_0`).
- ``"M0ad_withT"``: adiabatic-electron weighted 0-form mass operator (weight :math:`n_0/T_0`).
- ``"Id"``: identity operator.
stab_average : str, default=None
If not None, the stabilization becomes ``stab_mat @ (I - A)``, where ``A`` is the average
over the given logical directions, e.g. ``"eta3"`` or ``"eta2 eta3"``, weighted with
``stab_mat`` (:class:`~struphy.feec.mass.WeightedAverageProjection`; e.g. the flux-surface
average with the volume element for ``"M0"``), such that the operator stays symmetric.
With ``precond="MultiGrid"``, the average is re-built on every coarse level.
diffusion_mat : {"M1", "M1perp", None} or WeightedMassOperator, default="M1"
Diffusion metric in the bilinear form
``grad.T @ diffusion_mat @ grad``.
You can pass the name of a pre-built operator in ``mass_ops`` or a
custom ``WeightedMassOperator`` compatible with the codomain of
``grad``. With ``None`` there is no diffusion term and the system is
``sigma_1 * stab_mat`` only (a weighted L2 projection); then ``precond`` cannot be
``"MultiGrid"`` or ``"StiffnessPreconditioner"``.
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.OptsDiffusionPrecond, default=None
Name of the preconditioner:
- ``None``: no preconditioner.
- ``"MultiGrid"``: geometric multigrid V-cycle
(:class:`~struphy.linear_algebra.multigrid.preconditioner.MultiGridPreconditioner`,
requires ``solver="pcg"`` or ``"pcg_sr"``).
- ``"MassMatrixPreconditioner"``: Kronecker
approximation of the inverse of the 0-form mass matrix ``M0``
(see :mod:`struphy.feec.preconditioner`).
- ``"PolarMassMatrixPreconditioner"``: 2D x 1D Kronecker approximation of the inverse
of ``M0`` on polar splines.
- ``"StiffnessPreconditioner"``: Kronecker approximation of the inverse of
``grad.T @ diffusion_mat @ grad + sigma * stab_mat (I - A)``
(:class:`~struphy.feec.preconditioner.StiffnessPreconditioner` with ``mid=diffusion_mat``,
``stab=stab_mat`` and ``stab_average``), with ``sigma = sigma_1`` (divided by dt if
``divide_by_dt=True``). With polar splines, this is a
:class:`~struphy.feec.preconditioner.PolarStiffnessPreconditioner`.
precond_params : dict, default=None
Keyword arguments passed to the constructor of the mass-matrix or stiffness
preconditioner, e.g. ``{"weights": "unit"}`` for ``"StiffnessPreconditioner"``
or ``{"dim_reduce": None}`` for ``"MassMatrixPreconditioner"``.
multigrid : MultiGridOptions, default=None
Options of the multigrid preconditioner (if ``precond="MultiGrid"``).
If ``None``, defaults to ``MultiGridOptions()``. Set ``nullspace="constants"``
for (nearly) singular problems, e.g. a periodic Poisson problem with tiny ``sigma_1``; for an
exactly singular system (:attr:`lhs_is_singular`) it is set automatically if left ``None``.
solver_params : SolverParameters, default=None
Iterative-solver controls (for example ``tol``, ``maxiter``,
``verbose``, ``info``, ``recycle``).
If ``None``, defaults to ``SolverParameters()``.
enforce_compatibility : bool, default=False
If ``True`` and the diffusion operator is singular (no Dirichlet boundary conditions
and no polar splines, i.e. the constant functions are in its kernel), the right-hand side
is made compatible before the solve by removing its constant ("net charge") component,
.. math::
\mathbf b \leftarrow \mathbf b - \frac{\mathbf 1^\top \mathbf b}{\mathbf 1^\top \mathbb M^0 \mathbf 1}\, \mathbb M^0 \mathbf 1\,,
where :math:`\mathbf 1` (all ones) is the coefficient vector of the constant function
(partition of unity). This is equivalent to adding a uniform neutralizing background.
It is meant for (nearly) singular problems where ``sigma_1`` is only a small
stabilization; then ``phi`` stays of order one instead of growing like
``(net charge) / sigma_1``, which would destroy the accuracy of ``grad phi``.
For ``stab_mat="M0"`` and ``sigma_2 = 0`` it changes ``phi`` only by a constant, for any ``sigma_1``.
If the diffusion operator is singular and ``stab_average`` is set (e.g. the quasi-neutrality equation
with adiabatic electrons in a fully periodic domain), the stabilization ``stab_mat (I - A)`` vanishes
on the constants too, so the whole system is singular (see :attr:`lhs_is_singular`). Then the
right-hand side is always made compatible, whatever ``enforce_compatibility`` says (the constant
component of the averaged equation is the net charge, which must vanish), and the undetermined
constant of ``phi`` is fixed by :math:`\int \phi \, \textrm d \mathbf x = 0` after the solve.
filter_params : dict[PICVariable | SPHVariable, FilterParameters], default=None
If not None, specifies a filter to the accumulation of a specific variable.
Keyed by the ``pic_variable`` of the corresponding
:class:`~struphy.pic.accumulation.particles_to_grid.ParticlesToGrid` source in ``rho``.
Note
----
The accumulation kernel for a particle source is not configured here:
it is carried by the :class:`~struphy.pic.accumulation.particles_to_grid.ParticlesToGrid`
instance passed as (part of) ``rho``.
"""
# specific literals
OptsStabMat = Literal["M0", "M0ad", "M0ad_withT", "Id"]
OptsDiffusionMat = Literal["M1", "M1perp", None]
# propagator options
sigma_1: float = 1.0
sigma_2: float = 0.0
sigma_3: float = 1.0
divide_by_dt: bool = False
stab_mat: OptsStabMat = "M0"
stab_average: str = None
diffusion_mat: OptsDiffusionMat = "M1"
x0: StencilVector = None
solver: LiteralOptions.OptsSymmSolver = "pcg"
precond: LiteralOptions.OptsDiffusionPrecond = None
precond_params: dict = None
multigrid: MultiGridOptions = None
solver_params: SolverParameters = None
enforce_compatibility: bool = False
filter_params: dict[PICVariable | SPHVariable, FilterParameters] = 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.OptsDiffusionPrecond)
if self.precond is not None:
assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'."
if self.diffusion_mat is None:
assert self.precond not in ("MultiGrid", "StiffnessPreconditioner"), (
f"precond={self.precond!r} needs a diffusion term (diffusion_mat=None)."
)
assert self.stab_average is None, (
"Without a diffusion term the averaged part of phi is undetermined (stab_average must be None)."
)
preconditioner._parse_directions(self.stab_average)
# defaults
if self.precond_params is None:
self.precond_params = {}
if self.solver_params is None:
self.solver_params = SolverParameters()
if self.multigrid is None:
self.multigrid = MultiGridOptions()
@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):
# always stabilize
auto_stabilized = False
if xp.abs(self.options.sigma_1) < 1e-14:
self.options.sigma_1 = 1e-14
auto_stabilized = True
logger.warning(f"Stabilizing Poisson solve with {self.options.sigma_1 =}")
# The automatic stabilization only makes singular problems (e.g. periodic in all directions) solvable. With a
# Dirichlet boundary in some direction, grad.T @ M1 @ grad is already positive definite and 1e-14 * stab_mat
# changes the solution only at round-off, but applying it costs a whole stab_mat product per solver iteration;
# it is then left out of the solver's operator. (M1perp can be degenerate, so it keeps the stabilization.)
self._stabilize_lhs = not (
auto_stabilized
and self.options.diffusion_mat == "M1"
and any(any(sides) for sides in self.derham.dirichlet_bc)
)
if not self._stabilize_lhs:
logger.info("Dirichlet boundary present: the automatic stabilization is left out of the solver operator")
# model parameters
self._sigma_1 = self.options.sigma_1
self._sigma_2 = self.options.sigma_2
self._sigma_3 = self.options.sigma_3
self._divide_by_dt = self.options.divide_by_dt
phi = self.variables.phi.spline.vector
# collect rhs
def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector:
"""Perform preliminary operations on rho to compute the rhs and return the result."""
if rho is None:
rhs = phi.space.zeros()
elif isinstance(rho, FEECVariable):
assert rho.space == "H1"
rhs = rho
elif isinstance(rho, ParticlesToGrid):
filter_params = None
if self.options.filter_params is not None:
filter_params = self.options.filter_params.get(rho.pic_variable)
rhs = AccumulatorVector(
rho.pic_variable.particles,
rho.accum_space,
rho.accum_kernel,
Propagator.mass_ops,
Propagator.domain.args_domain,
filter_params=filter_params,
)
if not self._deposits_delta_f(rhs.particles) and rhs.space_id == "H1":
l2_proj = L2Projector("H1", Propagator.mass_ops)
# neutralising background -n0: the accumulator deposits the particle density (weights) and the
# charge (number and sign) is in the coefficient of the source, which the background shares
f0e = -1 * rhs.particles.f0
# get_dofs already returns (n0, Lambda^0)_L2, i.e. the weak-form rhs,
# so it is added as a StencilVector (no extra M0)
rho_eh = l2_proj.get_dofs(f0e.n, apply_bc=True)
return [rhs, rho_eh]
elif isinstance(rho, Callable):
rhs = L2Projector("H1", self.mass_ops).get_dofs(rho, apply_bc=True)
else:
raise TypeError(f"{type(rho) =} is not accepted.")
return [rhs]
rho_list = self.rho if isinstance(self.rho, list) else [self.rho]
# coeffs given by the user (copied, the user's list is not modified)
if self.rho_coeffs is None:
user_coeffs = [1.0 for r in rho_list]
elif isinstance(self.rho_coeffs, (list, tuple)):
user_coeffs = list(self.rho_coeffs)
else:
user_coeffs = [self.rho_coeffs]
# collect sources; a full-f accumulator adds a background source after itself
sources_per_rho = [verify_rhs(r) for r in rho_list]
self._sources = [src for tmp in sources_per_rho for src in tmp]
# coeffs of rhs: one per entry of rho (duplicated for the background source),
# or already one per collected source
if len(user_coeffs) == len(rho_list):
self._coeffs = [c for c, tmp in zip(user_coeffs, sources_per_rho) for src in tmp]
else:
self._coeffs = user_coeffs
assert len(self._coeffs) == len(self._sources), (
f"Got {len(user_coeffs)} rho_coeffs for {len(rho_list)} entries of rho ({len(self._sources)} 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(phi.space)
else:
stab_mat = getattr(self.mass_ops, self.options.stab_mat)
# stabilization without the (stab_mat-weighted) average over the given directions: stab_mat @ (I - A)
average_dirs = preconditioner._parse_directions(self.options.stab_average)
if average_dirs:
avg = WeightedAverageProjection(self.derham, stab_mat, average_dirs)
stab_mat = stab_mat @ (IdentityOperator(stab_mat.domain, stab_mat.codomain) - avg)
if self.options.diffusion_mat is None:
diffusion_mat = None
elif 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.grad.codomain
assert diffusion_mat.codomain == self.derham.grad.codomain
# Set lhs matrices (without dt)
self._stab_mat = stab_mat
if diffusion_mat is None:
self._diffusion_op = None
else:
# grad = B1 G B0^T already applies the boundary operator B1 of the 1-forms on both sides of diffusion_mat;
# if diffusion_mat applies the same one, it is left out there (the same operator, see
# MassWithoutBoundaryOps), except for the multigrid preconditioner, which coarsens the operator by its
# structure
mid = diffusion_mat
grad, grad_t = self.derham.grad, self.derham.grad.T
B1 = self.derham.boundary_ops["1"]
if self.options.precond != "MultiGrid" and all(
op is B1 for op in getattr(diffusion_mat, "boundary_ops", (None,))
):
mid = diffusion_mat.without_boundary_ops()
if self.options.precond != "MultiGrid":
# on the GPU, the gradient and its transpose with their boundary operators as one kernel each
grad, grad_t = fuse_boundary_gradient(grad), fuse_boundary_gradient(grad_t)
self._diffusion_op = grad_t @ mid @ grad
# preconditioner and solver for Ax=b
# the multigrid and stiffness preconditioners depend on sigma_1, they are updated in __call__
# if it changes (e.g. with dt); _pc_sig_1 is the sigma_1 they were built with
self._mg = None
self._stiffness_pc = None
self._pc_sig_1 = self._sigma_1
precond = self.options.precond
if precond is None:
pc = None
elif precond == "MultiGrid":
mg_options = self.options.multigrid
# singular system: the coarse levels must know the kernel, otherwise pcg needs many more iterations
if self.lhs_is_singular and mg_options.nullspace is None:
mg_options = dataclasses.replace(mg_options, nullspace="constants")
self._mg = MultiGridPreconditioner(
self._sigma_1 * stab_mat + self._diffusion_op,
self.derham,
self.domain,
mg_options,
mass_ops=self.mass_ops,
)
pc = self._mg
elif precond == "StiffnessPreconditioner":
self._stiffness_pc = self._build_stiffness_pc(self._sigma_1)
pc = self._stiffness_pc
else:
pc_class = getattr(preconditioner, precond)
pc = pc_class(self.mass_ops.M0, **self.options.precond_params)
# solver just with A_2, but will be set during call with dt
# (only "pcg" and "pcg_sr" take a preconditioner)
pc_kwargs = {"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {}
self._solver = inverse(
stab_mat if self._diffusion_op is None else self._diffusion_op,
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 = phi.space.zeros()
self._tmp_energy = phi.space.zeros()
self._rhs = phi.space.zeros()
self._rhs2 = phi.space.zeros()
self._tmp_src = phi.space.zeros()
# constant function and its weak form M0 @ 1, for removing the kernel component of the rhs
self._ones = None
self._fix_gauge = self.lhs_is_singular
if self._fix_gauge:
logger.info(
"Singular system (no Dirichlet boundary, stabilization without the average): "
"removing the net charge of the rhs and fixing the mean of phi to zero"
)
if (self.options.enforce_compatibility or self._fix_gauge) and self.operator_is_singular:
self._ones = phi.space.zeros()
owned = tuple(slice(s, e + 1) for s, e in zip(self._ones.starts, self._ones.ends))
self._ones[owned] = 1.0
self._ones.update_ghost_regions()
self._m0_ones = self.mass_ops.M0.dot(self._ones)
self._ones_m0_ones = self._ones.inner(self._m0_ones)
@staticmethod
def _is_full_f(particles) -> bool:
"""True if the marker weights are full-f weights f / s (no control variate, no delta-f markers)."""
return not particles.control_variate and not isinstance(particles, DeltaFParticles6D)
def _deposits_delta_f(self, particles) -> bool:
"""True if the particle source deposits f - f0 (so that no neutralizing background is needed)."""
return not self._is_full_f(particles) or self.delta_f_source
@property
def operator_is_singular(self) -> bool:
"""True if the diffusion operator ``grad.T @ diffusion_mat @ grad`` is known to be singular
with the constant functions (coefficient vector of all ones) in its kernel: no Dirichlet boundary
conditions and tensor-product splines. Polar splines are conservatively treated as regular."""
if self._diffusion_op is None:
return False
no_dirichlet = not any(any(bc) for bc in self.derham.dirichlet_bc)
return (
no_dirichlet
and not self.derham.polar_splines
and isinstance(self.variables.phi.spline.vector, StencilVector)
)
@property
def lhs_is_singular(self) -> bool:
r"""True if the whole system matrix ``sigma_1 * stab_mat (I - A) + grad.T @ diffusion_mat @ grad`` is
singular, with the constants in its kernel: the diffusion operator is singular (see
:attr:`operator_is_singular`) and the stabilization removes the average ``A`` (``stab_average``), which
reproduces the constants. Then the right-hand side is made compatible and the mean of ``phi`` is fixed.
The kernel can be larger than the constants if ``diffusion_mat`` is degenerate (``"M1perp"``,
``"M1gyro"``) and :math:`\mathbf B_0` is not along the averaged directions (e.g. ``stab_average="eta3"``
with :math:`\mathbf B_0` along :math:`\eta_1`); such a setup stays singular.
Polar splines are treated as regular, see :attr:`operator_is_singular`."""
return (
self.operator_is_singular
and self._stabilize_lhs
and bool(preconditioner._parse_directions(self.options.stab_average))
)
def _build_stiffness_pc(self, sig_1: float) -> preconditioner.StiffnessPreconditioner:
"""StiffnessPreconditioner for ``sig_1 * stab_mat (I - A) + grad.T @ diffusion_mat @ grad``."""
return preconditioner.StiffnessPreconditioner(
self.mass_ops,
"grad",
sigma=sig_1,
mid=self.options.diffusion_mat,
stab=self.options.stab_mat,
stab_average=self.options.stab_average,
**self.options.precond_params,
)
@property
def sources(self) -> list[StencilVector | 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.V0
assert value.space.symbolic_space == "H1", (
f"Right-hand side must be in H1, 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):
# set parameters
if self._divide_by_dt:
sig_1 = self._sigma_1 / dt
sig_2 = self._sigma_2 / dt
sig_3 = self._sigma_3 / dt
else:
sig_1 = self._sigma_1
sig_2 = self._sigma_2
sig_3 = self._sigma_3
# compute rhs
phin = self.variables.phi.spline.vector
if sig_2 == 0.0:
# no stab_mat product only to multiply it by zero (e.g. Poisson: sigma_2 = 0)
rhs = self._rhs
rhs *= 0.0
else:
rhs = self._stab_mat.dot(phin, out=self._rhs)
rhs *= sig_2
self._rhs2 *= 0.0
for src, coeff in zip(self.sources, self.coeffs):
if isinstance(src, StencilVector):
self._rhs2 += sig_3 * coeff * src
elif isinstance(src, FEECVariable):
v = src.spline.vector
self._rhs2 += sig_3 * coeff * self.mass_ops.M0.dot(v, out=self._tmp_src)
elif isinstance(src, AccumulatorVector):
# delta-f weights for full-f markers, restored after the deposit (delta_f_source)
restore_full_f = self.delta_f_source and self._is_full_f(src.particles)
if src.particles.control_variate or restore_full_f:
src.particles.update_weights()
if src.space_id == "H1":
src() # accumulate
vec = src.vectors[0]
elif src.space_id == "Hcurl":
# 1. update density at marker positions
eta = src.particles.positions
valid_mks = src.particles.valid_mks
first_free_idx = src.particles.first_free_idx
density = src.particles.f0.n0(eta)
src.particles.markers[first_free_idx, valid_mks] = density
# 2. accumulate
src()
# 3. take weak divergence
vec = -self.derham.grad.T.dot(src.vectors[0], out=self._tmp_src)
else:
raise ValueError(f"Unsupported source space {src.space_id}.")
if restore_full_f:
src.particles.weights = src.particles.weights0.copy()
self._rhs2 += sig_3 * coeff * vec
rhs += self._rhs2
# remove the kernel component (net charge) of the rhs: 1^T rhs = 0 (compatibility condition)
if self._ones is not None:
rhs.mul_iadd(-self._ones.inner(rhs) / self._ones_m0_ones, self._m0_ones)
if self.diagnostic is not None:
proj = L2Projector("H1", self.mass_ops)
self.diagnostic.spline.vector = proj.solve(rhs)
# compute lhs
if self._diffusion_op is None:
self._solver.linop = sig_1 * self._stab_mat
elif self._stabilize_lhs:
self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op
else:
self._solver.linop = self._diffusion_op
if sig_1 != self._pc_sig_1:
if self._mg is not None:
self._mg.update(self._solver.linop)
if self._stiffness_pc is not None:
self._stiffness_pc = self._build_stiffness_pc(sig_1)
self._solver.set_options(pc=self._stiffness_pc)
self._pc_sig_1 = sig_1
# solve
with ProfileManager.profile_region(self._solve_region, functions=[self._solver.solve]):
out = self._solver.solve(rhs, out=self._tmp)
info = self._solver._info
# singular system: fix the undetermined constant by a zero mean, 1^T M0 phi = 0
if self._fix_gauge:
out.mul_iadd(-self._m0_ones.inner(out) / self._ones_m0_ones, self._ones)
if self._info:
logger.warning(f"\nSolver info: {info}")
self.update_feec_variables(phi=out)
