diff --git a/feectools b/feectools index 2e3aa651d..a991b7182 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 2e3aa651d6d7fad5072dd39ff8c375fb7ed8ae1d +Subproject commit a991b7182b5efdeafc2b968ac016b1dc87de335d diff --git a/pyproject.toml b/pyproject.toml index 6538f59f7..d7bec3f8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "numpy<=2.5.0", "cunumpy>=0.2.0, <=0.3.0", "pyccel>=2.2.0, <=2.2.3", - "feectools>=0.1.11, <=0.2.0", + "feectools>=0.3.0, <=0.3.0", "scipy<=1.18.0", "h5py<=3.16.0", "h5netcdf<=1.8.1", diff --git a/src/struphy/feec/basis_projection_ops.py b/src/struphy/feec/basis_projection_ops.py index ff5ecf806..43d3d4c86 100644 --- a/src/struphy/feec/basis_projection_ops.py +++ b/src/struphy/feec/basis_projection_ops.py @@ -1851,6 +1851,56 @@ def dot(self, v, out=None, tol=1e-14, maxiter=1000): return out + @property + def is_reconstructible(self) -> bool: + """Whether the operator can be re-created from :meth:`to_dict` (e.g. on another Derham). + + False if a weight is given as values at the projection points (an array bound to the current Derham). + """ + return not any(isinstance(w, xp.ndarray) for row in self._weights for w in row) + + def to_dict(self) -> dict: + """Recipe for re-creating the operator with :meth:`from_dict` (on any Derham). + + Weights are stored as given (callables are kept as objects, hence the dictionary is in general not JSON serializable). + """ + if not self.is_reconstructible: + raise ValueError("BasisProjectionOperator with weights given as arrays cannot be serialized.") + V_id, W_id = ( + (self._codomain_symbolic_name, self._domain_symbolic_name) + if self._transposed + else (self._domain_symbolic_name, self._codomain_symbolic_name) + ) + return { + "type": self.__class__.__name__, + "params": { + "V_id": V_id, + "W_id": W_id, + "weights": [list(row) for row in self._weights], + "transposed": self._transposed, + "polar_shift": self._polar_shift, + "use_cache": self._use_cache, + }, + } + + @classmethod + def from_dict(cls, dct: dict, derham: Derham) -> "BasisProjectionOperator": + """Re-create a :class:`BasisProjectionOperator` from :meth:`to_dict` on the given Derham, + with the Derham's (global) commuting projector, extraction and boundary operators.""" + assert dct["type"] == cls.__name__ + params = dct["params"] + V_id, W_id = params["V_id"], params["W_id"] + return cls( + derham.projectors[W_id], + derham.fem_spaces[V_id], + [list(row) for row in params["weights"]], + V_extraction_op=derham.extraction_ops[V_id], + V_boundary_op=derham.boundary_ops[V_id], + transposed=params["transposed"], + polar_shift=params["polar_shift"], + use_cache=params["use_cache"], + ) + def transpose(self, conjugate=False): """ Returns the transposed operator. diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 963c123ca..b230a64c0 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -1254,6 +1254,23 @@ def f_call_matrix(e1, e2, e3): dry_run=dry_run, ) + # weights given at quadrature points or as spline functions are bound to this Derham + grid_bound = len(spline_functions) > 0 or ( + isinstance(weights, list) and any(isinstance(w, xp.ndarray) for row in weights for w in row) + ) + out._creation_info = ( + None + if grid_bound + else { + "V_id": V_id, + "W_id": W_id, + "name": name, + "weights": weights, + "transposed": transposed, + "is_transpose": False, + } + ) + if assemble and not dry_run: out.assemble() @@ -1465,6 +1482,9 @@ def __init__( self._name = name self._dry_run = dry_run + # recipe for re-creating the operator with WeightedMassOperators.create_weighted_mass, see to_dict() + self._creation_info: dict | None = None + assert not (dry_run and transposed), "dry_run=True is not supported for transposed operators." # spline functions that are used as weights in the operator, to be evaluated at quadrature points @@ -2066,6 +2086,9 @@ def transpose(self, conjugate=False): # weights of M in its own (transposed) block order M._weights = [[self._weights[n][m] for n in range(len(self._weights))] for m in range(len(self._weights[0]))] + if self._creation_info is not None: + M._creation_info = dict(self._creation_info, is_transpose=not self._creation_info["is_transpose"]) + if self._matrix_free: if self._symmetry is not None: M.assemble(weights=M._weights) @@ -2107,6 +2130,9 @@ def assemble(self, weights=None, clear=True): assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) + if weights is not None or not clear: + # the data no longer stems from the creation recipe + self._creation_info = None if self._matrix_free: if weights is not None: @@ -2325,6 +2351,67 @@ def assemble(self, weights=None, clear=True): logger.debug("Done.") + @property + def is_reconstructible(self) -> bool: + """Whether the operator can be re-created from :meth:`to_dict` (e.g. on another Derham). + + True for operators created by :meth:`WeightedMassOperators.create_weighted_mass` (and their transposes) + whose data has not been modified afterwards (by ``assemble(weights=...)``, in-place arithmetic, ...). + """ + return self._creation_info is not None + + def to_dict(self) -> dict: + """Recipe for re-creating the operator with :meth:`WeightedMassOperators.create_weighted_mass`. + + The weights are stored as given at creation. The dictionary is JSON serializable if they are + strings (``'Ginv'``, ``'sqrt_g'``, ...) or nested lists of numbers; callables are kept as objects. + Re-create the operator (on any Derham) with :meth:`from_dict`. + """ + if self._creation_info is None: + raise ValueError( + f"WeightedMassOperator {self.name!r} cannot be serialized: it was not created by " + "WeightedMassOperators.create_weighted_mass or its data was modified afterwards." + ) + params = dict(self._creation_info) + # tuple (1D product of weights) and 2D list (block weights) are different formats; store the + # tuple as a list (JSON) and record its type, since its entries may themselves be (3x3) lists + params["weights_is_tuple"] = isinstance(params["weights"], tuple) + if params["weights_is_tuple"]: + params["weights"] = list(params["weights"]) + return { + "type": self.__class__.__name__, + "params": params, + } + + @classmethod + def from_dict(cls, dct: dict, mass_ops: "WeightedMassOperators") -> "WeightedMassOperator": + """Re-create a :class:`WeightedMassOperator` from :meth:`to_dict` with the given collection. + + Parameters + ---------- + dct : dict + Output of :meth:`to_dict`. + + mass_ops : WeightedMassOperators + Collection providing the Derham, domain and matrix_free option of the new operator. + """ + assert dct["type"] == cls.__name__ + params = dct["params"] + name = params["name"] + weights = params["weights"] + if params["weights_is_tuple"]: + weights = tuple(weights) + + out = mass_ops.create_weighted_mass( + params["V_id"], + params["W_id"], + name=name, + weights=weights, + assemble=True, + transposed=params["transposed"], + ) + return out.T if params["is_transpose"] else out + def copy(self, out=None): """Create a copy of self, that can potentially be stored in a given WeightedMassOperator. @@ -2357,10 +2444,17 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + + if self._creation_info is None: + out._creation_info = None + else: + out._creation_info = dict(self._creation_info) # to create a separate dictionary + return out def __imul__(self, a): self._mat *= a + self._creation_info = None return self def __iadd__(self, M): @@ -2368,10 +2462,12 @@ def __iadd__(self, M): assert M.codomain is self.codomain if isinstance(M, WeightedMassOperator): + self._creation_info = None self._mat += M._mat return self elif isinstance(M, LinearOperator): + self._creation_info = None self._mat += M return self @@ -2383,10 +2479,12 @@ def __isub__(self, M): assert M.codomain is self.codomain if isinstance(M, WeightedMassOperator): + self._creation_info = None self._mat -= M._mat return self elif isinstance(M, LinearOperator): + self._creation_info = None self._mat -= M return self diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index e16af24c5..2b696bda1 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -563,6 +563,11 @@ class Derham: domain : Domain, optional The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients. + domain_decomposition : DomainDecomposition, optional + Prescribed MPI decomposition of the elements, e.g. ``fine_derham.domain_decomposition.coarsen(...)`` + for an aligned multigrid level. Must match ``grid.num_elements`` and the periodicity from ``options.bcs``, + and be built on ``comm``. If None (default), it is computed from ``comm`` and ``grid.mpi_dims_mask``. + Notes ----- The underlying base sequence is @@ -579,6 +584,8 @@ def __init__( options: DerhamOptions, comm: MPI.Intracomm = None, domain: Domain = None, + *, + domain_decomposition: DomainDecomposition | None = None, ): # inputs @@ -680,6 +687,7 @@ def __init__( comm=self.comm, mpi_dims_mask=mpi_dims_mask, use_feectools=use_feectools, + domain_decomposition=domain_decomposition, ) # FEM spaces @@ -1520,6 +1528,7 @@ def init_derham( comm=None, mpi_dims_mask: tuple[bool, bool, bool] = None, use_feectools: bool = True, + domain_decomposition: DomainDecomposition | None = None, ) -> DiscreteDerham: """Return a discrete Derham complex. Allows for the use of tiny-feectools. @@ -1543,12 +1552,29 @@ def init_derham( use_feectools: bool Use slimmed-down fork `feectools` of Psydac. + + domain_decomposition : DomainDecomposition, optional + Prescribed decomposition of the elements; if None, it is computed from ``comm`` and ``mpi_dims_mask``. """ if use_feectools: - self._domain_decomposition = DomainDecomposition( - num_elements, spl_kind, comm=comm, mpi_dims_mask=mpi_dims_mask - ) + if domain_decomposition is None: + self._domain_decomposition = DomainDecomposition( + num_elements, spl_kind, comm=comm, mpi_dims_mask=mpi_dims_mask + ) + else: + assert tuple(domain_decomposition.ncells) == tuple(num_elements), ( + f"{domain_decomposition.ncells = } does not match {num_elements = }." + ) + assert tuple(domain_decomposition.periods) == tuple(spl_kind), ( + f"{domain_decomposition.periods = } does not match {spl_kind = }." + ) + if domain_decomposition.comm is not None and comm is not None: + # (comm is None in the decomposition when feectools runs with MockMPI) + assert domain_decomposition.comm == comm, ( + "domain_decomposition must be built on the Derham communicator." + ) + self._domain_decomposition = domain_decomposition _derham = self._discretize_derham( num_elements, diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 72299fbbd..289404c52 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -76,6 +76,7 @@ class LiteralOptions: OptsSymmSolver = Literal["pcg", "cg"] OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] + OptsDiffusionPrecond = Literal["MultiGrid", "MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None] OptsSaddlePointSolver = Literal["uzawa"] OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"] OptsNonlinearSolver = Literal["Picard", "Newton"] diff --git a/src/struphy/linear_algebra/multigrid/__init__.py b/src/struphy/linear_algebra/multigrid/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/struphy/linear_algebra/multigrid/coarsen.py b/src/struphy/linear_algebra/multigrid/coarsen.py new file mode 100644 index 000000000..b33bce84e --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/coarsen.py @@ -0,0 +1,175 @@ +r"""Re-discretization of a linear operator on a coarser Derham. + +A (composite) operator on the fine level is walked as an expression tree. Composite nodes (sums, +compositions, scalings, powers, block operators) are rebuilt from their coarsened children, scalars are +kept. Leaves are re-created on the coarse Derham: + +* :class:`IdentityOperator`, :class:`ZeroOperator`, :class:`BoundaryOperator`: on the corresponding coarse spaces, +* :class:`DirectionalDerivativeOperator` (blocks of ``derham.grad``, ``curl``, ``div`` and their transposes), +* :class:`WeightedMassOperator` and :class:`BasisProjectionOperator`: from their ``to_dict()`` recipe. + +Further leaf types can be supported with :func:`register_coarsening`. +""" + +from collections.abc import Callable + +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ( + ComposedLinearOperator, + IdentityOperator, + LinearOperator, + PowerLinearOperator, + ScaledLinearOperator, + SumLinearOperator, + VectorSpace, + ZeroOperator, +) +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace + +from struphy.feec.basis_projection_ops import BasisProjectionOperator +from struphy.feec.linear_operators import BoundaryOperator +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.geometry.base import Domain + +_REGISTRY: dict[type, Callable[[LinearOperator, "OperatorCoarsener"], LinearOperator]] = {} + + +def register_coarsening(cls: type): + """Decorator registering ``fun(op, coarsener) -> LinearOperator`` as the coarsening rule for leaves of type ``cls``.""" + + def decorator(fun): + _REGISTRY[cls] = fun + return fun + + return decorator + + +class OperatorCoarsener: + r"""Maps linear operators on the coefficient spaces of ``fine`` to the corresponding operators on ``coarse``. + + Coarse leaves are cached (keyed by the fine leaf object), hence calling the coarsener again on an + operator that differs only in its scalars or composition (e.g. ``sigma * M0 + grad.T @ M1 @ grad`` + with a new ``sigma``) re-assembles nothing. + + Parameters + ---------- + fine, coarse : Derham + Fine and coarse level (same options, nested grids). + + domain : Domain + Mapping used for re-assembling mass matrices on the coarse level. + + matrix_free : bool + Whether coarse mass matrices are matrix-free. + """ + + def __init__(self, fine: Derham, coarse: Derham, domain: Domain, *, matrix_free: bool = False): + self._fine = fine + self._coarse = coarse + self._mass_ops = WeightedMassOperators(coarse, domain, matrix_free=matrix_free) + + # fine -> coarse coefficient spaces (also components of block spaces) + self._spaces: dict[int, tuple[VectorSpace, VectorSpace]] = {} + for form in ("0", "1", "2", "3", "v"): + Vf, Vc = fine.coeff_spaces[form], coarse.coeff_spaces[form] + self._spaces[id(Vf)] = (Vf, Vc) + if isinstance(Vf, BlockVectorSpace): + for vf, vc in zip(Vf.spaces, Vc.spaces): + self._spaces.setdefault(id(vf), (vf, vc)) + + self._cache: dict[int, tuple[LinearOperator, LinearOperator]] = {} + + @property + def fine(self) -> Derham: + return self._fine + + @property + def coarse(self) -> Derham: + return self._coarse + + @property + def mass_ops(self) -> WeightedMassOperators: + """Mass operators of the coarse level.""" + return self._mass_ops + + def space(self, V: VectorSpace) -> VectorSpace: + """Coarse counterpart of the fine coefficient space ``V``.""" + try: + return self._spaces[id(V)][1] + except KeyError: + raise ValueError(f"{V} is not a coefficient space of the fine Derham.") from None + + def __call__(self, A: LinearOperator) -> LinearOperator: + """Return the coarse-level version of the fine-level operator ``A``.""" + if isinstance(A, ScaledLinearOperator): + B = self(A.operator) + return ScaledLinearOperator(B.domain, B.codomain, A.scalar, B) + + if isinstance(A, SumLinearOperator): + addends = [self(a) for a in A.addends] + return SumLinearOperator(self.space(A.domain), self.space(A.codomain), *addends) + + if isinstance(A, ComposedLinearOperator): + factors = [self(a) for a in A.multiplicants] + return ComposedLinearOperator(self.space(A.domain), self.space(A.codomain), *factors) + + if isinstance(A, PowerLinearOperator): + return PowerLinearOperator(self.space(A.domain), self.space(A.codomain), self(A.operator), A.factorial) + + if isinstance(A, BlockLinearOperator): + blocks = {ij: self(A[ij]) for ij in A.nonzero_block_indices} + return BlockLinearOperator(self.space(A.domain), self.space(A.codomain), blocks=blocks) + + if isinstance(A, IdentityOperator): + return IdentityOperator(self.space(A.domain), self.space(A.codomain)) + + if isinstance(A, ZeroOperator): + return ZeroOperator(self.space(A.domain), self.space(A.codomain)) + + # leaves: cached + cached = self._cache.get(id(A)) + if cached is not None and cached[0] is A: + return cached[1] + + for cls in type(A).__mro__: + if cls in _REGISTRY: + B = _REGISTRY[cls](A, self) + break + else: + raise NotImplementedError( + f"Cannot re-discretize an operator of type {type(A).__name__} on a coarse grid; " + "use struphy operators or register a rule with register_coarsening." + ) + + assert B.domain is self.space(A.domain) and B.codomain is self.space(A.codomain), ( + f"Coarsening of {type(A).__name__} gave wrong (co)domain." + ) + self._cache[id(A)] = (A, B) + return B + + +@register_coarsening(DirectionalDerivativeOperator) +def _coarsen_derivative(A: DirectionalDerivativeOperator, c: OperatorCoarsener) -> LinearOperator: + return DirectionalDerivativeOperator( + c.space(A._spaceV), + c.space(A._spaceW), + A._diffdir, + negative=A._negative, + transposed=A._transposed, + ) + + +@register_coarsening(BoundaryOperator) +def _coarsen_boundary(A: BoundaryOperator, c: OperatorCoarsener) -> LinearOperator: + return BoundaryOperator(c.space(A.domain), A._space_id, A.bc) + + +@register_coarsening(WeightedMassOperator) +def _coarsen_mass(A: WeightedMassOperator, c: OperatorCoarsener) -> LinearOperator: + return WeightedMassOperator.from_dict(A.to_dict(), c.mass_ops) + + +@register_coarsening(BasisProjectionOperator) +def _coarsen_basis_projection(A: BasisProjectionOperator, c: OperatorCoarsener) -> LinearOperator: + return BasisProjectionOperator.from_dict(A.to_dict(), c.coarse) diff --git a/src/struphy/linear_algebra/multigrid/hierarchy.py b/src/struphy/linear_algebra/multigrid/hierarchy.py new file mode 100644 index 000000000..77bf8c8df --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/hierarchy.py @@ -0,0 +1,106 @@ +"""Hierarchy of nested Derham sequences for geometric multigrid.""" + +import logging + +import numpy as np +from feectools.ddm.cart import DomainDecomposition + +from struphy.feec.psydac_derham import Derham +from struphy.topology.grids import TensorProductGrid + +logger = logging.getLogger("struphy") + + +class MultiGridHierarchy: + r"""Nested Derham sequences :math:`V_0 \supset V_1 \supset \dots \supset V_{L-1}` obtained by + uniform coarsening of the user's (finest) Derham. + + From one level to the next, the number of elements is halved in every direction ``i`` where this + is possible, i.e. where + + * ``num_elements[i]`` is even, + * the coarse grid keeps at least ``max(min_cells, degree[i] + 1)`` elements, + * the MPI decomposition stays aligned: every process owns exactly the coarse elements covering + its fine elements (element starts and ends+1 are even), and owns at least ``degree[i]`` + coarse elements if the direction is split among several processes. + + Directions that cannot be halved are kept (semi-coarsening). Coarsening stops when no direction + can be halved or when ``max_levels`` is reached. All coarse Derhams share the communicator, + the process grid, the options (degree, boundary conditions, quadrature) and the domain of the finest one. + + Parameters + ---------- + derham : Derham + The finest level. + + max_levels : int | None + Maximal number of levels (including the finest); None means as many as possible. + + min_cells : int + Minimal number of elements per direction on the coarsest level. + """ + + def __init__(self, derham: Derham, *, max_levels: int | None = None, min_cells: int = 2): + if derham.polar_splines: + raise NotImplementedError("Multigrid is not yet implemented for polar splines.") + assert max_levels is None or max_levels >= 1 + assert min_cells >= 1 + + self._min_cells = min_cells + self._derhams: list[Derham] = [derham] + self._factors: list[tuple[int, int, int]] = [] + + while max_levels is None or len(self._derhams) < max_levels: + fine = self._derhams[-1] + factors = self._coarsening_factors(fine) + if all(f == 1 for f in factors): + break + + ddm = fine.domain_decomposition.coarsen(factors) + grid = TensorProductGrid( + num_elements=tuple(int(n) for n in ddm.ncells), + mpi_dims_mask=fine.grid.mpi_dims_mask, + ) + coarse = Derham(grid, fine.options, comm=fine.comm, domain=fine.domain, domain_decomposition=ddm) + + self._derhams.append(coarse) + self._factors.append(factors) + + logger.debug(f"Multigrid hierarchy: {[d.num_elements for d in self._derhams]}") + + def _coarsening_factors(self, derham: Derham) -> tuple[int, int, int]: + """Return 2 in every direction that can be halved (see class docstring), else 1.""" + ddm: DomainDecomposition = derham.domain_decomposition + factors = [] + for axis in range(3): + n = derham.num_elements[axis] + p = derham.degree[axis] + starts = np.asarray(ddm.global_element_starts[axis]) + ends = np.asarray(ddm.global_element_ends[axis]) + ok = n % 2 == 0 and n // 2 >= max(self._min_cells, p + 1) + ok = ok and np.all(starts % 2 == 0) and np.all((ends + 1) % 2 == 0) + if ok and ddm.nprocs[axis] > 1: + ok = np.all((ends - starts + 1) // 2 >= p) + factors.append(2 if ok else 1) + return tuple(factors) + + @property + def derhams(self) -> list[Derham]: + """Derham of each level, ``derhams[0]`` is the finest.""" + return self._derhams + + @property + def n_levels(self) -> int: + """Number of levels (including the finest).""" + return len(self._derhams) + + @property + def factors(self) -> list[tuple[int, int, int]]: + """``factors[l]`` is the coarsening factor in each direction from level ``l`` to ``l+1``.""" + return self._factors + + def __getitem__(self, level: int) -> Derham: + return self._derhams[level] + + def __len__(self) -> int: + return self.n_levels diff --git a/src/struphy/linear_algebra/multigrid/preconditioner.py b/src/struphy/linear_algebra/multigrid/preconditioner.py new file mode 100644 index 000000000..9764cae70 --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/preconditioner.py @@ -0,0 +1,483 @@ +r"""Geometric multigrid V-cycle as a preconditioner, and a multigrid-preconditioned CG solver. + +Given a symmetric positive (semi-)definite operator :math:`A` on one space of a :class:`Derham`, +for example the Poisson operator :math:`\sigma \mathbb M^0 + \mathbb G^\top \mathbb M^1 \mathbb G`: + +1. A hierarchy of coarser Derhams is built (:class:`MultiGridHierarchy`). +2. :math:`A` is re-discretized on every level from its expression tree (:class:`OperatorCoarsener`). +3. One application of the preconditioner is a V-cycle with zero initial guess:: + + x_l = 0 + smooth(A_l, b_l, x_l) # n_pre times + x_l += P_l V_cycle(l+1, R_l (b_l - A_l x_l)) + smooth(A_l, b_l, x_l) # n_post times + + with an exact solve on the coarsest level. With the same linear, :math:`A`-symmetric smoother before and + after the coarse-grid correction, the V-cycle is a symmetric positive definite preconditioner for CG. +""" + +import logging +from dataclasses import dataclass +from typing import Literal + +import numpy as np +import scipy.linalg as sla +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import IdentityOperator, LinearOperator, Vector +from feectools.linalg.block import BlockVector +from feectools.linalg.solvers import inverse + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.preconditioner import MassMatrixPreconditioner +from struphy.feec.psydac_derham import Derham +from struphy.geometry.base import Domain +from struphy.io.options import OptionsBase +from struphy.linear_algebra.multigrid.coarsen import OperatorCoarsener +from struphy.linear_algebra.multigrid.hierarchy import MultiGridHierarchy +from struphy.linear_algebra.multigrid.smoothers import ( + ChebyshevSmoother, + DiagonalComputer, + JacobiSmoother, + KrylovSmoother, + Smoother, + _owned_slice, + _stencil_blocks, + inverse_diagonal, +) +from struphy.linear_algebra.multigrid.transfer import SplineProlongation +from struphy.utils.utils import check_option + +logger = logging.getLogger("struphy") + +OptsSmoother = Literal["chebyshev", "jacobi", "cg"] +OptsSmootherPrecond = Literal["mass", "jacobi", "identity"] +OptsCoarseSolver = Literal["direct", "cg"] +OptsNullspace = Literal["constants"] + + +@dataclass +class MultiGridOptions(OptionsBase): + r"""Options of :class:`MultiGridPreconditioner`. + + Parameters + ---------- + smoother : str + "chebyshev" (default), "jacobi" (damped) or "cg" (fixed number of PCG steps, non-linear). + + smoother_precond : str + Preconditioner inside the Chebyshev and CG smoothers: "mass" (Kronecker-product approximation of + the inverse mass matrix of the space, default), "jacobi" (inverse diagonal of the operator) or "identity". + + smoother_degree : int + Polynomial degree of the Chebyshev smoother, number of sweeps of the Jacobi smoother, + or number of iterations of the CG smoother. + + n_pre, n_post : int + Number of smoother applications before and after the coarse-grid correction. + + jacobi_omega : float + Damping factor of the Jacobi smoother. + + chebyshev_bounds : tuple[float, float] + Smoothing interval of the Chebyshev smoother relative to the estimated largest eigenvalue. + + eig_iter : int + Number of Lanczos steps for the eigenvalue estimate of the Chebyshev smoother. + + max_levels : int | None + Maximal number of levels (None: coarsen as long as possible). + + min_cells : int + Minimal number of elements per coarsened direction on the coarsest level. + + coarse_solver : str + "direct" (LU factorization of the assembled coarsest operator, replicated on all processes) + or "cg" (PCG with the smoother preconditioner, to tolerance ``coarse_tol``). + + coarse_tol : float + Relative tolerance of the "cg" coarse solver. + + nullspace : str | None + "constants" if the operator is singular with the constant functions in its kernel (e.g. the + Poisson operator on 0-forms with periodic or Neumann boundary conditions). + + matrix_free_mass : bool + Whether re-discretized mass matrices on coarse levels are matrix-free. + """ + + smoother: OptsSmoother = "chebyshev" + smoother_precond: OptsSmootherPrecond = "mass" + smoother_degree: int = 3 + n_pre: int = 1 + n_post: int = 1 + jacobi_omega: float = 2.0 / 3.0 + chebyshev_bounds: tuple[float, float] = (0.1, 1.1) + eig_iter: int = 15 + max_levels: int | None = None + min_cells: int = 2 + coarse_solver: OptsCoarseSolver = "direct" + coarse_tol: float = 1e-10 + nullspace: OptsNullspace | None = None + matrix_free_mass: bool = False + + def __post_init__(self): + check_option(self.smoother, OptsSmoother) + check_option(self.smoother_precond, OptsSmootherPrecond) + check_option(self.coarse_solver, OptsCoarseSolver) + if self.nullspace is not None: + check_option(self.nullspace, OptsNullspace) + assert self.smoother_degree >= 1 + assert self.n_pre >= 0 and self.n_post >= 0 and self.n_pre + self.n_post >= 1 + + +class MultiGridPreconditioner(LinearOperator): + r"""One geometric multigrid V-cycle as an approximate inverse of ``A`` (see module docstring). + + Parameters + ---------- + A : LinearOperator + Symmetric positive (semi-)definite operator on ``derham.coeff_spaces[form]``, built from struphy + operators (Derham derivatives, boundary operators, :class:`WeightedMassOperator`, ...) with + ``+``, ``-``, ``*`` and ``@``. + + derham : Derham + The finest level. + + domain : Domain + Mapping, used to re-discretize mass matrices on the coarse levels. + + options : MultiGridOptions | None + Options (default: ``MultiGridOptions()``). + + mass_ops : WeightedMassOperators | None + Mass operators of the finest level (only used by the "mass" smoother preconditioner). + """ + + def __init__( + self, + A: LinearOperator, + derham: Derham, + domain: Domain, + options: MultiGridOptions | None = None, + *, + mass_ops: WeightedMassOperators | None = None, + ): + self._options = MultiGridOptions() if options is None else options + opts = self._options + self._derham = derham + self._domain_map = domain + + self._form = _find_form(derham, A.domain) + assert A.codomain is A.domain, "MultiGridPreconditioner requires a square operator." + if opts.nullspace == "constants": + assert self._form == "0", "nullspace='constants' is only implemented for 0-forms." + + self._hierarchy = MultiGridHierarchy(derham, max_levels=opts.max_levels, min_cells=opts.min_cells) + L = self._hierarchy.n_levels + logger.info(f"Multigrid levels: {[d.num_elements for d in self._hierarchy.derhams]}") + + self._P = [SplineProlongation(self._hierarchy[l + 1], self._hierarchy[l], self._form) for l in range(L - 1)] + self._R = [P.T for P in self._P] + self._coarseners = [ + OperatorCoarsener(self._hierarchy[l], self._hierarchy[l + 1], domain, matrix_free=opts.matrix_free_mass) + for l in range(L - 1) + ] + self._diag = [DiagonalComputer(d.degree) for d in self._hierarchy.derhams] + + if opts.smoother_precond == "mass": + fine_mass = WeightedMassOperators(derham, domain) if mass_ops is None else mass_ops + mass = [fine_mass] + [c.mass_ops for c in self._coarseners] + self._mass_pc = [MassMatrixPreconditioner(getattr(m, "M" + self._form)) for m in mass] + + # work vectors per level + self._b = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + self._x = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + self._r = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + self._e = [d.coeff_spaces[self._form].zeros() for d in self._hierarchy.derhams] + + self._A: list[LinearOperator] = [] + self.update(A) + + # ------------------------------------------------------------------ + @property + def domain(self): + return self._A[0].domain + + @property + def codomain(self): + return self._A[0].codomain + + @property + def dtype(self): + return self._A[0].dtype + + @property + def options(self) -> MultiGridOptions: + return self._options + + @property + def hierarchy(self) -> MultiGridHierarchy: + return self._hierarchy + + @property + def operators(self) -> list[LinearOperator]: + """System operator on each level, ``operators[0]`` is the given one.""" + return self._A + + @property + def smoothers(self) -> list[Smoother]: + return self._smoothers + + def transpose(self, conjugate: bool = False) -> "MultiGridPreconditioner": + assert all(s.is_symmetric for s in self._smoothers), "Only a symmetric V-cycle can be transposed." + assert self._options.n_pre == self._options.n_post + return self + + # ------------------------------------------------------------------ + def update(self, A: LinearOperator) -> None: + """Set a new fine-level operator (e.g. with changed scalars) and update all levels. + + Mass matrices and derivative operators of coarse levels are re-used if ``A`` is built from the same objects. + """ + assert A.domain is self._derham.coeff_spaces[self._form] + self._A = [A] + for c in self._coarseners: + self._A.append(c(self._A[-1])) + self._smoothers = [self._make_smoother(l) for l in range(self._hierarchy.n_levels - 1)] + self._setup_coarse_solver() + + def _smoother_precond(self, l: int) -> LinearOperator: + opts = self._options + if opts.smoother_precond == "mass": + return self._mass_pc[l] + if opts.smoother_precond == "jacobi": + return inverse_diagonal(self._diag[l](self._A[l])) + return IdentityOperator(self._A[l].domain) + + def _make_smoother(self, l: int) -> Smoother: + opts = self._options + A = self._A[l] + if opts.smoother == "chebyshev": + return ChebyshevSmoother( + A, + self._smoother_precond(l), + degree=opts.smoother_degree, + bounds=opts.chebyshev_bounds, + eig_iter=opts.eig_iter, + ) + if opts.smoother == "jacobi": + D_inv = inverse_diagonal(self._diag[l](A)) + return JacobiSmoother(A, D_inv, omega=opts.jacobi_omega, sweeps=opts.smoother_degree) + return KrylovSmoother(A, self._smoother_precond(l), iterations=opts.smoother_degree) + + def _setup_coarse_solver(self) -> None: + opts = self._options + A = self._A[-1] + if opts.coarse_solver == "cg": + pc = self._smoother_precond(len(self._A) - 1) if opts.smoother_precond != "identity" else None + self._coarse_cg = inverse(A, "pcg", pc=pc, tol=1e-300, maxiter=1000, recycle=False) + return + + Ad = _assemble_dense(A) + if opts.nullspace == "constants": + # A + s 1 1^T is regular; for b orthogonal to 1 its solution is the zero-mean solution of A x = b + n = Ad.shape[0] + Ad = Ad + np.mean(np.abs(np.diag(Ad))) / n * np.ones((n, n)) + else: + # rows/cols of Dirichlet dofs are zero: put ones on the diagonal + zero = np.flatnonzero(np.all(Ad == 0.0, axis=1)) + Ad[zero, zero] = 1.0 + self._coarse_lu = sla.lu_factor(Ad) + + # ------------------------------------------------------------------ + def dot(self, b: Vector, out: Vector | None = None) -> Vector: + """Apply one V-cycle (zero initial guess) to ``b``.""" + assert b.space is self.domain + if out is None: + out = self.domain.zeros() + b.copy(out=self._b[0]) + if self._options.nullspace == "constants": + _remove_mean(self._b[0]) + self._vcycle(0) + self._x[0].copy(out=out) + if self._options.nullspace == "constants": + _remove_mean(out) + return out + + def _vcycle(self, l: int) -> None: + b, x = self._b[l], self._x[l] + if l == len(self._A) - 1: + self._coarse_solve(b, x) + return + + x *= 0.0 + S = self._smoothers[l] + for _ in range(self._options.n_pre): + S.smooth(b, x) + + r = S.residual(b, x, self._r[l]) + self._R[l].dot(r, out=self._b[l + 1]) + self._vcycle(l + 1) + self._P[l].dot(self._x[l + 1], out=self._e[l]) + x += self._e[l] + + for _ in range(self._options.n_post): + S.smooth(b, x) + + def _coarse_solve(self, b: Vector, x: Vector) -> None: + if self._options.coarse_solver == "cg": + nb = np.sqrt(b.inner(b)) + if nb == 0.0: + x *= 0.0 + return + self._coarse_cg._options["tol"] = self._options.coarse_tol * nb + self._coarse_cg.dot(b, out=x) + return + + bg = _gather(b) + if self._options.nullspace == "constants": + bg -= bg.mean() + _scatter(sla.lu_solve(self._coarse_lu, bg), x) + + +class MultiGridSolver(LinearOperator): + r"""Conjugate gradient method preconditioned with :class:`MultiGridPreconditioner`. + + Parameters + ---------- + A, derham, domain, options, mass_ops : + See :class:`MultiGridPreconditioner`. + + tol : float + Relative tolerance, the iteration stops when :math:`\|b - A x\|_2 \leq \mathrm{tol}\, \|b\|_2`. + + maxiter : int + Maximal number of CG iterations. + + verbose : bool + Print the residual in every iteration. + """ + + def __init__( + self, + A: LinearOperator, + derham: Derham, + domain: Domain, + options: MultiGridOptions | None = None, + *, + mass_ops: WeightedMassOperators | None = None, + tol: float = 1e-8, + maxiter: int = 100, + verbose: bool = False, + ): + self._pc = MultiGridPreconditioner(A, derham, domain, options, mass_ops=mass_ops) + self._tol = tol + self._solver = inverse(A, "pcg", pc=self._pc, tol=tol, maxiter=maxiter, verbose=verbose, recycle=False) + + @property + def domain(self): + return self._pc.domain + + @property + def codomain(self): + return self._pc.codomain + + @property + def dtype(self): + return self._pc.dtype + + @property + def preconditioner(self) -> MultiGridPreconditioner: + return self._pc + + @property + def info(self) -> dict: + """Information of the last solve: ``niter``, ``success``, ``res_norm``.""" + return self._solver._info + + def update(self, A: LinearOperator) -> None: + """Set a new operator (see :meth:`MultiGridPreconditioner.update`).""" + self._pc.update(A) + self._solver.linop = A + + def transpose(self, conjugate: bool = False) -> "MultiGridSolver": + return self + + def dot(self, b: Vector, out: Vector | None = None, x0: Vector | None = None) -> Vector: + """Solve ``A x = b`` (initial guess ``x0``, zero by default).""" + nb = np.sqrt(b.inner(b)) + self._solver._options["tol"] = self._tol * nb if nb > 0.0 else self._tol + self._solver._options["x0"] = x0 if x0 is not None else self.domain.zeros() + return self._solver.dot(b, out=out) + + +# ---------------------------------------------------------------------------------------------------- +def _find_form(derham: Derham, V) -> str: + for form in ("0", "1", "2", "3", "v"): + if derham.coeff_spaces[form] is V: + return form + raise ValueError("The operator does not act on a coefficient space of the given Derham.") + + +def _comm(v: Vector): + blk = _stencil_blocks(v)[0] + return blk.space.cart.comm if blk.space.parallel else None + + +def _gather(v: Vector) -> np.ndarray: + """Global coefficient array of ``v`` (same on all processes).""" + a = v.toarray() + comm = _comm(v) + if comm is not None and comm.size > 1: + comm.Allreduce(MPI.IN_PLACE, a, op=MPI.SUM) + return a + + +def _scatter(a: np.ndarray, v: Vector) -> None: + """Write the owned part of the global array ``a`` into ``v``.""" + offset = 0 + for blk in _stencil_blocks(v): + V = blk.space + n = int(np.prod(V.npts)) + glob = a[offset : offset + n].reshape(tuple(int(m) for m in V.npts)) + blk._data[_owned_slice(blk)] = glob[tuple(slice(s, e + 1) for s, e in zip(V.starts, V.ends))] + blk.ghost_regions_in_sync = False + offset += n + + +def _assemble_dense(A: LinearOperator) -> np.ndarray: + """Global dense matrix of ``A`` (same on all processes), by applying it to all unit vectors.""" + e = A.domain.zeros() + blocks = _stencil_blocks(e) + sizes = [int(np.prod(b.space.npts)) for b in blocks] + N = sum(sizes) + out = np.zeros((N, N)) + y = A.codomain.zeros() + col = 0 + for blk, n in zip(blocks, sizes): + V = blk.space + for flat in range(n): + gidx = np.unravel_index(flat, tuple(int(m) for m in V.npts)) + for b in blocks: + b._data[...] = 0.0 + b.ghost_regions_in_sync = False + if all(s <= i <= e_ for i, s, e_ in zip(gidx, V.starts, V.ends)): + loc = tuple(int(i - s + p * m) for i, s, p, m in zip(gidx, V.starts, V.pads, V.shifts)) + blk._data[loc] = 1.0 + A.dot(e, out=y) + out[:, col] = _gather(y) + col += 1 + return out + + +def _remove_mean(v: Vector) -> None: + """Subtract the mean of all coefficients (the projection orthogonal to the constant vector).""" + comm = _comm(v) + s = sum(float(np.sum(b._data[_owned_slice(b)])) for b in _stencil_blocks(v)) + N = sum(int(np.prod(b.space.npts)) for b in _stencil_blocks(v)) + if comm is not None and comm.size > 1: + s = comm.allreduce(s, op=MPI.SUM) + mean = s / N + for b in _stencil_blocks(v): + b._data[_owned_slice(b)] -= mean + b.ghost_regions_in_sync = False diff --git a/src/struphy/linear_algebra/multigrid/smoothers.py b/src/struphy/linear_algebra/multigrid/smoothers.py new file mode 100644 index 000000000..67a684332 --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/smoothers.py @@ -0,0 +1,425 @@ +r"""Smoothers for geometric multigrid. + +A smoother approximately solves :math:`A x = b` by the update :math:`x \leftarrow x + S (b - A x)`. +:class:`ChebyshevSmoother` and :class:`JacobiSmoother` are linear and :math:`A`-symmetric, hence a V-cycle +with the same smoother before and after the coarse-grid correction is a symmetric preconditioner, suitable +for the conjugate gradient method. :class:`KrylovSmoother` is non-linear (use with care). +""" + +from abc import ABC, abstractmethod + +import numpy as np +from feectools.linalg.basic import ( + IdentityOperator, + LinearOperator, + ScaledLinearOperator, + SumLinearOperator, + Vector, + ZeroOperator, +) +from feectools.linalg.block import BlockVector +from feectools.linalg.stencil import StencilVector + + +class Smoother(ABC): + """Base class of multigrid smoothers for the operator ``A``.""" + + def __init__(self, A: LinearOperator): + assert A.domain is A.codomain + self._A = A + self._r = A.codomain.zeros() + + @property + def A(self) -> LinearOperator: + return self._A + + @property + def is_symmetric(self) -> bool: + """Whether the smoother is linear and A-symmetric (needed for a symmetric V-cycle).""" + return True + + def residual(self, b: Vector, x: Vector, out: Vector) -> Vector: + """``out = b - A x``.""" + self._A.dot(x, out=out) + out *= -1.0 + out += b + return out + + @abstractmethod + def smooth(self, b: Vector, x: Vector) -> None: + """Improve the approximate solution ``x`` of ``A x = b`` in place.""" + + +class JacobiSmoother(Smoother): + r"""Damped Jacobi, :math:`x \leftarrow x + \omega D^{-1}(b - A x)`, repeated ``sweeps`` times. + + Parameters + ---------- + A : LinearOperator + System operator. + + diag_inv : LinearOperator + Inverse diagonal :math:`D^{-1}` of ``A`` (e.g. from :func:`inverse_diagonal`). + + omega : float + Damping factor. + + sweeps : int + Number of iterations per call. + """ + + def __init__(self, A: LinearOperator, diag_inv: LinearOperator, *, omega: float = 2.0 / 3.0, sweeps: int = 1): + super().__init__(A) + self._D_inv = diag_inv + self._omega = omega + self._sweeps = sweeps + self._z = A.domain.zeros() + + def smooth(self, b: Vector, x: Vector) -> None: + for _ in range(self._sweeps): + self.residual(b, x, self._r) + self._D_inv.dot(self._r, out=self._z) + x.mul_iadd(self._omega, self._z) + + +class ChebyshevSmoother(Smoother): + r"""Chebyshev polynomial smoother of degree ``degree`` for :math:`M^{-1} A`. + + The polynomial damps the eigenmodes of :math:`M^{-1} A` in the interval + :math:`[\alpha \lambda_\max, \beta \lambda_\max]` (``bounds = (alpha, beta)``), where + :math:`\lambda_\max` is estimated with a few Lanczos (PCG) steps. Each call costs ``degree`` + applications of ``A`` and of ``M_inv``. No inner products are computed. + + Parameters + ---------- + A : LinearOperator + Symmetric positive (semi-)definite system operator. + + M_inv : LinearOperator + Symmetric positive definite preconditioner, e.g. an approximate mass-matrix inverse or + an inverse diagonal of ``A``. + + degree : int + Polynomial degree. + + bounds : tuple[float, float] + Lower and upper end of the smoothing interval relative to the estimated :math:`\lambda_\max`. + + eig_iter : int + Number of Lanczos steps for estimating :math:`\lambda_\max`. + + lambda_max : float | None + Largest eigenvalue of :math:`M^{-1} A`, estimated if None. + """ + + def __init__( + self, + A: LinearOperator, + M_inv: LinearOperator, + *, + degree: int = 3, + bounds: tuple[float, float] = (0.1, 1.1), + eig_iter: int = 15, + lambda_max: float | None = None, + ): + super().__init__(A) + assert degree >= 1 + assert 0.0 < bounds[0] < bounds[1] + self._M_inv = M_inv + self._degree = degree + + if lambda_max is None: + lambda_max = estimate_lambda_max(A, M_inv, n_iter=eig_iter) + self._lambda_max = lambda_max + a, b = bounds[0] * lambda_max, bounds[1] * lambda_max + self._theta = 0.5 * (b + a) + self._delta = 0.5 * (b - a) + + self._d = A.domain.zeros() + self._z = A.domain.zeros() + + @property + def lambda_max(self) -> float: + return self._lambda_max + + def smooth(self, b: Vector, x: Vector) -> None: + # Saad, Iterative Methods for Sparse Linear Systems, Alg. 12.1 (preconditioned) + theta, delta = self._theta, self._delta + sigma = theta / delta + rho = 1.0 / sigma + r, d, z = self._r, self._d, self._z + + self.residual(b, x, r) + self._M_inv.dot(r, out=d) + d *= 1.0 / theta + for k in range(self._degree): + x += d + if k == self._degree - 1: + break + self._A.dot(d, out=z) + r -= z + rho_new = 1.0 / (2.0 * sigma - rho) + self._M_inv.dot(r, out=z) + d *= rho_new * rho + d.mul_iadd(2.0 * rho_new / delta, z) + rho = rho_new + + +class KrylovSmoother(Smoother): + r"""A fixed number of (preconditioned) conjugate gradient iterations, warm-started from ``x``. + + Each call to :meth:`smooth` performs exactly ``iterations`` PCG steps for :math:`A x = b`, + starting from the current ``x``. No convergence tolerance is used; the iteration stops early only on + breakdown, i.e. when :math:`r^\top M^{-1} r = 0` (e.g. zero residual) or :math:`p^\top A p = 0`. + + Note that this smoother is non-linear; a V-cycle using it is not a fixed linear preconditioner. + + Parameters + ---------- + A : LinearOperator + System operator (symmetric positive definite). + + M_inv : LinearOperator | None + Preconditioner :math:`M^{-1}` (symmetric positive definite). If None, the identity is used. + + iterations : int + Number of PCG steps per call (at least 1). + """ + + def __init__(self, A: LinearOperator, M_inv: LinearOperator | None = None, *, iterations: int = 3): + super().__init__(A) + assert iterations >= 1, f"KrylovSmoother needs at least one iteration, got {iterations}." + self._M_inv = IdentityOperator(A.domain) if M_inv is None else M_inv + self._iterations = iterations + self._z = A.domain.zeros() + self._p = A.domain.zeros() + self._q = A.domain.zeros() + + @property + def is_symmetric(self) -> bool: + return False + + def smooth(self, b: Vector, x: Vector) -> None: + """Perform ``iterations`` PCG steps for ``A x = b``, updating ``x`` in place.""" + r, z, p, q = self._r, self._z, self._p, self._q + self.residual(b, x, r) + self._M_inv.dot(r, out=z) + z.copy(out=p) + rz = r.inner(z) + for k in range(self._iterations): + # inner products are global reductions, hence all ranks break consistently + if rz == 0.0: + break + self._A.dot(p, out=q) + pq = p.inner(q) + if pq == 0.0: + break + alpha = rz / pq + x.mul_iadd(alpha, p) + if k == self._iterations - 1: + break + r.mul_iadd(-alpha, q) + self._M_inv.dot(r, out=z) + rz_new = r.inner(z) + p *= rz_new / rz + p += z + rz = rz_new + + +def estimate_lambda_max(A: LinearOperator, M_inv: LinearOperator, *, n_iter: int = 15, seed: int = 1234) -> float: + r"""Estimate the largest eigenvalue of :math:`M^{-1} A` with ``n_iter`` Lanczos steps (via PCG coefficients). + + The estimate is a lower bound that converges quickly to :math:`\lambda_\max`. + """ + b = A.domain.zeros() + _fill_random(b, seed) + + x_r = b.copy() + z = M_inv.dot(x_r) + p = z.copy() + q = A.domain.zeros() + rz = x_r.inner(z) + alphas, betas = [], [] + for _ in range(n_iter): + A.dot(p, out=q) + pq = p.inner(q) + if pq <= 0.0 or rz <= 0.0: + break + alpha = rz / pq + x_r.mul_iadd(-alpha, q) + M_inv.dot(x_r, out=z) + rz_new = x_r.inner(z) + beta = rz_new / rz + alphas.append(alpha) + betas.append(beta) + if rz_new <= 1e-30 * rz: + break + p *= beta + p += z + rz = rz_new + + k = len(alphas) + assert k > 0, "Lanczos breakdown in the eigenvalue estimate (is A positive semi-definite?)." + T = np.zeros((k, k)) + for j in range(k): + T[j, j] = 1.0 / alphas[j] + (betas[j - 1] / alphas[j - 1] if j > 0 else 0.0) + if j + 1 < k: + T[j, j + 1] = T[j + 1, j] = np.sqrt(betas[j]) / alphas[j] + return float(np.linalg.eigvalsh(T).max()) + + +def _stencil_blocks(v: Vector) -> list[StencilVector]: + return list(v.blocks) if isinstance(v, BlockVector) else [v] + + +def _fill_random(v: Vector, seed: int) -> None: + """Fill the owned entries of ``v`` with uniform random numbers in [-1, 1] (different on each process).""" + for n, blk in enumerate(_stencil_blocks(v)): + V = blk.space + rng = np.random.default_rng([seed, n] + [int(s) for s in V.starts]) + idx = _owned_slice(blk) + blk._data[idx] = rng.uniform(-1.0, 1.0, size=blk._data[idx].shape) + blk.ghost_regions_in_sync = False + + +def _owned_slice(v: StencilVector) -> tuple[slice, ...]: + V = v.space + return tuple(slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(V.pads, V.shifts, V.starts, V.ends)) + + +# ---------------------------------------------------------------------------------------------------- +# diagonal of composite operators +# ---------------------------------------------------------------------------------------------------- +class DiagonalOperator(LinearOperator): + """Pointwise multiplication by the entries of a vector.""" + + def __init__(self, d: Vector): + self._d = d + self._space = d.space + + @property + def domain(self): + return self._space + + @property + def codomain(self): + return self._space + + @property + def dtype(self): + return self._space.dtype + + @property + def vector(self) -> Vector: + return self._d + + def dot(self, v: Vector, out: Vector | None = None) -> Vector: + if out is None: + out = self._space.zeros() + for vb, db, ob in zip(_stencil_blocks(v), _stencil_blocks(self._d), _stencil_blocks(out)): + idx = _owned_slice(ob) + ob._data[idx] = db._data[idx] * vb._data[idx] + ob.ghost_regions_in_sync = False + return out + + def transpose(self, conjugate: bool = False) -> "DiagonalOperator": + return self + + +class DiagonalComputer: + r"""Exact diagonal of (composite) operators, cached per operator object. + + Sums and scalings are combined from the diagonals of their parts; identity and zero operators are + trivial; every other operator (a leaf, or a product such as :math:`G^\top M G`) is probed with + colored unit vectors: dofs with equal index modulo :math:`c_d \geq 2 w_d + 1` in every direction do + not couple, so one application of the operator per color gives the diagonal entries of all dofs of + that color. This costs :math:`\prod_d c_d` operator applications (times the number of components). + + Parameters + ---------- + widths : tuple[int, int, int] + Upper bound for the coupling distance (in index units) of the operators in each direction, + e.g. the spline degrees for mass and stiffness matrices. + """ + + def __init__(self, widths: tuple[int, int, int]): + self._widths = tuple(widths) + self._cache: dict[int, tuple[LinearOperator, Vector]] = {} + + def __call__(self, A: LinearOperator) -> Vector: + """Diagonal of ``A`` as a vector in ``A.domain``.""" + assert A.domain is A.codomain + if isinstance(A, SumLinearOperator): + d = A.domain.zeros() + for a in A.addends: + d += self(a) + return d + if isinstance(A, ScaledLinearOperator): + return self(A.operator) * A.scalar + if isinstance(A, IdentityOperator): + d = A.domain.zeros() + for blk in _stencil_blocks(d): + blk._data[...] = 1.0 + return d + if isinstance(A, ZeroOperator): + return A.domain.zeros() + + cached = self._cache.get(id(A)) + if cached is not None and cached[0] is A: + return cached[1] + d = self._probe(A) + self._cache[id(A)] = (A, d) + return d + + def _probe(self, A: LinearOperator) -> Vector: + d = A.domain.zeros() + e = A.domain.zeros() + y = A.domain.zeros() + d_blocks, e_blocks = _stencil_blocks(d), _stencil_blocks(e) + + for n, (db, eb) in enumerate(zip(d_blocks, e_blocks)): + V = eb.space + colors = [ + _n_colors(int(npts), 2 * w + 1, bool(per)) for npts, w, per in zip(V.npts, self._widths, V.periods) + ] + glob = [np.arange(s, e + 1) for s, e in zip(V.starts, V.ends)] + idx = _owned_slice(eb) + for color in np.ndindex(*colors): + mask = np.ones([len(g) for g in glob], dtype=bool) + for axis, (g, c, k) in enumerate(zip(glob, colors, color)): + shape = [1, 1, 1] + shape[axis] = len(g) + mask = mask & (g % c == k).reshape(shape) + for blk in e_blocks: + blk._data[...] = 0.0 + blk.ghost_regions_in_sync = False + eb._data[idx] = mask.astype(float) + A.dot(e, out=y) + yb = _stencil_blocks(y)[n] + db._data[idx][mask] = yb._data[idx][mask] + return d + + +def _n_colors(n: int, c: int, periodic: bool) -> int: + """Number of colors in one direction with ``n`` dofs such that equal colors are at least ``c`` apart.""" + if c >= n: + return n + if not periodic: + return c + # periodic: c must divide n to avoid coupling across the periodic boundary + for k in range(c, n + 1): + if n % k == 0: + return k + return n + + +def inverse_diagonal(diag: Vector) -> DiagonalOperator: + """Inverse of a diagonal (entries equal to zero, e.g. Dirichlet dofs, are mapped to zero).""" + inv = diag.copy() + for blk in _stencil_blocks(inv): + idx = _owned_slice(blk) + data = blk._data[idx] + out = np.zeros_like(data) + np.divide(1.0, data, out=out, where=data != 0.0) + blk._data[idx] = out + return DiagonalOperator(inv) diff --git a/src/struphy/linear_algebra/multigrid/transfer.py b/src/struphy/linear_algebra/multigrid/transfer.py new file mode 100644 index 000000000..68299a116 --- /dev/null +++ b/src/struphy/linear_algebra/multigrid/transfer.py @@ -0,0 +1,257 @@ +r"""Grid-transfer operators between nested spline spaces. + +On uniformly refined grids the spline spaces are nested, :math:`V_H \subset V_h`, hence every coarse +basis function is a linear combination of fine ones. The prolongation :math:`P: V_H \to V_h` maps the +coefficients of a coarse spline to the coefficients of the *same* function in the fine basis; the +restriction is its transpose :math:`R = P^\top`. Both are tensor products of 1D matrices, applied +component-wise for vector-valued spaces. +""" + +import numpy as np +import scipy.sparse as spa +from feectools.core.bsplines import collocation_matrix +from feectools.fem.splines import SplineSpace +from feectools.fem.tensor import TensorFemSpace +from feectools.linalg.basic import IdentityOperator, LinearOperator, Vector +from feectools.linalg.block import BlockVector +from feectools.linalg.stencil import StencilVector, StencilVectorSpace + +from struphy.feec.psydac_derham import Derham + + +def prolongation_matrix_1d(coarse: SplineSpace, fine: SplineSpace) -> np.ndarray: + r"""Dense 1D prolongation matrix :math:`P \in \mathbb R^{n_h \times n_H}` between nested spline spaces. + + Column ``j`` holds the fine coefficients of the coarse basis function ``j``, i.e. + :math:`\Lambda^H_j = \sum_i P_{ij} \Lambda^h_i`. It is computed by collocation at ``degree + 1`` + points per fine cell (an overdetermined, exactly solvable system). Works for periodic and clamped + splines and for both normalizations (B-splines and D-splines/M-splines). + + Parameters + ---------- + coarse, fine : SplineSpace + 1D spaces of the same degree, kind and normalization; the breaks of ``coarse`` are a subset of those of ``fine``. + """ + assert coarse.degree == fine.degree + assert coarse.periodic == fine.periodic + assert coarse.basis == fine.basis + + if coarse.ncells == fine.ncells: + return np.eye(fine.nbasis) + + breaks = np.asarray(fine.breaks) + nq = fine.degree + 1 + s = (np.arange(nq) + 0.5) / nq + x = (breaks[:-1, None] + np.diff(breaks)[:, None] * s[None, :]).ravel() + + Bh = collocation_matrix(fine.knots, fine.degree, fine.periodic, fine.basis, x) + BH = collocation_matrix(coarse.knots, coarse.degree, coarse.periodic, coarse.basis, x) + P = np.linalg.lstsq(Bh, BH, rcond=None)[0] + + assert np.allclose(Bh @ P, BH, atol=1e-10), "Spline spaces are not nested." + P[np.abs(P) < 1e-13 * np.abs(P).max()] = 0.0 + return P + + +def local_matrix_1d( + A: np.ndarray, + out_start: int, + out_end: int, + in_start: int, + in_end: int, + in_ghost: int, + periodic: bool, +) -> spa.csr_matrix: + r"""Restrict a global 1D matrix to the rows owned by this process and to local (ghosted) input columns. + + Row ``r`` of the result is row ``out_start + r`` of ``A``. Global column ``c`` is mapped to the index + of the local ghosted input array, ``c - in_start + in_ghost``, using periodic images if ``periodic``. + If a column is present several times (owned and as ghost) the owned copy is used. + + Parameters + ---------- + A : numpy.ndarray + Global matrix of shape ``(n_out, n_in)``. + + out_start, out_end : int + Global indices of the first and last output entry owned by this process. + + in_start, in_end : int + Global indices of the first and last input entry owned by this process. + + in_ghost : int + Width of the ghost region on each side of the local input array (``pads * shifts``). + + periodic : bool + Whether the input index is periodic. + """ + n_in = A.shape[1] + n_loc = in_end - in_start + 1 + 2 * in_ghost + + rows, cols, vals = [], [], [] + for r, row in enumerate(range(out_start, out_end + 1)): + for c in np.flatnonzero(A[row]): + images = [c - n_in, c, c + n_in] if periodic else [c] + local = [g - in_start + in_ghost for g in images] + owned = [l for l in local if in_ghost <= l < n_loc - in_ghost] + ghost = [l for l in local if 0 <= l < n_loc] + if owned: + loc = owned[0] + elif ghost: + loc = ghost[0] + else: + raise ValueError(f"Column {c} of row {row} is outside the local ghost region.") + rows.append(r) + cols.append(loc) + vals.append(A[row, c]) + + return spa.csr_matrix((vals, (rows, cols)), shape=(out_end - out_start + 1, n_loc)) + + +class _KronTransfer: + """Tensor product of three local 1D matrices mapping a ghosted input StencilVector to the owned part of the output.""" + + def __init__(self, mats: list[spa.csr_matrix], W: StencilVectorSpace): + self._mats = mats + self._out_slice = tuple( + slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(W.pads, W.shifts, W.starts, W.ends) + ) + + def dot(self, v: StencilVector, out: StencilVector) -> None: + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + x = v._data + for axis, L in enumerate(self._mats): + x = np.moveaxis(x, axis, 0) + shp = x.shape + x = (L @ x.reshape(shp[0], -1)).reshape((L.shape[0],) + shp[1:]) + x = np.moveaxis(x, 0, axis) + out._data[...] = 0.0 + out._data[self._out_slice] = x + out.ghost_regions_in_sync = False + + +def _scalar_spaces(V) -> list[TensorFemSpace]: + """Scalar components of a (vector) FEM space.""" + return [V] if isinstance(V, TensorFemSpace) else list(V.spaces) + + +def _coeff_spaces(W) -> list[StencilVectorSpace]: + """Scalar components of a (block) coefficient space.""" + return [W] if isinstance(W, StencilVectorSpace) else list(W.spaces) + + +class SplineProlongation(LinearOperator): + r"""Prolongation :math:`P: V_H \to V_h` (or, with ``transposed=True``, restriction :math:`R = P^\top`) + between the same space of two nested Derham sequences. + + The MPI decompositions must be aligned (each process owns the coarse elements covering its fine + elements), as produced by :meth:`DomainDecomposition.coarsen`. With homogeneous Dirichlet boundary + conditions, the operator is :math:`\mathbb B_h P \mathbb B_H^\top` (resp. its transpose), which is the + exact embedding of the coarse into the fine space with boundary conditions. + + Parameters + ---------- + coarse, fine : Derham + Coarse and fine level. + + space_id : str + Space key of ``Derham.fem_spaces`` ("0", "1", "2", "3", "v" or "H1", "Hcurl", "Hdiv", "L2", "H1vec"). + + transposed : bool + If True, the restriction :math:`R = P^\top` (fine to coarse) is created. + """ + + def __init__(self, coarse: Derham, fine: Derham, space_id: str, *, transposed: bool = False): + if coarse.polar_splines or fine.polar_splines: + raise NotImplementedError("Grid transfer is not yet implemented for polar splines.") + + self._coarse = coarse + self._fine = fine + self._space_id = space_id + self._transposed = transposed + + form = coarse.space_to_form.get(space_id, space_id) + self._form = form + VH = coarse.fem_spaces[form] + Vh = fine.fem_spaces[form] + + self._Bc = coarse.boundary_ops[form] + self._Bf = fine.boundary_ops[form] + self._apply_bc = not (isinstance(self._Bc, IdentityOperator) and isinstance(self._Bf, IdentityOperator)) + + if transposed: + self._domain, self._codomain = fine.coeff_spaces[form], coarse.coeff_spaces[form] + V_in, V_out = Vh, VH + else: + self._domain, self._codomain = coarse.coeff_spaces[form], fine.coeff_spaces[form] + V_in, V_out = VH, Vh + + self._kron = [] + for cH, ch, Win, Wout in zip( + _scalar_spaces(VH), + _scalar_spaces(Vh), + _coeff_spaces(self._domain), + _coeff_spaces(self._codomain), + ): + mats = [] + for axis, (sH, sh) in enumerate(zip(cH.spaces, ch.spaces)): + P = prolongation_matrix_1d(sH, sh) + A = P.T if transposed else P + mats.append( + local_matrix_1d( + A, + int(Wout.starts[axis]), + int(Wout.ends[axis]), + int(Win.starts[axis]), + int(Win.ends[axis]), + int(Win.pads[axis] * Win.shifts[axis]), + sh.periodic, + ) + ) + self._kron.append(_KronTransfer(mats, Wout)) + + @property + def domain(self): + return self._domain + + @property + def codomain(self): + return self._codomain + + @property + def dtype(self): + return self._domain.dtype + + @property + def space_id(self) -> str: + return self._space_id + + @property + def transposed(self) -> bool: + return self._transposed + + def dot(self, v: Vector, out: Vector | None = None) -> Vector: + """Apply the operator, ``out = P v`` (or ``out = R v`` if transposed).""" + assert isinstance(v, Vector) and v.space == self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) and out.space == self.codomain + + B_in, B_out = (self._Bf, self._Bc) if self._transposed else (self._Bc, self._Bf) + if self._apply_bc: + v = B_in.dot(v) + + if isinstance(v, BlockVector): + for k, vk, ok in zip(self._kron, v.blocks, out.blocks): + k.dot(vk, ok) + else: + self._kron[0].dot(v, out) + + if self._apply_bc: + B_out.dot(out, out=out) + return out + + def transpose(self, conjugate: bool = False) -> "SplineProlongation": + return SplineProlongation(self._coarse, self._fine, self._space_id, transposed=not self._transposed) diff --git a/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py new file mode 100644 index 000000000..2b6ad256e --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_multigrid_coarsen.py @@ -0,0 +1,152 @@ +import json + +import numpy as np +import pytest +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.feec.utilities import create_equal_random_arrays +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.multigrid.coarsen import OperatorCoarsener +from struphy.linear_algebra.multigrid.hierarchy import MultiGridHierarchy +from struphy.linear_algebra.multigrid.transfer import SplineProlongation +from struphy.topology.grids import TensorProductGrid + +BCS = [ + (None, None, None), + (("dirichlet", "dirichlet"), None, ("free", "dirichlet")), +] + + +def _derham(num_elements, degree, bcs): + domain = Cuboid(l1=0.0, r1=2.0, l2=0.0, r2=1.0, l3=0.0, r3=3.0) + derham = Derham( + TensorProductGrid(num_elements=num_elements), + DerhamOptions(degree=degree, bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domain, + ) + return derham, domain + + +def _max_diff(a, b): + comm = MPI.COMM_WORLD + err = np.max(np.abs((a - b).toarray())) + ref = np.max(np.abs(b.toarray())) + if comm.Get_size() > 1: + err = comm.allreduce(err, op=MPI.MAX) + ref = comm.allreduce(ref, op=MPI.MAX) + return err / ref + + +@pytest.mark.parametrize("bcs", BCS) +def test_mass_to_dict(bcs): + derham, domain = _derham((8, 6, 4), (2, 2, 1), bcs) + mass_ops = WeightedMassOperators(derham, domain) + _, u = create_equal_random_arrays(derham.fem_spaces["1"], seed=1) + + # predefined operator with string weights: JSON serializable + M1 = mass_ops.M1 + dct = M1.to_dict() + assert dct["type"] == "WeightedMassOperator" + assert dct["params"]["weights"] == ["Ginv", "sqrt_g"] + json.dumps(dct) + assert _max_diff(WeightedMassOperator.from_dict(dct, mass_ops).dot(u), M1.dot(u)) < 1e-14 + + # callable weights and transposes + Mc = mass_ops.create_weighted_mass( + "Hcurl", "Hdiv", weights=("Ginv", lambda e1, e2, e3: 1.0 + e1 * e2), name="Mc", assemble=True + ) + McT = Mc.T + assert McT.to_dict()["params"]["is_transpose"] + _, w = create_equal_random_arrays(derham.fem_spaces["2"], seed=2) + assert _max_diff(WeightedMassOperator.from_dict(McT.to_dict(), mass_ops).dot(w), McT.dot(w)) < 1e-14 + + # constant 3x3 matrix as first tuple entry (must not be mistaken for block weights) + Mm = mass_ops.create_weighted_mass( + "Hcurl", "Hcurl", weights=([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]], "sqrt_g"), assemble=True + ) + dct = json.loads(json.dumps(Mm.to_dict())) + assert _max_diff(WeightedMassOperator.from_dict(dct, mass_ops).dot(u), Mm.dot(u)) < 1e-14 + + # modified data cannot be re-created + M = mass_ops.create_weighted_mass("H1", "H1", weights=("sqrt_g",), assemble=True) + assert M.is_reconstructible + M *= 2.0 + assert not M.is_reconstructible + with pytest.raises(ValueError): + M.to_dict() + + +def test_basis_projection_to_dict(): + from struphy.feec.basis_projection_ops import BasisProjectionOperator + + derham, domain = _derham((8, 6, 4), (2, 2, 1), BCS[0]) + fun = [[lambda e1, e2, e3: 1.0 + e1 * e3]] + K = BasisProjectionOperator( + derham.projectors["L2"], + derham.fem_spaces["H1"], + fun, + V_extraction_op=derham.extraction_ops["H1"], + V_boundary_op=derham.boundary_ops["H1"], + ) + for op, V in [(K, "0"), (K.T, "3")]: + _, u = create_equal_random_arrays(derham.fem_spaces[V], seed=3) + op2 = BasisProjectionOperator.from_dict(op.to_dict(), derham) + assert _max_diff(op2.dot(u), op.dot(u)) < 1e-14 + + +@pytest.mark.parametrize("bcs", BCS) +@pytest.mark.parametrize("degree", [(2, 3, 1), (3, 1, 2)]) +def test_coarsen_poisson(bcs, degree): + r"""Re-discretization of :math:`\sigma M_0 + G^\top M_1 G` equals the Galerkin product :math:`R A P` (Cuboid, exact quadrature).""" + derham, domain = _derham((8, 8, 4), degree, bcs) + h = MultiGridHierarchy(derham, max_levels=2) + mass_ops = WeightedMassOperators(derham, domain) + A = 0.7 * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + + C = OperatorCoarsener(h[0], h[1], domain) + Ac = C(A) + assert Ac.domain is h[1].coeff_spaces["0"] and Ac.codomain is h[1].coeff_spaces["0"] + + # same as building it directly on the coarse level + mass_c = WeightedMassOperators(h[1], domain) + Ad = 0.7 * mass_c.M0 + h[1].grad.T @ mass_c.M1 @ h[1].grad + _, u = create_equal_random_arrays(h[1].fem_spaces["0"], seed=5) + assert _max_diff(Ac.dot(u), Ad.dot(u)) < 1e-13 + + # Galerkin property + P = SplineProlongation(h[1], h[0], "H1") + assert _max_diff(P.T.dot(A.dot(P.dot(u))), Ac.dot(u)) < 1e-12 + + # leaves are cached: a new scalar re-assembles nothing + A2 = 2.0 * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + Ac2 = C(A2) + leaves = lambda op: [a for a in op.addends] + assert leaves(Ac2)[0].operator is leaves(Ac)[0].operator + + +@pytest.mark.parametrize("bcs", BCS) +def test_coarsen_curl_curl(bcs): + r"""Re-discretization of :math:`C^\top M_2 C + M_1` equals the Galerkin product on 1-forms.""" + derham, domain = _derham((8, 8, 4), (2, 2, 1), bcs) + h = MultiGridHierarchy(derham, max_levels=2) + mass_ops = WeightedMassOperators(derham, domain) + A = derham.curl.T @ mass_ops.M2 @ derham.curl + mass_ops.M1 + + Ac = OperatorCoarsener(h[0], h[1], domain)(A) + P = SplineProlongation(h[1], h[0], "Hcurl") + _, u = create_equal_random_arrays(h[1].fem_spaces["1"], seed=6) + assert _max_diff(P.T.dot(A.dot(P.dot(u))), Ac.dot(u)) < 1e-12 + + +def test_coarsen_unknown_leaf(): + from feectools.linalg.stencil import StencilMatrix + + derham, domain = _derham((8, 8, 4), (2, 2, 1), BCS[0]) + h = MultiGridHierarchy(derham, max_levels=2) + S = StencilMatrix(derham.coeff_spaces["0"], derham.coeff_spaces["0"]) + with pytest.raises(NotImplementedError): + OperatorCoarsener(h[0], h[1], domain)(S) diff --git a/src/struphy/linear_algebra/tests/test_multigrid_solver.py b/src/struphy/linear_algebra/tests/test_multigrid_solver.py new file mode 100644 index 000000000..3cd809231 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_multigrid_solver.py @@ -0,0 +1,202 @@ +import numpy as np +import pytest +from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.basic import LinearOperator + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.feec.utilities import create_equal_random_arrays +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.multigrid.preconditioner import ( + MultiGridOptions, + MultiGridPreconditioner, + MultiGridSolver, + _assemble_dense, +) +from struphy.linear_algebra.multigrid.smoothers import ( + ChebyshevSmoother, + DiagonalComputer, + JacobiSmoother, + KrylovSmoother, + inverse_diagonal, +) +from struphy.topology.grids import TensorProductGrid + +DIRICHLET = (("dirichlet", "dirichlet"), ("dirichlet", "dirichlet"), None) +PERIODIC = (None, None, None) + + +def _poisson(n, p, bcs, sigma=0.0): + domain = Cuboid(l1=0.0, r1=1.0, l2=0.0, r2=2.0, l3=0.0, r3=1.0) + derham = Derham( + TensorProductGrid(num_elements=(n, n, 1)), + DerhamOptions(degree=(p, p, 1), bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domain, + ) + mass_ops = WeightedMassOperators(derham, domain) + A = derham.grad.T @ mass_ops.M1 @ derham.grad + if sigma != 0.0: + A = sigma * mass_ops.M0 + A + return derham, domain, mass_ops, A + + +class _SmootherAsOperator(LinearOperator): + """x = S b (one smoother call from zero initial guess).""" + + def __init__(self, S): + self._S = S + + domain = property(lambda self: self._S.A.domain) + codomain = property(lambda self: self._S.A.domain) + dtype = property(lambda self: float) + + def transpose(self, conjugate=False): + return self + + def dot(self, b, out=None): + x = self.domain.zeros() + self._S.smooth(b, x) + if out is None: + return x + x.copy(out=out) + return out + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("bcs", [DIRICHLET, PERIODIC]) +def test_diagonal(bcs): + """Probed diagonal equals the diagonal of the assembled operator.""" + derham, _, mass_ops, A = _poisson(8, 3, bcs, sigma=0.3) + d = DiagonalComputer(derham.degree)(A) + assert np.allclose(d.toarray(), np.diag(_assemble_dense(A)), atol=1e-14) + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("kind", ["chebyshev_jacobi", "chebyshev_mass", "jacobi"]) +def test_smoother_symmetric(kind): + """The linear smoothers are symmetric (as matrices from right-hand side to iterate).""" + from struphy.feec.preconditioner import MassMatrixPreconditioner + + derham, _, mass_ops, A = _poisson(8, 2, PERIODIC, sigma=0.5) + D_inv = inverse_diagonal(DiagonalComputer(derham.degree)(A)) + if kind == "chebyshev_jacobi": + S = ChebyshevSmoother(A, D_inv, degree=3) + elif kind == "chebyshev_mass": + S = ChebyshevSmoother(A, MassMatrixPreconditioner(mass_ops.M0), degree=3) + else: + S = JacobiSmoother(A, D_inv, sweeps=3) + Sd = _assemble_dense(_SmootherAsOperator(S)) + assert np.abs(Sd - Sd.T).max() < 1e-12 * np.abs(Sd).max() + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("iterations", [1, 2, 4]) +def test_krylov_smoother(iterations): + """KrylovSmoother performs exactly ``iterations`` CG steps and is safe for a zero residual.""" + derham, _, mass_ops, A = _poisson(8, 2, PERIODIC, sigma=0.5) + _, b = create_equal_random_arrays(derham.fem_spaces["0"], seed=3) + + # zero right-hand side with zero initial guess: x stays zero (no NaN) + x = A.domain.zeros() + KrylovSmoother(A, iterations=iterations).smooth(A.domain.zeros(), x) + assert np.all(x.toarray() == 0.0) + + # the k-th CG iterate minimizes the A-norm error over the k-th Krylov space, so the error + # decreases strictly with the number of iterations; compare with one call of k - 1 iterations + x_ref = A.domain.zeros() + if iterations > 1: + KrylovSmoother(A, iterations=iterations - 1).smooth(b, x_ref) + x = A.domain.zeros() + KrylovSmoother(A, iterations=iterations).smooth(b, x) + + Ad = _assemble_dense(A) + x_ex = np.linalg.solve(Ad, b.toarray()) + + def err(y): + e = y.toarray() - x_ex + return e @ Ad @ e + + assert err(x) < err(x_ref) + + # one step from zero equals the steepest descent step alpha * b + if iterations == 1: + alpha = b.inner(b) / b.inner(A.dot(b)) + assert np.allclose(x.toarray(), alpha * b.toarray(), rtol=1e-12, atol=1e-14) + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("bcs, nullspace", [(DIRICHLET, None), (PERIODIC, "constants")]) +@pytest.mark.parametrize("smoother_precond", ["mass", "jacobi"]) +def test_vcycle_spd(bcs, nullspace, smoother_precond): + """The V-cycle is symmetric positive definite (on the complement of the null space) and contracts.""" + derham, domain, mass_ops, A = _poisson(16, 2, bcs) + pc = MultiGridPreconditioner( + A, derham, domain, MultiGridOptions(smoother_precond=smoother_precond, nullspace=nullspace), mass_ops=mass_ops + ) + B = _assemble_dense(pc) + Ad = _assemble_dense(A) + assert np.abs(B - B.T).max() < 1e-12 * np.abs(B).max() + + # restrict to the dofs/modes that matter: interior dofs (Dirichlet) or zero-mean vectors (periodic) + N = Ad.shape[0] + if nullspace == "constants": + Q = np.linalg.qr(np.eye(N) - np.ones((N, N)) / N)[0][:, : N - 1] + else: + Q = np.eye(N)[:, np.flatnonzero(np.diag(Ad) != 0.0)] + assert np.linalg.eigvalsh(Q.T @ B @ Q).min() > 0.0 + E = Q.T @ (np.eye(N) - B @ Ad) @ Q + assert np.abs(np.linalg.eigvals(E)).max() < 0.5 + + +@pytest.mark.parametrize("bcs, nullspace", [(DIRICHLET, None), (PERIODIC, "constants")]) +@pytest.mark.parametrize("p", [2, 3]) +@pytest.mark.parametrize( + "smoother, smoother_precond", + [("chebyshev", "mass"), ("chebyshev", "jacobi"), ("jacobi", "identity"), ("cg", "mass")], +) +def test_poisson_h_independent(bcs, nullspace, p, smoother, smoother_precond): + """MG-preconditioned CG converges in a small, mesh-independent number of iterations.""" + niter = [] + for n in (16, 32): + derham, domain, mass_ops, A = _poisson(n, p, bcs) + _, u = create_equal_random_arrays(derham.fem_spaces["0"], seed=3) + b = A.dot(u) + opts = MultiGridOptions(smoother=smoother, smoother_precond=smoother_precond, nullspace=nullspace) + solver = MultiGridSolver(A, derham, domain, opts, mass_ops=mass_ops, tol=1e-8) + x = solver.dot(b) + r = b - A.dot(x) + assert solver.info["success"] + assert np.sqrt(r.inner(r)) <= 1e-8 * np.sqrt(b.inner(b)) + niter.append(solver.info["niter"]) + assert max(niter) <= 15 + assert niter[1] <= niter[0] + 2 + + +def test_update(): + """Changing a scalar of the operator re-uses the coarse operators and still converges.""" + derham, domain, mass_ops, A = _poisson(16, 2, DIRICHLET, sigma=2.0) + solver = MultiGridSolver(A, derham, domain, mass_ops=mass_ops, tol=1e-10) + coarse_M0 = solver.preconditioner.operators[1].addends[0].operator + + A2 = 100.0 * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + solver.update(A2) + assert solver.preconditioner.operators[1].addends[0].operator is coarse_M0 + + _, u = create_equal_random_arrays(derham.fem_spaces["0"], seed=4) + b = A2.dot(u) + x = solver.dot(b) + r = b - A2.dot(x) + assert np.sqrt(r.inner(r)) <= 1e-10 * np.sqrt(b.inner(b)) + assert solver.info["niter"] <= 15 + + +@pytest.mark.mpi_skip +def test_options(): + opts = MultiGridOptions(smoother="jacobi", max_levels=3) + assert MultiGridOptions.from_dict(opts.to_dict()) == opts + with pytest.raises(AssertionError): + MultiGridOptions(smoother="gauss-seidel") + with pytest.raises(AssertionError): + MultiGridOptions(n_pre=0, n_post=0) diff --git a/src/struphy/linear_algebra/tests/test_multigrid_transfer.py b/src/struphy/linear_algebra/tests/test_multigrid_transfer.py new file mode 100644 index 000000000..2584f75a4 --- /dev/null +++ b/src/struphy/linear_algebra/tests/test_multigrid_transfer.py @@ -0,0 +1,113 @@ +import numpy as np +import pytest +from feectools.ddm.mpi import mpi as MPI + +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.psydac_derham import Derham +from struphy.feec.utilities import create_equal_random_arrays +from struphy.geometry.domains import Cuboid +from struphy.io.options import DerhamOptions +from struphy.linear_algebra.multigrid.hierarchy import MultiGridHierarchy +from struphy.linear_algebra.multigrid.transfer import SplineProlongation, prolongation_matrix_1d +from struphy.topology.grids import TensorProductGrid + +BCS = [ + (None, None, None), + (("dirichlet", "dirichlet"), None, ("free", "dirichlet")), +] + + +def _hierarchy(num_elements, degree, bcs, max_levels=None): + domain = Cuboid(l1=0.0, r1=2.0, l2=0.0, r2=1.0, l3=0.0, r3=3.0) + derham = Derham( + TensorProductGrid(num_elements=num_elements), + DerhamOptions(degree=degree, bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domain, + ) + return MultiGridHierarchy(derham, max_levels=max_levels), domain + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("degree", [1, 2, 3, 4]) +@pytest.mark.parametrize("periodic", [True, False]) +@pytest.mark.parametrize("basis", ["B", "M"]) +def test_prolongation_matrix_1d(degree, periodic, basis): + """The coarse basis is reproduced exactly by the prolongated coefficients; partition of unity is kept.""" + from feectools.fem.splines import SplineSpace + + def space(n): + return SplineSpace(degree, grid=np.linspace(0.0, 1.0, n + 1), periodic=periodic, basis=basis) + + coarse, fine = space(8), space(16) + P = prolongation_matrix_1d(coarse, fine) + assert P.shape == (fine.nbasis, coarse.nbasis) + + if basis == "B": + # partition of unity: the constant function has coefficients 1 on both grids + assert np.allclose(P @ np.ones(coarse.nbasis), 1.0) + + # each fine row couples to at most ceil((p+2)/2) coarse functions + assert np.max(np.count_nonzero(P, axis=1)) <= (degree + 3) // 2 + + +@pytest.mark.parametrize("num_elements, degree", [((16, 8, 8), (3, 2, 1)), ((8, 16, 1), (2, 3, 1))]) +@pytest.mark.parametrize("bcs", BCS) +def test_hierarchy(num_elements, degree, bcs): + """Coarse levels have aligned decompositions and at least degree+1 cells per coarsened direction.""" + h, _ = _hierarchy(num_elements, degree, bcs) + assert h.n_levels >= 2 + assert len(h.factors) == h.n_levels - 1 + for l, f in enumerate(h.factors): + fine, coarse = h[l], h[l + 1] + for axis in range(3): + assert fine.num_elements[axis] == f[axis] * coarse.num_elements[axis] + assert fine.domain_decomposition.starts[axis] == f[axis] * coarse.domain_decomposition.starts[axis] + if f[axis] == 2: + assert coarse.num_elements[axis] >= degree[axis] + 1 + assert coarse.options is fine.options + # the coarsest level cannot be coarsened further + assert h._coarsening_factors(h[-1]) == (1, 1, 1) + + +@pytest.mark.parametrize("bcs", BCS) +@pytest.mark.parametrize("space_id", ["H1", "Hcurl", "Hdiv", "L2", "H1vec"]) +def test_transfer(bcs, space_id): + r"""Restriction is the transpose of the prolongation, and P is the exact embedding: R M_h P = M_H.""" + h, domain = _hierarchy((16, 8, 8), (3, 2, 1), bcs, max_levels=3) + comm = MPI.COMM_WORLD + form = h[0].space_to_form[space_id] + + for l in range(h.n_levels - 1): + P = SplineProlongation(h[l + 1], h[l], space_id) + R = P.T + assert R.domain is P.codomain and R.codomain is P.domain + + _, u = create_equal_random_arrays(h[l + 1].fem_spaces[form], seed=1) + _, w = create_equal_random_arrays(h[l].fem_spaces[form], seed=2) + assert np.isclose(R.dot(w).inner(u), w.inner(P.dot(u)), rtol=1e-12) + + Mh = getattr(WeightedMassOperators(h[l], domain), "M" + form) + MH = getattr(WeightedMassOperators(h[l + 1], domain), "M" + form) + a = R.dot(Mh.dot(P.dot(u))) + b = MH.dot(u) + err = np.max(np.abs((a - b).toarray())) + ref = np.max(np.abs(b.toarray())) + if comm.Get_size() > 1: + err = comm.allreduce(err, op=MPI.MAX) + ref = comm.allreduce(ref, op=MPI.MAX) + assert err < 1e-12 * ref + + +@pytest.mark.parametrize("bcs", BCS) +def test_prolongation_of_spline(bcs): + """The prolongated coefficients represent the same function (point evaluation).""" + h, _ = _hierarchy((8, 8, 4), (2, 3, 1), bcs, max_levels=2) + P = SplineProlongation(h[1], h[0], "H1") + _, u = create_equal_random_arrays(h[1].fem_spaces["0"], seed=4) + u = h[1].boundary_ops["0"].dot(u) + + fH = h[1].create_spline_function("fH", "H1", coeffs=u) + fh = h[0].create_spline_function("fh", "H1", coeffs=P.dot(u)) + e = np.linspace(0.0, 1.0, 7) + assert np.allclose(fH(e, e, e), fh(e, e, e), atol=1e-12) diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index e211e127b..2af8fcd03 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -11,6 +11,7 @@ from struphy.feec.mass import L2Projector, 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 @@ -182,10 +183,16 @@ class Options(OptionsBase): Name of the symmetric iterative solver passed to :func:`psydac.linalg.solvers.inverse`. - precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner" + precond : LiteralOptions.OptsDiffusionPrecond, default="MassMatrixPreconditioner" Name of the preconditioner configuration. - Currently this class sets ``pc=None`` internally, so this option is - reserved for compatibility and future extensions. + ``"MultiGrid"`` uses a geometric multigrid V-cycle + (:class:`~struphy.linear_algebra.multigrid.preconditioner.MultiGridPreconditioner`, + requires ``solver="pcg"``). The other values currently result in ``pc=None``. + + 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``. solver_params : SolverParameters, default=None Iterative-solver controls (for example ``tol``, ``maxiter``, @@ -216,7 +223,8 @@ class Options(OptionsBase): diffusion_mat: OptsDiffusionMat = "M1" x0: StencilVector = None solver: LiteralOptions.OptsSymmSolver = "pcg" - precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" + precond: LiteralOptions.OptsDiffusionPrecond = "MassMatrixPreconditioner" + multigrid: MultiGridOptions = None solver_params: SolverParameters = None filter_params: dict[PICVariable | SPHVariable, FilterParameters] = None @@ -225,11 +233,15 @@ def __post_init__(self): 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.OptsMassPrecond) + check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) + if self.precond == "MultiGrid": + assert self.solver == "pcg", "precond='MultiGrid' requires solver='pcg'." # defaults if self.solver_params is None: self.solver_params = SolverParameters() + if self.multigrid is None: + self.multigrid = MultiGridOptions() @property def options(self) -> Options: @@ -345,10 +357,20 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: self._diffusion_op = self.derham.grad.T @ diffusion_mat @ self.derham.grad # preconditioner and solver for Ax=b - if self.options.precond is None: - pc = None + self._mg = None + if self.options.precond == "MultiGrid": + # the operator is updated in __call__ if sigma_1 changes (e.g. with dt) + self._mg_sig_1 = self._sigma_1 + self._mg = MultiGridPreconditioner( + self._sigma_1 * stab_mat + self._diffusion_op, + self.derham, + self.domain, + self.options.multigrid, + mass_ops=self.mass_ops, + ) + pc = self._mg else: - # TODO: waiting for multigrid preconditioner + # TODO: mass-matrix preconditioners are not effective for this operator pc = None # solver just with A_2, but will be set during call with dt @@ -459,6 +481,9 @@ def __call__(self, dt): # compute lhs self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + if self._mg is not None and sig_1 != self._mg_sig_1: + self._mg.update(self._solver.linop) + self._mg_sig_1 = sig_1 # solve with ProfileManager.profile_region(self._solve_region, functions=[self._solver.solve]): diff --git a/src/struphy/propagators/poisson_solve.py b/src/struphy/propagators/poisson_solve.py index cb7a1c3be..9e71cb73c 100644 --- a/src/struphy/propagators/poisson_solve.py +++ b/src/struphy/propagators/poisson_solve.py @@ -5,6 +5,7 @@ from feectools.linalg.stencil import StencilVector from struphy.io.options import LiteralOptions, OptionsBase +from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable, PICVariable, SPHVariable from struphy.pic.accumulation.filter import FilterParameters @@ -66,10 +67,14 @@ class Options(OptionsBase): Name of the symmetric iterative solver passed to :func:`psydac.linalg.solvers.inverse`. - precond : LiteralOptions.OptsMassPrecond, default="MassMatrixPreconditioner" - Name of the preconditioner configuration. - Currently this class inherits the same behavior as - :class:`ImplicitDiffusion`, where ``pc=None`` is used internally. + precond : LiteralOptions.OptsDiffusionPrecond, default="MassMatrixPreconditioner" + Name of the preconditioner configuration, see :class:`ImplicitDiffusion` + (``"MultiGrid"`` for geometric multigrid). + + multigrid : MultiGridOptions, default=None + Options of the multigrid preconditioner (if ``precond="MultiGrid"``). + For periodic or Neumann boundary conditions with ``stab_eps = 0``, use + ``MultiGridOptions(nullspace="constants")``. solver_params : SolverParameters, default=None Iterative-solver controls (for example ``tol``, ``maxiter``, @@ -96,7 +101,8 @@ class Options(OptionsBase): diffusion_mat: OptsDiffusionMat = "M1" x0: StencilVector = None solver: LiteralOptions.OptsSymmSolver = "pcg" - precond: LiteralOptions.OptsMassPrecond = "MassMatrixPreconditioner" + precond: LiteralOptions.OptsDiffusionPrecond = "MassMatrixPreconditioner" + multigrid: MultiGridOptions = None solver_params: SolverParameters = None filter_params: dict[PICVariable | SPHVariable, FilterParameters] = None @@ -105,11 +111,15 @@ def __post_init__(self): 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.OptsMassPrecond) + check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) + if self.precond == "MultiGrid": + assert self.solver == "pcg", "precond='MultiGrid' requires solver='pcg'." # defaults if self.solver_params is None: self.solver_params = SolverParameters() + if self.multigrid is None: + self.multigrid = MultiGridOptions() # Poisson solve (-> set some params of parent class) self.sigma_1 = self.stab_eps diff --git a/src/struphy/propagators/tests/test_poisson.py b/src/struphy/propagators/tests/test_poisson.py index dafd9bb48..d1c70e48e 100644 --- a/src/struphy/propagators/tests/test_poisson.py +++ b/src/struphy/propagators/tests/test_poisson.py @@ -786,6 +786,101 @@ def rho2_pulled(e1, e2, e3): assert error2 < err_lim +@pytest.mark.parametrize("degree", [[2, 2, 1], [3, 3, 1]]) +@pytest.mark.parametrize("bc_type", ["periodic", "dirichlet", "neumann"]) +def test_poisson_2d_multigrid(degree, bc_type): + """PoissonSolve with precond="MultiGrid" agrees with the unpreconditioned solve, in few iterations.""" + from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions + + domain = domains.Colella(Lx=4.0, Ly=2.0, alpha=0.1, Lz=1.0) + bcs = { + "periodic": (None, None, None), + "dirichlet": (("dirichlet", "dirichlet"), None, None), + "neumann": (("free", "free"), None, None), + }[bc_type] + derham = Derham(TensorProductGrid(num_elements=[32, 32, 1]), DerhamOptions(degree=degree, bcs=bcs), comm=comm) + mass_ops = WeightedMassOperators(derham, domain) + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + def rho(e1, e2, e3): + return xp.cos(2 * xp.pi * e1) * xp.sin(2 * xp.pi * e2) + 0.3 * xp.sin(4 * xp.pi * e2) + + phis = [] + infos = [] + for precond in ["MassMatrixPreconditioner", "MultiGrid"]: + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + solver = PoissonSolve(rho=rho) + solver.variables.phi = phi + solver.options = solver.Options( + stab_eps=1e-8, + solver="pcg", + precond=precond, + # the Jacobi smoother is robust w.r.t. the mapping (the default mass smoother is robust w.r.t. the degree); + # no null space: the system is regularized by stab_eps + multigrid=MultiGridOptions(smoother_precond="jacobi"), + solver_params=SolverParameters(tol=1e-11, maxiter=3000, recycle=False), + ) + solver.allocate() + solver(1.0) + phis.append(phi.spline.vector.toarray()) + infos.append(solver._solver._info) + + # global coefficient arrays (toarray only fills the local part) + if comm.Get_size() > 1: + phis = [comm.allreduce(p, op=MPI.SUM) for p in phis] + if bc_type != "dirichlet": + # solutions are defined up to a constant (the stabilization is tiny) + phis = [p - xp.mean(p) for p in phis] + assert xp.max(xp.abs(phis[0] - phis[1])) < 1e-6 * xp.max(xp.abs(phis[0])) + assert infos[1]["niter"] <= 25 + assert infos[1]["niter"] < infos[0]["niter"] + + +def test_implicit_diffusion_multigrid_dt(): + """With divide_by_dt, the multigrid preconditioner follows changes of dt.""" + from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions + from struphy.propagators.implicit_diffusion import ImplicitDiffusion + + domain = domains.Cuboid(l1=0.0, r1=2.0, l2=0.0, r2=1.0, l3=0.0, r3=1.0) + derham = Derham( + TensorProductGrid(num_elements=[32, 16, 1]), + DerhamOptions(degree=[2, 2, 1], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=comm, + ) + mass_ops = WeightedMassOperators(derham, domain) + Propagator.derham = derham + Propagator.domain = domain + Propagator.mass_ops = mass_ops + + phi = FEECVariable(space="H1") + phi.allocate(derham=derham, domain=domain) + phi.spline.vector = derham.P0(lambda e1, e2, e3: xp.sin(xp.pi * e1) * xp.cos(2 * xp.pi * e2)) + + prop = ImplicitDiffusion() + prop.variables.phi = phi + prop.options = prop.Options( + sigma_1=1.0, + sigma_2=1.0, + sigma_3=0.0, + divide_by_dt=True, + precond="MultiGrid", + multigrid=MultiGridOptions(), + solver_params=SolverParameters(tol=1e-12, maxiter=100, recycle=False), + ) + prop.allocate() + + for dt in [0.1, 0.1, 0.01]: + rhs = (1.0 / dt) * mass_ops.M0.dot(phi.spline.vector) + prop(dt) + A = (1.0 / dt) * mass_ops.M0 + derham.grad.T @ mass_ops.M1 @ derham.grad + r = rhs - A.dot(phi.spline.vector) + assert xp.sqrt(r.inner(r)) < 1e-9 * xp.sqrt(rhs.inner(rhs)) + assert prop._solver._info["niter"] <= 15 + + if __name__ == "__main__": # direction = 0 # bc_type = "dirichlet"