diff --git a/feectools b/feectools index a99769785..5385af968 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit a997697858cb6be57f7066d30e077f0b7c755a5b +Subproject commit 5385af968237c96f7d3ca643a5da4f4e4e31916d diff --git a/pyproject.toml b/pyproject.toml index d1f709064..d14c35a11 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "cunumpy >= 0.6.1, <= 0.6.1", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", - "feectools>=0.6.0, <=0.6.0", + "feectools>=0.7.0, <=0.7.0", "scipy<=1.18.0", "h5py<=3.16.0", "h5netcdf<=1.8.1", diff --git a/src/struphy/feec/memory.py b/src/struphy/feec/memory.py index bbfc17c3e..5b43bf3d5 100644 --- a/src/struphy/feec/memory.py +++ b/src/struphy/feec/memory.py @@ -58,9 +58,10 @@ def linop_nbytes(op, _seen: set = None) -> int: if isinstance(op, BlockLinearOperator): return sum(linop_nbytes(block, _seen) for row in op.blocks for block in row) - for attr in ("multiplicants", "addends", "mats"): - if hasattr(op, attr): - return sum(linop_nbytes(child, _seen) for child in getattr(op, attr)) + # (a ComposedKroneckerStencilMatrix has both multiplicands and mats, the factors of the product) + children = [attr for attr in ("multiplicands", "addends", "mats") if hasattr(op, attr)] + if children: + return sum(linop_nbytes(child, _seen) for attr in children for child in getattr(op, attr)) for attr in ("operator", "tp_operator"): if hasattr(op, attr): diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index ac413ba59..f671c39c2 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -1,49 +1,81 @@ +from __future__ import annotations + +import itertools import logging +from collections.abc import Callable +from functools import reduce import cunumpy as xp +import numpy as np from feectools.api.essential_bc import apply_essential_bc_stencil from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.feec.derivatives import DirectionalDerivativeOperator from feectools.fem.tensor import TensorFemSpace -from feectools.linalg.basic import ComposedLinearOperator, LinearOperator, Vector -from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.basic import ( + ComposedLinearOperator, + LinearOperator, + ScaledLinearOperator, + SumLinearOperator, + Vector, + VectorSpace, +) +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace from feectools.linalg.direct_solvers import BandedSolver, SparseSolver -from feectools.linalg.kron import KroneckerLinearSolver, KroneckerStencilMatrix -from feectools.linalg.stencil import StencilMatrix, StencilVectorSpace +from feectools.linalg.kron import KroneckerLinearSolver, KroneckerStencilMatrix, KroneckerSumSolver +from feectools.linalg.stencil import StencilDiagonalMatrix, StencilMatrix, StencilVectorSpace from line_profiler import profile from maybempi import MPI, SerialComm from scipy import sparse from struphy.feec.linear_operators import BoundaryOperator -from struphy.feec.mass import WeightedMassOperator +from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators logger = logging.getLogger("struphy") -class MassMatrixPreconditioner(LinearOperator): - """ - Preconditioner for inverting 3d weighted mass matrices. +class KroneckerPreconditioner(LinearOperator): + r""" + Base class for preconditioners built from Kronecker approximations. - The mass matrix is approximated by a Kronecker product of 1d mass matrices - in each direction with correct boundary conditions (block diagonal in case of vector-valued spaces). - In this process, the 3d weight function is appoximated by a 1d counterpart in the dim_reduce direction - (default 1st direction) at the fixed point (0.5) in the other directions. The inversion is then - performed with a Kronecker solver. + The operator to be preconditioned has the form :math:`A = B E \, A_c \, E^T B^T`, where + :math:`A_c` acts on the tensor-product coefficient space and :math:`B E` are the + boundary and extraction operators of a mass operator (the composition ``M0`` or ``M`` + of ``mass_operator``). The preconditioner applies + + .. math:: + + P = B E \, S \tilde A_c^{-1} S \, E^T B^T \,, + + where :math:`\tilde A_c` (``matrix``) approximates :math:`A_c` and can be inverted exactly + (``solver``), e.g. as a Kronecker product. With ``diagonal_scaling``, the diagonal operator + :math:`S = \hat D^{1/2} D^{-1/2}` with :math:`D = \mathrm{diag}(A_c)` and + :math:`\hat D = \mathrm{diag}(\tilde A_c)` corrects the approximation pointwise + (Loli, Sangalli, Tani, Comp. Math. Appl. 116, 2022); otherwise :math:`S = I`. + + Subclasses implement ``_build_approximation`` and ``transpose``, and ``_core_diagonal`` + for diagonal scaling. Parameters ---------- - mass_operator : struphy.feec.mass.WeightedMassOperator - The weighted mass operator for which the approximate inverse is needed. + mass_operator : WeightedMassOperator + Mass operator whose composition defines the spaces and :math:`B E`; its core + (``mass_operator.matrix``) is replaced by :math:`S \tilde A_c^{-1} S`. apply_bc : bool Whether to include boundary operators. - dim_reduce : int - Along which axis to take the approximate value of the weight + diagonal_scaling : bool + Whether to correct the approximation with the diagonal scaling :math:`S`. """ - def __init__(self, mass_operator, apply_bc=True, dim_reduce=0): + def __init__( + self, + mass_operator: WeightedMassOperator, + apply_bc: bool = True, + diagonal_scaling: bool = False, + ): assert isinstance(mass_operator, WeightedMassOperator) - assert mass_operator.domain == mass_operator.codomain, "Only square mass matrices can be inverted!" + assert mass_operator.domain == mass_operator.codomain, "Only square operators can be inverted!" self._mass_operator = mass_operator self._femspace = mass_operator.domain_femspace @@ -52,329 +84,78 @@ def __init__(self, mass_operator, apply_bc=True, dim_reduce=0): self._codomain = mass_operator.codomain self._domain = mass_operator.domain self._apply_bc = apply_bc - self._dim_reduce = dim_reduce + self._diagonal_scaling = diagonal_scaling - # 3d Kronecker stencil matrices and solvers - solverblocks = [] - matrixblocks = [] + assert self._femspace.ldim == 3 # other dims not yet implemented - # collect TensorFemSpaces in a tuple - if isinstance(self._femspace, TensorFemSpace): - femspaces = (self._femspace,) - else: - femspaces = self._femspace.spaces + # boundary conditions are only imposed if the mass operator has a BoundaryOperator + self._bc = _boundary_conditions(mass_operator, apply_bc) - n_comps = len(femspaces) - n_dims = self._femspace.ldim + # approximation of the core operator and its exact inverse + self._matrix, self._solver = self._build_approximation(self._bc) - assert n_dims == 3 # other dims not yet implemented - assert dim_reduce < n_dims + # diagonal scaling S (a sequence of diagonal operators, applied in this order before the + # solver and in reverse order after it) + self._scaling = self._build_scaling() if diagonal_scaling else () + self._tmp_core = tuple(self._matrix.codomain.zeros() for _ in range(2)) if self._scaling else () - # get boundary conditions list from BoundaryOperator in ComposedLinearOperator M0 of mass operator - if apply_bc and isinstance(mass_operator.M0, ComposedLinearOperator): - if isinstance(mass_operator.M0.multiplicants[-1], BoundaryOperator): - bc = mass_operator.M0.multiplicants[-1].bc - else: - apply_bc = False - bc = None - else: - apply_bc = False - bc = None + # operator to be inverted (with boundary operators B, E if apply_bc), needed in solve + self._M, self._mass_index, self._tmp_vectors = _operator_to_invert(mass_operator, self._bc is not None) - # define subcomm to gather 1d weight info along dim_reduce - derham = mass_operator.derham - logger.debug(f"{derham.num_elements = }, {derham.bcs = }, {derham.degree = }") - comm = derham.comm - dom_dec = derham.domain_decomposition - gather_weights = not isinstance(comm, (SerialComm, type(None))) - if gather_weights: - rank = comm.Get_rank() - # in the directions other than dim_reduce, the weight is taken at the global mid point; - # select the ranks owning the mid element there (exactly one rank per slab along dim_reduce) - is_selected = all( - dom_dec.starts[i] <= derham.num_elements[i] // 2 <= dom_dec.ends[i] - for i in range(n_dims) - if i != dim_reduce - ) - color = 0 if is_selected else MPI.UNDEFINED - subcomm = comm.Split(color=color, key=rank) - root = comm.allreduce(rank if is_selected else comm.Get_size(), op=MPI.MIN) - logger.debug(f"Rank {rank} selected for gathering 1d weight info in dimension {dim_reduce}: {is_selected}") - - # loop over components - for c in range(n_comps): - # 1d mass matrices and solvers - solvercells = [] - matrixcells = [] - - # loop over spatial directions - for d in range(n_dims): - # weight function only along in first direction - if d == dim_reduce: - # pts = [0.5] * (n_dims - 1) - loc_weights = mass_operator.weights[c][c] - if callable(loc_weights): - - def fun(e): - # make input in meshgrid format to be able to use it with general functions - s = e.shape[0] - newshape = tuple([1 if i != d else s for i in range(n_dims)]) - f = e.reshape(newshape) - return xp.atleast_1d( - loc_weights( - *[xp.array(xp.full_like(f, 0.5)) if i != d else xp.array(f) for i in range(n_dims)], - ).squeeze(), - ) - elif isinstance(loc_weights, xp.ndarray): - s = loc_weights.shape - logger.debug(f"{loc_weights.shape = } for component {c} and direction {d}.") - npts = derham.num_elements[d] * derham.nquads[d] - fun = xp.zeros(npts, dtype=float) - # local index of the global mid quadrature point in the other directions - # (clipped on non-selected ranks, which receive the gathered weight below) - mid = [0] * n_dims - for i in range(n_dims): - if i != d: - nq_i = s[i] // dom_dec.local_ncells[i] - mid_i = (derham.num_elements[i] * nq_i) // 2 - dom_dec.starts[i] * nq_i - mid[i] = min(max(mid_i, 0), s[i] - 1) - if d == 0: - local_fun = loc_weights[:, mid[1], mid[2]] - elif d == 1: - local_fun = loc_weights[mid[0], :, mid[2]] - elif d == 2: - local_fun = loc_weights[mid[0], mid[1], :] - local_fun = xp.ascontiguousarray(local_fun, dtype=float) - logger.debug( - f"{fun.size = } for component {c} and direction {d} before gathering on all processes." - ) - if gather_weights: - # local sizes differ if num_elements[d] is not divisible by the number of processes - if subcomm != MPI.COMM_NULL: - counts = subcomm.allgather(local_fun.size) - displs = [sum(counts[:j]) for j in range(len(counts))] - subcomm.Allgatherv(local_fun, [fun, counts, displs, MPI.DOUBLE]) - comm.Bcast(fun, root=root) - else: - fun[:] = local_fun - logger.debug( - f"{fun.shape = } for component {c} and direction {d} after gathering on all processes." - ) - elif loc_weights is None: - fun = lambda e: xp.ones(e.size, dtype=float) - else: - raise TypeError( - "weights needs to be callable, xp.ndarray or None but is{}".format(type(loc_weights)), - ) - fun = [[fun]] - else: - fun = [[lambda e: xp.ones(e.size, dtype=float)]] - - # get 1D FEM space (serial, not distributed) and quadrature order - if femspaces[c].spaces[d].basis == "B": - femspace_1d_tensor = mass_operator.derham.H1_1d_serial[d] - else: - femspace_1d_tensor = mass_operator.derham.L2_1d_serial[d] - - domain_decompos_1d = femspace_1d_tensor.domain_decomposition - qu_order_1d = (mass_operator.derham.nquads[d],) - - M = WeightedMassOperator( - mass_operator.derham, - femspace_1d_tensor, - femspace_1d_tensor, - weights_info=fun, - nquads=qu_order_1d, - ) - M.assemble() - M = M.matrix - - # apply boundary conditions - if apply_bc: - if mass_operator._domain_symbolic_name not in ("H1H1H1", "H1vec"): - if femspaces[c].spaces[d].basis == "B": - if bc[d][0]: - apply_essential_bc_stencil( - M, - axis=0, - ext=-1, - order=0, - identity=True, - ) - if bc[d][1]: - apply_essential_bc_stencil( - M, - axis=0, - ext=+1, - order=0, - identity=True, - ) - else: - if c == d: - if bc[d][0]: - apply_essential_bc_stencil( - M, - axis=0, - ext=-1, - order=0, - identity=True, - ) - if bc[d][1]: - apply_essential_bc_stencil( - M, - axis=0, - ext=+1, - order=0, - identity=True, - ) - - M_arr = M.toarray() - - # create 1d solver for mass matrix - if is_circulant(M_arr): - solvercells += [FFTSolver(M_arr)] - else: - solvercells += [SparseSolver(M.tosparse())] - - # === NOTE: for KroneckerStencilMatrix being built correctly, 1d matrices must be local to process! === - periodic = femspaces[c].coeff_space.periods[d] - - n = femspaces[c].coeff_space.npts[d] - p = femspaces[c].coeff_space.pads[d] - s = femspaces[c].coeff_space.starts[d] - e = femspaces[c].coeff_space.ends[d] - - cart_decomp_1d = CartDecomposition( - domain_decompos_1d, - [n], - [[s]], - [[e]], - [p], - [1], - ) - - V_local = StencilVectorSpace(cart_decomp_1d) - - M_local = StencilMatrix(V_local, V_local) - - row_indices, col_indices = xp.nonzero(M_arr) - - for row_i, col_i in zip(row_indices, col_indices): - # only consider row indices on process - if row_i in range(V_local.starts[0], V_local.ends[0] + 1): - row_i_loc = row_i - s - - M_local._data[ - row_i_loc + p, - (col_i + p - row_i) % M_arr.shape[1], - ] = M_arr[row_i, col_i] - - # check if stencil matrix was built correctly - assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) - - matrixcells += [M_local.copy()] - # ======================================================================================================= - - if isinstance(self._femspace, TensorFemSpace): - matrixblocks += [ - KroneckerStencilMatrix( - self._femspace.coeff_space, - self._femspace.coeff_space, - *matrixcells, - ), - ] - solverblocks += [ - KroneckerLinearSolver( - self._femspace.coeff_space, - self._femspace.coeff_space, - solvercells, - ), - ] - else: - matrixblocks += [ - KroneckerStencilMatrix( - self._femspace.coeff_space[c], - self._femspace.coeff_space[c], - *matrixcells, - ), - ] - solverblocks += [ - KroneckerLinearSolver( - self._femspace.coeff_space[c], - self._femspace.coeff_space[c], - solvercells, - ), - ] - - # build final matrix and solver - if isinstance(self._femspace, TensorFemSpace): - self._matrix = matrixblocks[0] - self._solver = solverblocks[0] - else: - blocks = [ - [matrixblocks[0], None, None], - [None, matrixblocks[1], None], - [None, None, matrixblocks[2]], - ] - - self._matrix = BlockLinearOperator( - self._femspace.coeff_space, - self._femspace.coeff_space, - blocks=blocks, - ) - - sblocks = [ - [solverblocks[0], None, None], - [None, solverblocks[1], None], - [None, None, solverblocks[2]], - ] + # -------------------------------------- + # To be implemented by subclasses + # -------------------------------------- + def _build_approximation(self, bc: list | None) -> tuple[LinearOperator, LinearOperator]: + """ + Approximation of the core operator and its exact inverse. - self._solver = BlockLinearOperator( - self._femspace.coeff_space, - self._femspace.coeff_space, - blocks=sblocks, - ) + Parameters + ---------- + bc : list | None + Boundary conditions from ``_boundary_conditions``; None for no boundary conditions. - # save mass operator to be inverted (needed in solve method) - if apply_bc: - self._M = mass_operator.M0 - else: - self._M = mass_operator.M + Returns + ------- + matrix : LinearOperator + The approximation :math:`\\tilde A_c`. - self._is_composed = isinstance(self._M, ComposedLinearOperator) + solver : LinearOperator + Its exact inverse. + """ + raise NotImplementedError - # temporary vectors for dot product - if self._is_composed: - tmp_vectors = [] - for op in self._M.multiplicants[1:]: - tmp_vectors.append(op.codomain.zeros()) + def _core_diagonal(self) -> LinearOperator: + """Diagonal :math:`D` of the core operator as diagonal operator (for diagonal scaling).""" + raise NotImplementedError(f"{type(self).__name__} does not support diagonal scaling.") - self._tmp_vectors = tuple(tmp_vectors) - else: - self._tmp_vector = self._M.codomain.zeros() + def transpose(self, conjugate: bool = False) -> KroneckerPreconditioner: + raise NotImplementedError + # -------------------------------------- + # Common interface + # -------------------------------------- @property - def space(self): + def space(self) -> VectorSpace: """Stencil-/BlockVectorSpace or PolarDerhamSpace.""" return self._space @property - def matrix(self): - """Approximation of the input mass matrix as KroneckerStencilMatrix.""" + def matrix(self) -> LinearOperator: + """Approximation of the core operator (Kronecker products, block diagonal for vector-valued spaces).""" return self._matrix @property - def solver(self): - """KroneckerLinearSolver or BlockDiagonalSolver for exactly inverting the approximate mass matrix self.matrix.""" + def solver(self) -> LinearOperator: + """Exact inverse of the approximation self.matrix.""" return self._solver @property - def domain(self): + def domain(self) -> VectorSpace: """The domain of the linear operator - an element of Vectorspace""" return self._space @property - def codomain(self): + def codomain(self) -> VectorSpace: """The codomain of the linear operator - an element of Vectorspace""" return self._codomain @@ -382,16 +163,43 @@ def codomain(self): def dtype(self): return self._dtype - def transpose(self, conjugate=False): - """ - Returns the transposed operator. - """ - return MassMatrixPreconditioner(self._mass_operator.transpose(), self._apply_bc, self._dim_reduce) + @property + def diagonal_scaling(self) -> bool: + """Whether the approximation is corrected by diagonal scaling.""" + return self._diagonal_scaling + + def _build_scaling(self) -> tuple[LinearOperator, ...]: + """Diagonal scaling :math:`S = \\hat D^{1/2} D^{-1/2}` as one diagonal operator.""" + D = _local_diagonal(self._core_diagonal()) + D_hat = _local_diagonal(self._matrix) + S = [xp.where(d > 0, xp.sqrt(xp.abs(dh) / xp.where(d > 0, d, 1.0)), 1.0) for d, dh in zip(D, D_hat)] + return (_diagonal_operator(self._matrix.domain, S),) + + def _apply_inverse(self, x: Vector, out: Vector) -> Vector: + """Apply :math:`S \\tilde A_c^{-1} S` (the replacement of the core operator) to x.""" + if not self._scaling: + return self._solver.dot(x, out=out) + + buf = self._tmp_core + k = 0 + y = x + for op in self._scaling: + y = op.dot(y, out=buf[k]) + k = 1 - k + y = self._solver.dot(y, out=buf[k]) + k = 1 - k + for op in reversed(self._scaling[1:]): + y = op.dot(y, out=buf[k]) + k = 1 - k + return self._scaling[0].dot(y, out=out) @profile - def solve(self, rhs, out=None): + def solve(self, rhs: Vector, out: Vector | None = None) -> Vector: """ - Computes (B * E * M^(-1) * E^T * B^T) * rhs as an approximation for an inverse mass matrix. + Computes :math:`B E \\, S \\tilde A_c^{-1} S \\, E^T B^T` rhs, an approximation of :math:`A^{-1}` rhs. + + The operators of the composition are applied from right to left; the core operator + is replaced by :math:`S \\tilde A_c^{-1} S`. Parameters ---------- @@ -404,60 +212,145 @@ def solve(self, rhs, out=None): Returns ------- out : feectools.linalg.basic.Vector - The result of (B * E * M^(-1) * E^T * B^T) * rhs. + The result. """ - assert isinstance(rhs, Vector) assert rhs.space == self._space + if out is not None: + assert isinstance(out, Vector) + assert out.space == self._space - # successive dot products with all but last operator - if self._is_composed: - x = rhs - for i in range(len(self._tmp_vectors)): - y = self._tmp_vectors[-1 - i] - A = self._M.multiplicants[-1 - i] - if isinstance(A, (StencilMatrix, BlockLinearOperator)): - self.solver.dot(x, out=y) - else: - A.dot(x, out=y) - x = y - - # last operator - A = self._M.multiplicants[0] - if out is None: - out = A.dot(x) - else: - assert isinstance(out, Vector) - assert out.space == self._space - A.dot(x, out=out) - - else: - if out is None: - out = self._tmp_vector.copy() - self.solver.dot(rhs, out=out) - - return out + return _apply_composed(self._M, self._mass_index, self._apply_inverse, self._tmp_vectors, rhs, out) - def dot(self, v, out=None): + def dot(self, v: Vector, out: Vector | None = None) -> Vector: """Apply linear operator to Vector v. Result is written to Vector out, if provided.""" - assert isinstance(v, Vector) assert v.space == self.domain - - # newly created output vector - if out is None: - out = self.solve(v) - - # in-place dot-product (result is written to out) - else: + if out is not None: assert isinstance(out, Vector) assert out.space == self.codomain - self.solve(v, out=out) + return self.solve(v, out=out) - return out +class MassMatrixPreconditioner(KroneckerPreconditioner): + r""" + Preconditioner for inverting 3d weighted mass matrices. + + The mass matrix is approximated by a Kronecker product of 1d mass matrices + (block diagonal for vector-valued spaces), which is inverted exactly with a + :class:`~feectools.linalg.kron.KroneckerLinearSolver`: + + * In the direction ``dim_reduce``, the 1d mass matrix carries a 1d weight + obtained from the diagonal block ``(c, c)`` of the 3d weight: its value at + the mid point (0.5) of the other two directions, or its mean over them + (see ``weight_reduction``). With ``dim_reduce=None``, no direction is weighted + (mass matrix on the logical cube). + * In the other directions, the 1d mass matrices are unweighted. + * Essential boundary conditions of the mass operator are imposed on the 1d + matrices (identity rows). + + The preconditioner applies :math:`B E \tilde M^{-1} E^T B^T`, where + :math:`B E \dots E^T B^T` is the composition of the mass operator (if any) + and :math:`\tilde M` is the Kronecker approximation (see :class:`KroneckerPreconditioner`). + + Parameters + ---------- + mass_operator : WeightedMassOperator + The weighted mass operator for which the approximate inverse is needed. + + apply_bc : bool + Whether to include boundary operators. + + dim_reduce : int | None + Axis along which the weight is kept; None for unit weights in all directions. + + weight_reduction : str + How the weight is reduced in the other axes: ``"midpoint"`` (value at 0.5, default) + or ``"average"`` (mean over the axes). Which one gives the better preconditioner + depends on the weight. + + diagonal_scaling : bool + Whether to correct the approximation with the diagonals of the mass matrix and of + its approximation (see :class:`KroneckerPreconditioner`). + """ + + def __init__( + self, + mass_operator: WeightedMassOperator, + apply_bc: bool = True, + dim_reduce: int | None = 0, + weight_reduction: str = "midpoint", + diagonal_scaling: bool = False, + ): + assert dim_reduce is None or dim_reduce < 3 + assert weight_reduction in ("midpoint", "average"), f"Unknown weight_reduction {weight_reduction!r}." + self._dim_reduce = dim_reduce + self._weight_reduction = weight_reduction + + super().__init__(mass_operator, apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) -class MassMatrixDiagonalPreconditioner(LinearOperator): + def _build_approximation(self, bc): + mass_operator = self._mass_operator + derham = mass_operator.derham + logger.debug(f"{derham.num_elements = }, {derham.bcs = }, {derham.degree = }") + + def weight_1d(c: int, d: int) -> Callable | xp.ndarray: + if d == self._dim_reduce: + return _reduced_weight_1d(mass_operator.weights[c][c], c, d, derham, self._weight_reduction) + return _ones_1d + + return _kronecker_approximation(mass_operator, bc, weight_1d) + + def _core_diagonal(self): + return self._mass_operator.matrix.diagonal() + + @property + def dim_reduce(self) -> int | None: + """Axis along which the weight is kept (None: unit weights in all directions).""" + return self._dim_reduce + + @property + def weight_reduction(self) -> str: + """How the weight is reduced in the directions other than ``dim_reduce``: ``"midpoint"`` or ``"average"``.""" + return self._weight_reduction + + def update_mass_operator(self, mass_operator: WeightedMassOperator) -> None: + """ + Update the mass operator (same spaces and structure, e.g. new weights) to recycle the preconditioner. + + The Kronecker approximation is rebuilt only if it depends on the weights (``dim_reduce`` is + not None); the diagonal scaling is updated. + """ + assert isinstance(mass_operator, WeightedMassOperator) + assert mass_operator.domain == mass_operator.codomain, "Only square mass matrices can be inverted!" + assert mass_operator.domain == self.domain, "Update needs to have the same domain and codomain" + + self._mass_operator = mass_operator + + # the composition has the same structure as before, so the temporary vectors can be reused + M, mass_index, _ = _operator_to_invert(mass_operator, self._bc is not None) + assert mass_index == self._mass_index, "The updated mass operator must have the same structure." + self._M = M + + if self._dim_reduce is not None: + self._matrix, self._solver = self._build_approximation(self._bc) + if self._diagonal_scaling: + self._scaling = self._build_scaling() + + def transpose(self, conjugate: bool = False) -> MassMatrixPreconditioner: + """ + Returns the transposed operator. + """ + return MassMatrixPreconditioner( + self._mass_operator.transpose(), + self._apply_bc, + self._dim_reduce, + self._weight_reduction, + self._diagonal_scaling, + ) + + +class MassMatrixDiagonalPreconditioner(MassMatrixPreconditioner): r""" Preconditioner for inverting 3d weighted mass matrices. The mass matrix is approximated by @@ -467,6 +360,9 @@ class MassMatrixDiagonalPreconditioner(LinearOperator): Where $D$ is the diagonal of the matrix to invert, :math:`\hat M` is the mass matrix on the logical domain that is a Kronecker product (fastly inverted) and :math:`\hat D^{-1/2}` is the diagonal of :math:`\hat M`. + This is ``MassMatrixPreconditioner(mass_operator, apply_bc, dim_reduce=None, diagonal_scaling=True)``; + the class is kept for its name (e.g. in solver options). + Notes ----- @@ -481,407 +377,866 @@ class MassMatrixDiagonalPreconditioner(LinearOperator): Whether to include boundary operators. """ - def __init__(self, mass_operator, apply_bc=True): - assert isinstance(mass_operator, WeightedMassOperator) - assert mass_operator.domain == mass_operator.codomain, "Only square mass matrices can be inverted!" + def __init__(self, mass_operator: WeightedMassOperator, apply_bc: bool = True): + super().__init__(mass_operator, apply_bc=apply_bc, dim_reduce=None, diagonal_scaling=True) - self._mass_operator = mass_operator - self._femspace = mass_operator.domain_femspace - self._space = mass_operator.domain - self._dtype = mass_operator.dtype - self._codomain = mass_operator.codomain - self._domain = mass_operator.domain - self._apply_bc = apply_bc + def transpose(self, conjugate: bool = False) -> MassMatrixDiagonalPreconditioner: + """ + Returns the transposed operator. + """ + return MassMatrixDiagonalPreconditioner(self._mass_operator.transpose(), self._apply_bc) - # 3d Kronecker stencil matrices and solvers - solverblocks = [] - matrixblocks = [] - # collect TensorFemSpaces in a tuple - if isinstance(self._femspace, TensorFemSpace): - femspaces = (self._femspace,) - else: - femspaces = self._femspace.spaces +class StiffnessPreconditioner(KroneckerPreconditioner): + r""" + Preconditioner for the stabilized stiffness operators - n_comps = len(femspaces) - n_dims = self._femspace.ldim + .. math:: - assert n_dims == 3 # other dims not yet implemented + A = \mathbb d^T \mathbb M_{k+1} \mathbb d + \sigma \, \mathbb M_k \,, + \qquad \mathbb d \in \{\mathbb G, \mathbb C, \mathbb D\} \ (k = 0, 1, 2)\,, - # get boundary conditions list from BoundaryOperator in ComposedLinearOperator M0 of mass operator - if apply_bc and isinstance(mass_operator.M0, ComposedLinearOperator): - if isinstance(mass_operator.M0.multiplicants[-1], BoundaryOperator): - bc = mass_operator.M0.multiplicants[-1].bc - else: - apply_bc = False - bc = None - else: - apply_bc = False - bc = None - - # loop over components - for c in range(n_comps): - # 1d mass matrices and solvers - solvercells = [] - matrixcells = [] - - # loop over spatial directions - for d in range(n_dims): - fun = [[lambda e: xp.ones(e.size, dtype=float)]] - - # get 1D FEM space (serial, not distributed) and quadrature order - if femspaces[c].spaces[d].basis == "B": - femspace_1d_tensor = mass_operator.derham.H1_1d_serial[d] - else: - femspace_1d_tensor = mass_operator.derham.L2_1d_serial[d] - - domain_decompos_1d = femspace_1d_tensor.domain_decomposition - qu_order_1d = (mass_operator.derham.nquads[d],) - - M = WeightedMassOperator( - self._mass_operator.derham, - femspace_1d_tensor, - femspace_1d_tensor, - weights_info=fun, - nquads=qu_order_1d, - ) - M.assemble() - M = M.matrix - - # apply boundary conditions - if apply_bc: - if mass_operator._domain_symbolic_name not in ("H1H1H1", "H1vec"): - if femspaces[c].spaces[d].basis == "B": - if bc[d][0]: - apply_essential_bc_stencil( - M, - axis=0, - ext=-1, - order=0, - identity=True, - ) - if bc[d][1]: - apply_essential_bc_stencil( - M, - axis=0, - ext=+1, - order=0, - identity=True, - ) - else: - if c == d: - if bc[d][0]: - apply_essential_bc_stencil( - M, - axis=0, - ext=-1, - order=0, - identity=True, - ) - if bc[d][1]: - apply_essential_bc_stencil( - M, - axis=0, - ext=+1, - order=0, - identity=True, - ) - - M_arr = M.toarray() - - # create 1d solver for mass matrix - if is_circulant(M_arr): - solvercells += [FFTSolver(M_arr)] - else: - solvercells += [SparseSolver(M.tosparse())] - - # === NOTE: for KroneckerStencilMatrix being built correctly, 1d matrices must be local to process! === - periodic = femspaces[c].coeff_space.periods[d] - - n = femspaces[c].coeff_space.npts[d] - p = femspaces[c].coeff_space.pads[d] - s = femspaces[c].coeff_space.starts[d] - e = femspaces[c].coeff_space.ends[d] - - cart_decomp_1d = CartDecomposition( - domain_decompos_1d, - [n], - [[s]], - [[e]], - [p], - [1], - ) + i.e. :math:`\mathbb G^T \mathbb M_1 \mathbb G + \sigma \mathbb M_0`, + :math:`\mathbb C^T \mathbb M_2 \mathbb C + \sigma \mathbb M_1` and + :math:`\mathbb D^T \mathbb M_3 \mathbb D + \sigma \mathbb M_2`. - V_local = StencilVectorSpace(cart_decomp_1d) - - M_local = StencilMatrix(V_local, V_local) - - row_indices, col_indices = xp.nonzero(M_arr) - - for row_i, col_i in zip(row_indices, col_indices): - # only consider row indices on process - if row_i in range(V_local.starts[0], V_local.ends[0] + 1): - row_i_loc = row_i - s - - M_local._data[ - row_i_loc + p, - (col_i + p - row_i) % M_arr.shape[1], - ] = M_arr[row_i, col_i] - - # check if stencil matrix was built correctly - assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) - - matrixcells += [M_local.copy()] - # ======================================================================================================= - - if isinstance(self._femspace, TensorFemSpace): - matrixblocks += [ - KroneckerStencilMatrix( - self._femspace.coeff_space, - self._femspace.coeff_space, - *matrixcells, - ), - ] - solverblocks += [ - KroneckerLinearSolver( - self._femspace.coeff_space, - self._femspace.coeff_space, - solvercells, - ), - ] - else: - matrixblocks += [ - KroneckerStencilMatrix( - self._femspace.coeff_space[c], - self._femspace.coeff_space[c], - *matrixcells, - ), - ] - solverblocks += [ - KroneckerLinearSolver( - self._femspace.coeff_space[c], - self._femspace.coeff_space[c], - solvercells, - ), - ] - - # build final matrix and solver - if isinstance(self._femspace, TensorFemSpace): - self._matrix = matrixblocks[0] - self._solver = solverblocks[0] - else: - blocks = [ - [matrixblocks[0], None, None], - [None, matrixblocks[1], None], - [None, None, matrixblocks[2]], - ] - - self._matrix = BlockLinearOperator( - self._femspace.coeff_space, - self._femspace.coeff_space, - blocks=blocks, - ) + With unit weights (the operator on the logical cube), the mass matrices are + block-diagonal Kronecker products of 1d mass matrices :math:`M_d`, and each diagonal block of + :math:`A` is a sum of Kronecker products, - sblocks = [ - [solverblocks[0], None, None], - [None, solverblocks[1], None], - [None, None, solverblocks[2]], - ] - self._solver = BlockLinearOperator( - self._femspace.coeff_space, - self._femspace.coeff_space, - blocks=sblocks, - ) + .. math:: - # save mass operator to be inverted (needed in solve method) - if apply_bc: - self._M = mass_operator.M0 - else: - self._M = mass_operator.M + A_{cc} \approx \sum_{d \in N_c} M_1 \otimes \dots \otimes S_d \otimes \dots \otimes M_3 + + \sigma \, M_1 \otimes M_2 \otimes M_3 \,, - self._is_composed = isinstance(self._M, ComposedLinearOperator) + with the 1d stiffness matrices :math:`S_d = \mathbb d_d^T M^D_d \mathbb d_d` in the directions + :math:`N_c` in which component c is differentiated (its B-spline directions). These blocks + are inverted exactly by fast diagonalization + (:class:`~feectools.linalg.kron.KroneckerSumSolver`). For the gradient this is the exact + inverse on the logical cube. For curl and div, the off-diagonal blocks are neglected + (block Jacobi, :math:`P_{BJ}`). - # temporary vectors for dot product - if self._is_composed: - tmp_vectors = [] - for op in self._M.multiplicants[1:]: - tmp_vectors.append(op.codomain.zeros()) + Block Jacobi overestimates :math:`A` on the kernel of :math:`\mathbb d` (gradients for + curl, curls for div), where :math:`A = \sigma \mathbb M_k`. With ``kernel_correction`` + (default, needs :math:`\sigma > 0`), the preconditioner is - self._tmp_vectors = tuple(tmp_vectors) - else: - self._tmp_vector = self._M.codomain.zeros() + .. math:: - # Need to assemble the logical mass matrix to extract the coefficients - fun = [ - [(lambda e1, e2, e3: xp.ones_like(e1, dtype=float)) if i == j else None for j in range(3)] for i in range(3) - ] - log_M = WeightedMassOperator( - self._mass_operator.derham, - self._femspace, - self._femspace, - weights_info=fun, - ) - log_M.assemble() - self._logM_srqt_diag = log_M.matrix.diagonal(sqrt=True) - self._M_invsrqt_diag = self._mass_operator.matrix.diagonal(inverse=True, sqrt=True) + P = P_{BJ} + \sigma^{-1} \, \mathbb d_- \, P_- \, \mathbb d_-^T \,, - self._tmp_vector_no_bc = [self._mass_operator.matrix.codomain.zeros() for i in range(2)] + where :math:`\mathbb d_-` is the previous derivative (:math:`\mathbb G` for curl, + :math:`\mathbb C` for div) and :math:`P_-` approximates the inverse of + :math:`\mathbb d_-^T \mathbb M_k \mathbb d_-` (the grad preconditioner for curl, the curl + block Jacobi for div). On the range of :math:`\mathbb d_-^T` the kernel of + :math:`\mathbb d_-` does not matter, since :math:`\mathbb G^T \mathbb C^T = 0`. - @property - def space(self): - """Stencil-/BlockVectorSpace or PolarDerhamSpace.""" - return self._space + The geometry is included with ``weights="average"`` (default) and by diagonal scaling + (default, see :class:`KroneckerPreconditioner`). For the weights, fast diagonalization needs the same 1d mass matrix in + each direction for all terms, so the weights :math:`w` of the mass matrices are approximated + by one common separable shape :math:`\prod_d \phi_d(\eta_d)` (mean 1; :math:`\phi_d` is the mean + over the other directions, averaged over the diagonal blocks of :math:`\mathbb M_{k+1}`), times + the mean of the weight of each term (the block of :math:`\mathbb M_{k+1}` for the stiffness + term of each direction, the block of :math:`\mathbb M_k` for the mass term). This is exact + e.g. for constant but anisotropic weights. + + Polar splines are not supported yet. + + Parameters + ---------- + mass_ops : WeightedMassOperators + The mass operators :math:`\mathbb M_k`. + + derivative : str + ``"grad"``, ``"curl"`` or ``"div"``. + + sigma : float + Coefficient of the mass term. + + apply_bc : bool + Whether to include boundary operators. + + diagonal_scaling : bool + Whether to correct the approximation with the diagonals of :math:`A` and of its + approximation. + + kernel_correction : bool + For curl and div: whether to add the correction on the kernel of the derivative (needs + ``sigma > 0``). Ignored for grad. + + weights : str + ``"average"`` (separable approximation of the weights, see above; default) or + ``"unit"`` (operator on the logical cube). + """ + + _FORMS = {"grad": 0, "curl": 1, "div": 2} + + def __init__( + self, + mass_ops: WeightedMassOperators, + derivative: str = "grad", + sigma: float = 0.0, + apply_bc: bool = True, + diagonal_scaling: bool = True, + kernel_correction: bool = True, + weights: str = "average", + ): + assert derivative in self._FORMS, f"derivative must be one of {tuple(self._FORMS)}, got {derivative!r}." + assert sigma >= 0.0 + assert weights in ("unit", "average"), f"weights must be 'unit' or 'average', got {weights!r}." + kernel_correction = kernel_correction and derivative != "grad" + if kernel_correction and sigma == 0.0: + raise ValueError("The kernel correction needs sigma > 0 (or kernel_correction=False).") + + derham = mass_ops.derham + if derham.polar_splines: + raise NotImplementedError("StiffnessPreconditioner does not support polar splines yet.") + + k = self._FORMS[derivative] + self._mass_ops = mass_ops + self._derivative = derivative + self._sigma = sigma + self._weights = weights + self._mass_mid = getattr(mass_ops, f"M{k + 1}") + # tensor-product derivative without boundary operators (no polar splines) + self._d = getattr(derham, f"{derivative}_bcfree") + + super().__init__(getattr(mass_ops, f"M{k}"), apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) + + # correction on the kernel of the derivative: sigma^{-1} d_prev P_prev d_prev^T + self._kernel_correction = None + if kernel_correction: + prev = "grad" if derivative == "curl" else "curl" + d_prev = getattr(derham, prev) + P_prev = StiffnessPreconditioner( + mass_ops, + prev, + apply_bc=apply_bc, + diagonal_scaling=diagonal_scaling, + kernel_correction=False, + weights=weights, + ) + self._kernel_correction = (1.0 / sigma) * (d_prev @ P_prev @ d_prev.T) + self._tmp_kernel = self.codomain.zeros() @property - def matrix(self): - """Mass matrix on the logical domain as KroneckerStencilMatrix.""" - return self._matrix + def derivative(self) -> str: + """``"grad"``, ``"curl"`` or ``"div"``.""" + return self._derivative @property - def solver(self): - """KroneckerLinearSolver or BlockDiagonalSolver for exactly inverting the approximate mass matrix self.matrix.""" - return self._solver + def sigma(self) -> float: + """Coefficient of the mass term.""" + return self._sigma @property - def domain(self): - """The domain of the linear operator - an element of Vectorspace""" - return self._space + def weights(self) -> str: + """``"unit"`` or ``"average"``.""" + return self._weights @property - def codomain(self): - """The codomain of the linear operator - an element of Vectorspace""" - return self._codomain + def kernel_correction(self) -> LinearOperator | None: + """The correction on the kernel of the derivative, or None.""" + return self._kernel_correction + + def solve(self, rhs: Vector, out: Vector | None = None) -> Vector: + """Apply the preconditioner (block Jacobi part plus kernel correction, if any) to rhs.""" + out = super().solve(rhs, out=out) + if self._kernel_correction is not None: + out += self._kernel_correction.dot(rhs, out=self._tmp_kernel) + return out @property - def dtype(self): - return self._dtype + def core_operator(self) -> LinearOperator: + """The core operator :math:`\\mathbb d^T \\mathbb M_{k+1} \\mathbb d + \\sigma \\mathbb M_k` on the tensor-product spaces.""" + A = self._d.T @ self._mass_mid.matrix @ self._d + if self._sigma != 0.0: + A = A + self._sigma * self._mass_operator.matrix + return A + + def _build_approximation(self, bc): + femspace = self._femspace + is_scalar = isinstance(femspace, TensorFemSpace) + comps = (femspace,) if is_scalar else femspace.spaces + + # directions in which each component is differentiated, and the codomain block + stiff_dirs = _derivative_directions(self._d) + + # weights: common shape per direction and means per block (unit or separable average) + n_mid = len(self._mass_mid.weights) + n_dom = len(self._mass_operator.weights) + if self._weights == "average": + mid_weights = [self._mass_mid.weights[r][r] for r in range(n_mid)] + shapes, mid_means = _separable_weight(mid_weights, self._mass_operator.derham) + _, dom_means = _separable_weight( + [self._mass_operator.weights[c][c] for c in range(n_dom)], self._mass_operator.derham + ) + else: + shapes, mid_means, dom_means = [_ones_1d] * 3, [1.0] * n_mid, [1.0] * n_dom - def update_mass_operator(self, mass_operator): - """Update the mass operator to enable recycling the preconditioner""" - assert isinstance(mass_operator, WeightedMassOperator) - assert mass_operator.domain == mass_operator.codomain, "Only square mass matrices can be inverted!" - assert mass_operator.domain == self.domain, "Update needs to have the same domain and codomain" + matrixblocks = [] + solverblocks = [] + for c, comp in enumerate(comps): + coeff_space = femspace.coeff_space if is_scalar else femspace.coeff_space[c] + + stiffness, mass, local_mass, local_stiffness = [], [], [], [] + for d in range(3): + basis = comp.spaces[d].basis + M_d, domain_decomposition = _dense_mass_1d(self._mass_operator, basis, d, shapes[d]) + + S_d = None + if d in stiff_dirs[c]: + assert basis == "B", "Only B-spline directions are differentiated." + M_D, _ = _dense_mass_1d(self._mass_operator, "M", d, shapes[d]) + D_d = _difference_matrix_1d(M_d.shape[0], M_D.shape[0]) + S_d = mid_means[stiff_dirs[c][d]] * (D_d.T @ M_D @ D_d) + + if bc is not None and basis == "B": + M_d = _apply_bc_dense(M_d, bc[d]) + S_d = None if S_d is None else _apply_bc_dense(S_d, bc[d]) + + stiffness.append(S_d) + mass.append(M_d) + local_mass.append(_process_local_matrix_1d(xp.asarray(M_d), comp.coeff_space, d, domain_decomposition)) + local_stiffness.append( + None + if S_d is None + else _process_local_matrix_1d(xp.asarray(S_d), comp.coeff_space, d, domain_decomposition) + ) - if self._is_composed: - if self._apply_bc: - assert isinstance(mass_operator.M0, ComposedLinearOperator) - else: - assert isinstance(mass_operator.M, ComposedLinearOperator) + sigma_c = self._sigma * dom_means[c] + solverblocks.append(KroneckerSumSolver(coeff_space, stiffness, mass, sigma=sigma_c)) + matrixblocks.append(_kronecker_sum(coeff_space, local_stiffness, local_mass, sigma_c)) - self._mass_operator = mass_operator + if is_scalar: + return matrixblocks[0], solverblocks[0] + return ( + _block_diagonal(femspace.coeff_space, matrixblocks), + _block_diagonal(femspace.coeff_space, solverblocks), + ) - if self._apply_bc: - self._M = mass_operator.M0 - else: - self._M = mass_operator.M - self._M_invsrqt_diag = self._mass_operator.matrix.diagonal(inverse=True, sqrt=True, out=self._M_invsrqt_diag) + def _core_diagonal(self): + space = self._matrix.domain + return _diagonal_operator(space, _probe_diagonal(self.core_operator, space)) - def transpose(self, conjugate=False): - """ - Returns the transposed operator. - """ - return MassMatrixDiagonalPreconditioner(self._mass_operator.transpose(), self._apply_bc) + def transpose(self, conjugate: bool = False) -> StiffnessPreconditioner: + """The operator is symmetric, so is the preconditioner.""" + return self - def _solve_no_bc(self, rhs, out): - r""" - Computes M^(-1) * rhs as an approximation for an inverse mass matrix. - With $M = D^{1/2} * \hat D^{-1/2} * \hat M * \hat D^{-1/2} * D^{1/2}$ - Should only be called by the solve method that will handle the bcs. - Parameters - ---------- - rhs : feectools.linalg.basic.Vector - The right-hand side vector. +# -------------------------------------------------------------------------------------- +# Helper functions for building Kronecker approximations of mass matrices +# -------------------------------------------------------------------------------------- +def _ones_1d(e: xp.ndarray) -> xp.ndarray: + """Unit weight for 1d mass matrices.""" + return xp.ones(e.size, dtype=float) - out : feectools.linalg.basic.Vector - The output vector will be written into this vector in-place. - Returns - ------- - out : feectools.linalg.basic.Vector - The result of M^(-1) * rhs. - """ +def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: + """ + Local (rows owned by this process) diagonal of ``op``, one array per block of its domain. - assert isinstance(rhs, Vector) - assert rhs.space == self._mass_operator.matrix.domain + Supported are StencilMatrix, StencilDiagonalMatrix, KroneckerStencilMatrix (1d factors), + sums, scalings and BlockLinearOperators of these (missing diagonal blocks give zeros). + """ + if isinstance(op, BlockLinearOperator): + diags = [] + for i, V in enumerate(op.domain.spaces): + if op[i, i] is None: + # zero block (e.g. a mass matrix block with zero weight) + diags.append(xp.zeros(tuple(e - s + 1 for s, e in zip(V.starts, V.ends)), dtype=float)) + else: + diags.extend(_local_diagonal(op[i, i])) + return diags + if isinstance(op, StencilDiagonalMatrix): + return [op._data] + if isinstance(op, StencilMatrix): + return [op.diagonal()._data] + if isinstance(op, KroneckerStencilMatrix): + # outer product of the diagonals of the (process-local) factors on the local rows + diags = [] + for A, s, e in zip(op.mats, op.codomain.starts, op.codomain.ends): + assert A.domain.ndim == 1, "Only 1d factors are supported." + # the factor may own more rows than the codomain on this process + off = A.codomain.pads[0] * A.codomain.shifts[0] + s - A.codomain.starts[0] + diags.append(A._data[off : off + e - s + 1, A.pads[0]]) + out = diags[0] + for d in diags[1:]: + out = xp.multiply.outer(out, d) + return [out] + if isinstance(op, SumLinearOperator): + parts = [_local_diagonal(a) for a in op.addends] + return [sum(blocks) for blocks in zip(*parts)] + if isinstance(op, ScaledLinearOperator): + return [op.scalar * a for a in _local_diagonal(op.operator)] + raise NotImplementedError(f"Diagonal of {type(op).__name__} is not supported.") + + +def _diagonal_operator(space: VectorSpace, diags: list[xp.ndarray]) -> LinearOperator: + """Diagonal operator on ``space`` (Stencil- or BlockVectorSpace) with the given local diagonals.""" + if isinstance(space, BlockVectorSpace): + n = len(space.spaces) + blocks = [ + [StencilDiagonalMatrix(V, V, diags[i]) if i == j else None for j in range(n)] + for i, V in enumerate(space.spaces) + ] + return BlockLinearOperator(space, space, blocks=blocks) + return StencilDiagonalMatrix(space, space, diags[0]) - # M^-1 ~ D^{-1/2} \hat D^{1/2} \hat M ^{-1} \hat D^{1/2} D^{-1/2} - Dmr = self._M_invsrqt_diag.dot(rhs, out=self._tmp_vector_no_bc[0]) - DhDmr = self._logM_srqt_diag.dot(Dmr, out=self._tmp_vector_no_bc[1]) - invMr = self.solver.dot(DhDmr, out=self._tmp_vector_no_bc[0]) - DhiMr = self._logM_srqt_diag.dot(invMr, out=self._tmp_vector_no_bc[1]) - out = self._M_invsrqt_diag.dot(DhiMr, out=out) - return out +def _boundary_conditions(mass_operator: WeightedMassOperator, apply_bc: bool) -> list | None: + """ + Boundary conditions of the mass operator, taken from the BoundaryOperator + at the right end of the composition ``mass_operator.M0``. - @profile - def solve(self, rhs, out=None): - r""" - Computes :math:`(B * E * M^{-1} * E^T * B^T) * rhs` as an approximation for an inverse mass matrix, - with :math:`M = D^{1/2} * \hat D^{-1/2} * \hat M * \hat D^{-1/2} * D^{1/2}`. + Returns + ------- + list | None + ``bc[d] = (left, right)`` flags of essential boundary conditions per direction, + or None if ``apply_bc`` is False or ``M0`` has no BoundaryOperator. + """ + if not apply_bc or not isinstance(mass_operator.M0, ComposedLinearOperator): + return None + last = mass_operator.M0.multiplicands[-1] + if not isinstance(last, BoundaryOperator): + return None + return last.bc - Parameters - ---------- - rhs : Vector - The right-hand side vector. - out : Vector, optional - If given, the output vector will be written into this vector in-place. +def _operator_to_invert( + mass_operator: WeightedMassOperator, apply_bc: bool +) -> tuple[ComposedLinearOperator | LinearOperator, int | None, tuple[Vector, ...]]: + """ + The operator approximately inverted by the preconditioners: ``M0`` (with boundary + operators) if ``apply_bc``, else ``M``. - Returns - ------- - out : Vector - The result of :math:`(B * E * M^{-1} * E^T * B^T) * rhs`. - """ + Returns + ------- + M : ComposedLinearOperator | LinearOperator + The operator. - assert isinstance(rhs, Vector) - assert rhs.space == self._space + mass_index : int | None + Position of the mass matrix ``mass_operator.matrix`` in ``M.multiplicands``, or + None if ``M`` is not composed (then ``M`` is the mass matrix itself). - # successive dot products with all but last operator - if self._is_composed: - x = rhs - for i in range(len(self._tmp_vectors)): - y = self._tmp_vectors[-1 - i] - A = self._M.multiplicants[-1 - i] - if isinstance(A, (StencilMatrix, BlockLinearOperator)): - self._solve_no_bc(x, out=y) - else: - A.dot(x, out=y) - x = y - - # last operator - A = self._M.multiplicants[0] - if out is None: - out = A.dot(x) - else: - assert isinstance(out, Vector) - assert out.space == self._space - A.dot(x, out=out) + tmp_vectors : tuple[Vector, ...] + Temporary vectors for ``_apply_composed``: the codomains of all but the first + factor of ``M`` (empty if ``M`` is not composed). + """ + M = mass_operator.M0 if apply_bc else mass_operator.M + if not isinstance(M, ComposedLinearOperator): + return M, None, () + + mass_index = next((k for k, op in enumerate(M.multiplicands) if op is mass_operator.matrix), None) + assert mass_index is not None, "The mass matrix was not found in the composed mass operator." + tmp_vectors = tuple(op.codomain.zeros() for op in M.multiplicands[1:]) + return M, mass_index, tmp_vectors + + +def _apply_composed( + M: LinearOperator, + mass_index: int | None, + apply_inverse: Callable[[Vector, Vector], Vector], + tmp_vectors: tuple[Vector, ...], + rhs: Vector, + out: Vector | None, +) -> Vector: + """ + Apply the operator ``M`` (see ``_operator_to_invert``) to ``rhs``, with the mass matrix + replaced by ``apply_inverse(x, out)`` (an approximate inverse). + The factors of a composition are applied from right to left; the result is written + to ``out`` (allocated if None). + """ + if mass_index is None: + if out is None: + out = M.codomain.zeros() + return apply_inverse(rhs, out) + + ops = M.multiplicands + if out is None: + out = ops[0].codomain.zeros() + + x = rhs + for k in range(len(ops) - 1, -1, -1): + y = out if k == 0 else tmp_vectors[k - 1] + if k == mass_index: + apply_inverse(x, y) else: - if out is None: - out = self._tmp_vector.copy() - self._solve_no_bc(rhs, out=out) + ops[k].dot(x, out=y) + x = y + return out - return out - def dot(self, v, out=None): - """Apply linear operator to Vector v. Result is written to Vector out, if provided.""" +# number of Gauss-Legendre points per direction for averaging callable weights +_N_AVG_CALLABLE = 16 - assert isinstance(v, Vector) - assert v.space == self.domain - # newly created output vector - if out is None: - out = self.solve(v) +def _reduced_weight_1d(weight, c: int, d: int, derham, reduction: str = "midpoint") -> Callable | xp.ndarray: + r""" + 1d weight along direction ``d``, obtained from the 3d weight by reducing the other two directions + :math:`i, j \neq d`: + + * ``"midpoint"``: the weight at the mid point, :math:`\bar w(\eta_d) = w(\eta_d; \eta_i = \eta_j = 0.5)`, + * ``"average"``: the mean, :math:`\bar w(\eta_d) = \int_0^1 \int_0^1 w \, \textnormal d\eta_i \, \textnormal d\eta_j`. + + Parameters + ---------- + weight : callable | xp.ndarray | None + Block ``(c, c)`` of the weights of the 3d mass operator: a function of the three + logical coordinates, its values at the local quadrature points, or None (unit weight). + + c, d : int + Component and direction (``c`` is only used for logging). + + derham : Derham + Discrete de Rham sequence of the mass operator. - # in-place dot-product (result is written to out) + reduction : str + ``"midpoint"`` or ``"average"``. + + Returns + ------- + Callable | xp.ndarray + For a callable weight, a function of the 1d points (for ``"average"``, the mean is + approximated with Gauss-Legendre quadrature on ``_N_AVG_CALLABLE`` points per direction). + For an array weight, the reduced weight at all global 1d quadrature points; this is + collective on ``derham.comm``. For the mid point (array weight) it is the value at the + global mid quadrature point; for the mean, Gauss quadrature of the derham is used. + """ + assert reduction in ("midpoint", "average"), f"Unknown weight reduction {reduction!r}." + n_dims = 3 + i, j = (k for k in range(n_dims) if k != d) + + if weight is None: + return _ones_1d + + if callable(weight): + if reduction == "midpoint": + + def fun(e): + # evaluate the 3d weight on the "meshgrid" (0.5, ..., e, ..., 0.5) + s = e.shape[0] + f = e.reshape(tuple([1 if k != d else s for k in range(n_dims)])) + return xp.atleast_1d( + weight( + *[xp.array(xp.full_like(f, 0.5)) if k != d else xp.array(f) for k in range(n_dims)] + ).squeeze(), + ) + + return fun + + # Gauss-Legendre points and weights on [0, 1] + nodes, wts = np.polynomial.legendre.leggauss(_N_AVG_CALLABLE) + pts = xp.asarray((nodes + 1.0) / 2.0) + wts = xp.asarray(wts / 2.0) + + def fun(e): + # integrate over direction i on the grid (e, pts) for each point in direction j + e_d, e_i = xp.meshgrid(xp.ravel(e), pts, indexing="ij") + total = xp.zeros(e_d.shape[0], dtype=float) + args = [None] * n_dims + args[d] = e_d + args[i] = e_i + for eta_j, w_j in zip(pts, wts): + args[j] = xp.full_like(e_d, eta_j) + total += w_j * (xp.reshape(weight(*args), e_d.shape) @ wts) + return total + + return fun + + if isinstance(weight, xp.ndarray): + logger.debug(f"{weight.shape = } for component {c} and direction {d}.") + dom_dec = derham.domain_decomposition + nq_d = derham.nquads[d] + start = dom_dec.starts[d] * nq_d + stop = start + weight.shape[d] + + # every process writes its contribution at its rows along d; the sum over all processes + # gives the reduced weight (row 0) and the normalization (row 1) + sums = xp.zeros((2, derham.num_elements[d] * nq_d), dtype=float) + + if reduction == "midpoint": + # only the process owning the global mid point in directions i and j contributes + # (exactly one process per slab along d) + mid = {} + for k in (i, j): + nq_k = weight.shape[k] // dom_dec.local_ncells[k] + mid[k] = (derham.num_elements[k] * nq_k) // 2 - dom_dec.starts[k] * nq_k + if all(0 <= mid[k] < weight.shape[k] for k in (i, j)): + cut = tuple(slice(None) if k == d else mid[k] for k in range(n_dims)) + sums[0, start:stop] = weight[cut] + sums[1, start:stop] = 1.0 else: - assert isinstance(out, Vector) - assert out.space == self.codomain - self.solve(v, out=out) + # Gauss weights at the local quadrature points (same grid as the array weight) + qwts = [xp.ravel(w) for w in derham.spline_attributes["H1"].quad_grid_wts[0]] + assert weight.shape == tuple(w.size for w in qwts), ( + f"Array weight of shape {weight.shape} does not match the local quadrature grid." + ) + weighted = weight * qwts[i].reshape([-1 if k == i else 1 for k in range(n_dims)]) + weighted = weighted * qwts[j].reshape([-1 if k == j else 1 for k in range(n_dims)]) + sums[0, start:stop] = weighted.sum(axis=(i, j)) + sums[1, start:stop] = qwts[i].sum() * qwts[j].sum() - return out + comm = derham.comm + if not isinstance(comm, (SerialComm, type(None))): + local_sums = sums.copy() + comm.Allreduce(local_sums, sums, op=MPI.SUM) + + return sums[0] / sums[1] + + raise TypeError(f"weights needs to be callable, xp.ndarray or None but is {type(weight)}") + + +def _has_essential_bc(mass_operator: WeightedMassOperator, basis: str, c: int, d: int) -> bool: + """ + Whether the boundary conditions of direction ``d`` apply to the 1d matrix of component ``c``. + + For H1 vector spaces (``H1H1H1``, ``H1vec``) this is the case if ``c == d``; for the other + spaces, in all directions with B-splines (``basis == "B"``). + """ + if mass_operator._domain_symbolic_name in ("H1H1H1", "H1vec"): + return c == d + return basis == "B" + + +def _mass_matrix_1d( + mass_operator: WeightedMassOperator, basis: str, d: int, weight: Callable | xp.ndarray +) -> tuple[StencilMatrix, DomainDecomposition]: + """ + Assemble the 1d mass matrix in direction ``d`` on the serial (not distributed) 1d space. + + Parameters + ---------- + basis : str + ``"B"`` for B-splines (H1), else M-splines (L2). + + weight : Callable | xp.ndarray + 1d weight, see ``_reduced_weight_1d``. + + Returns + ------- + M : StencilMatrix + The 1d mass matrix. + + domain_decomposition : DomainDecomposition + Domain decomposition of the serial 1d space. + """ + derham = mass_operator.derham + femspace_1d = derham.H1_1d_serial[d] if basis == "B" else derham.L2_1d_serial[d] + + M = WeightedMassOperator( + derham, + femspace_1d, + femspace_1d, + weights_info=[[weight]], + nquads=(derham.nquads[d],), + ) + M.assemble() + return M.matrix, femspace_1d.domain_decomposition + + +def _apply_bc_1d(M: StencilMatrix, bc_d: tuple[bool, bool]) -> None: + """Impose essential boundary conditions (identity rows) at the left/right end of the 1d matrix M (in place).""" + for ext, is_essential in zip((-1, +1), bc_d): + if is_essential: + apply_essential_bc_stencil(M, axis=0, ext=ext, order=0, identity=True) + + +def _solver_1d(M: StencilMatrix, M_arr: xp.ndarray) -> FFTSolver | SparseSolver: + """Direct solver for the 1d matrix M (dense copy M_arr): FFT if circulant, else sparse LU.""" + if is_circulant(M_arr): + return FFTSolver(M_arr) + return SparseSolver(M.tosparse()) + + +def _process_local_matrix_1d( + M_arr: xp.ndarray, coeff_space: StencilVectorSpace, d: int, domain_decomposition: DomainDecomposition +) -> StencilMatrix: + """ + Process-local 1d factor of a KroneckerStencilMatrix on ``coeff_space``. + + The factor lives on a 1d space without communicator that owns the same rows + (``starts[d]`` to ``ends[d]``) as ``coeff_space`` in direction ``d`` on this process, + as required by KroneckerStencilMatrix. + + Parameters + ---------- + M_arr : xp.ndarray + The global 1d matrix as dense array. + + coeff_space : StencilVectorSpace + The (distributed) 3d coefficient space. + + d : int + Direction. + + domain_decomposition : DomainDecomposition + Domain decomposition of the serial 1d space. + """ + n = coeff_space.npts[d] + p = coeff_space.pads[d] + s = coeff_space.starts[d] + e = coeff_space.ends[d] + + cart_1d = CartDecomposition(domain_decomposition, [n], [[s]], [[e]], [p], [1]) + V_local = StencilVectorSpace(cart_1d) + M_local = StencilMatrix(V_local, V_local) + + # copy the rows owned by this process: entry (i, j) is stored at row i - s + p and diagonal + # index k + p with the offset k = j - i; in periodic directions k is taken nearest to 0 + # (in [-n/2, n/2)), so that the main diagonal is at k = 0 also if 2p + 1 > n + rows, cols = xp.nonzero(M_arr) + on_process = (rows >= s) & (rows <= e) + rows, cols = rows[on_process], cols[on_process] + k = cols - rows + if coeff_space.periods[d]: + k = (k + n // 2) % n - n // 2 + assert xp.all(xp.abs(k) <= p), "Entries outside the band of the stencil matrix." + M_local._data[rows - s + p, k + p] = M_arr[rows, cols] + + # check if stencil matrix was built correctly + assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) + + return M_local + + +def _kronecker_approximation( + mass_operator: WeightedMassOperator, + bc: list | None, + weight_1d: Callable[[int, int], Callable | xp.ndarray], +) -> tuple[KroneckerStencilMatrix | BlockLinearOperator, KroneckerLinearSolver | BlockLinearOperator]: + """ + Approximate the mass matrix by Kronecker products of 1d mass matrices + (block diagonal for vector-valued spaces) and build their exact inverse. + + Parameters + ---------- + mass_operator : WeightedMassOperator + The mass operator to approximate. + + bc : list | None + Boundary conditions from ``_boundary_conditions``; None for no boundary conditions. + + weight_1d : callable + ``weight_1d(c, d)`` returns the 1d weight of component ``c`` in direction ``d``. + + Returns + ------- + matrix : KroneckerStencilMatrix | BlockLinearOperator + The approximate mass matrix. + + solver : KroneckerLinearSolver | BlockLinearOperator + Its exact inverse. + """ + femspace = mass_operator.domain_femspace + is_scalar = isinstance(femspace, TensorFemSpace) + femspaces = (femspace,) if is_scalar else femspace.spaces + n_dims = femspace.ldim + + matrixblocks = [] + solverblocks = [] + for c, femspace_c in enumerate(femspaces): + coeff_space = femspace.coeff_space if is_scalar else femspace.coeff_space[c] + + # 1d mass matrices (process-local) and solvers in each direction + matrixcells = [] + solvercells = [] + for d in range(n_dims): + basis = femspace_c.spaces[d].basis + M, domain_decomposition = _mass_matrix_1d(mass_operator, basis, d, weight_1d(c, d)) + + if bc is not None and _has_essential_bc(mass_operator, basis, c, d): + _apply_bc_1d(M, bc[d]) + + M_arr = M.toarray() + solvercells.append(_solver_1d(M, M_arr)) + matrixcells.append(_process_local_matrix_1d(M_arr, femspace_c.coeff_space, d, domain_decomposition)) + + matrixblocks.append(KroneckerStencilMatrix(coeff_space, coeff_space, *matrixcells)) + solverblocks.append(KroneckerLinearSolver(coeff_space, coeff_space, solvercells)) + + if is_scalar: + return matrixblocks[0], solverblocks[0] + return ( + _block_diagonal(femspace.coeff_space, matrixblocks), + _block_diagonal(femspace.coeff_space, solverblocks), + ) + + +def _block_diagonal(space: BlockVectorSpace, blocks: list[LinearOperator]) -> BlockLinearOperator: + """Block-diagonal operator on ``space`` with the given diagonal blocks.""" + n = len(blocks) + return BlockLinearOperator( + space, + space, + blocks=[[blocks[i] if i == j else None for j in range(n)] for i in range(n)], + ) + + +# -------------------------------------------------------------------------------------- +# Helper functions for stiffness preconditioners +# -------------------------------------------------------------------------------------- +def _derivative_directions(derivative: LinearOperator) -> list[dict[int, int]]: + """ + For each component c of the domain of a (tensor-product) derivative, the directions in + which it is differentiated, mapped to the codomain block of that derivative. + + The blocks ``(r, c)`` of ``derivative`` (a BlockLinearOperator, or a single operator) are + DirectionalDerivativeOperators, possibly negated. + """ + if isinstance(derivative, BlockLinearOperator): + nrows, ncols = derivative.n_block_rows, derivative.n_block_cols + blocks = {(r, c): derivative[r, c] for r in range(nrows) for c in range(ncols) if derivative[r, c] is not None} + else: + ncols = 1 + blocks = {(0, 0): derivative} + + directions = [{} for _ in range(ncols)] + for (r, c), op in blocks.items(): + assert isinstance(op, DirectionalDerivativeOperator), f"Unexpected block {type(op).__name__} in derivative." + assert op.diffdir not in directions[c], "Component differentiated twice in the same direction." + directions[c][op.diffdir] = r + return directions + + +def _separable_weight(weights: list, derham) -> tuple[list[xp.ndarray], list[float]]: + r""" + Common separable shape of 3d weights, :math:`w_i \approx \bar w_i \prod_d \phi_d(\eta_d)`. + + Parameters + ---------- + weights : list + 3d weights (callable, array at the local quadrature points, or None for 1). + + derham : Derham + Discrete de Rham sequence (quadrature grid). + + Returns + ------- + shapes : list[xp.ndarray] + :math:`\phi_d` at the global 1d quadrature points: the mean over the other directions, + divided by the mean, averaged over the weights (mean 1). + + means : list[float] + The means :math:`\bar w_i` of the weights over the logical cube. + + Collective on ``derham.comm`` for array weights. + """ + grids = [derham.H1_1d_serial[d].get_assembly_grids(derham.nquads[d])[0] for d in range(3)] + pts = [xp.asarray(np.ravel(g.points)) for g in grids] + wts = [xp.asarray(np.ravel(g.weights)) for g in grids] + + profiles, means = [], [] + for w in weights: + prof = [] + for d in range(3): + p = _reduced_weight_1d(w, 0, d, derham, "average") + prof.append(xp.asarray(p(pts[d]) if callable(p) else p, dtype=float)) + mean = float(xp.sum(prof[0] * wts[0]) / xp.sum(wts[0])) + assert mean > 0.0, "The weights must have a positive mean." + profiles.append(prof) + means.append(mean) + + shapes = [sum(prof[d] / mean for prof, mean in zip(profiles, means)) / len(weights) for d in range(3)] + return shapes, means + + +def _dense_mass_1d( + mass_operator: WeightedMassOperator, basis: str, d: int, weight: Callable | xp.ndarray +) -> tuple[np.ndarray, DomainDecomposition]: + """Global 1d mass matrix in direction d (B- or M-splines) as dense host array, see ``_mass_matrix_1d``.""" + M, domain_decomposition = _mass_matrix_1d(mass_operator, basis, d, weight) + return np.asarray(xp.to_numpy(M.toarray()), dtype=float), domain_decomposition + + +def _difference_matrix_1d(n_B: int, n_M: int) -> np.ndarray: + """ + 1d derivative matrix from B-spline to M-spline coefficients, :math:`(\\mathbb d c)_r = c_{r+1} - c_r` + (periodic if ``n_M == n_B``, else ``n_M == n_B - 1``). + """ + assert n_M in (n_B, n_B - 1) + D = np.zeros((n_M, n_B)) + for r in range(n_M): + D[r, r] = -1.0 + D[r, (r + 1) % n_B] = 1.0 + return D + + +def _apply_bc_dense(A: np.ndarray, bc_d: tuple[bool, bool]) -> np.ndarray: + """Impose essential boundary conditions on a dense 1d matrix: zero row and column, 1 on the diagonal.""" + A = A.copy() + for i, is_essential in zip((0, -1), bc_d): + if is_essential: + A[i, :] = 0.0 + A[:, i] = 0.0 + A[i, i] = 1.0 + return A + + +def _kronecker_sum( + space: StencilVectorSpace, + stiffness: list[StencilMatrix | None], + mass: list[StencilMatrix], + sigma: float, +) -> LinearOperator: + """ + The operator :math:`\\sum_d M_1 \\otimes \\dots \\otimes S_d \\otimes \\dots \\otimes M_3 + \\sigma M_1 \\otimes M_2 \\otimes M_3` + as sum of KroneckerStencilMatrix (process-local 1d factors; directions with ``S_d = None`` have no term). + """ + terms = [] + for d, S_d in enumerate(stiffness): + if S_d is not None: + factors = [S_d.copy() if e == d else M_e.copy() for e, M_e in enumerate(mass)] + terms.append(KroneckerStencilMatrix(space, space, *factors)) + if sigma != 0.0: + terms.append(KroneckerStencilMatrix(space, space, *[M_e.copy() for M_e in mass]) * sigma) + assert terms, "The operator has no terms." + return reduce(lambda a, b: a + b, terms) + + +def _probe_diagonal(op: LinearOperator, space: VectorSpace) -> list[xp.ndarray]: + """ + Local diagonal of ``op`` (one array per block of ``space``) by probing with colored unit vectors. + + In each direction, the degrees of freedom are colored with a stride larger than the coupling + width of ``op`` (assumed at most ``pads + 1``); applying ``op`` to the sum of the unit vectors + of one color gives the diagonal entries of that color. In periodic directions, the stride + divides the number of points (so that no two dofs of one color couple across the boundary). + This needs ``prod(stride)`` applications of ``op`` per block; it is collective. + """ + is_block = isinstance(space, BlockVectorSpace) + spaces = space.spaces if is_block else (space,) + x = space.zeros() + y = space.zeros() + xs = x.blocks if is_block else (x,) + ys = y.blocks if is_block else (y,) + + diags = [] + for c, V in enumerate(spaces): + local = tuple(slice(p * m, p * m + e - s + 1) for p, m, s, e in zip(V.pads, V.shifts, V.starts, V.ends)) + shape = tuple(e - s + 1 for s, e in zip(V.starts, V.ends)) + + strides = [] + for n, p, periodic in zip(V.npts, V.pads, V.periods): + st = min(p + 2, n) + while periodic and n % st: + st += 1 + strides.append(st) + indices = [xp.arange(s, e + 1) for s, e in zip(V.starts, V.ends)] + + diag = xp.zeros(shape, dtype=float) + for color in itertools.product(*(range(st) for st in strides)): + mask = xp.ones(shape, dtype=bool) + for d, (g, st, r) in enumerate(zip(indices, strides, color)): + mask = mask & (g % st == r).reshape([-1 if e == d else 1 for e in range(len(shape))]) + for xb in xs: + xb._data[:] = 0.0 + xs[c]._data[local] = mask + x.ghost_regions_in_sync = False + op.dot(x, out=y) + diag[mask] = ys[c]._data[local][mask] + diags.append(diag) + return diags class FFTSolver(BandedSolver): @@ -899,8 +1254,19 @@ def __init__(self, circmat): assert isinstance(circmat, xp.ndarray) assert is_circulant(circmat) + from scipy.linalg import solve_circulant # deferred: scipy.linalg is slow to import + self._space = xp.ndarray - self._column = circmat[:, 0] + # copy: circmat is still used by the caller (e.g. for the process-local stencil matrix) + self._column = xp.array(circmat[:, 0], copy=True) + + # stabilize a singular matrix once, here, so that all solves use the same matrix + try: + solve_circulant(self._column, xp.ones_like(self._column)) + except xp.linalg.LinAlgError: + eps = 1e-4 + logger.info(f"Stabilizing singular preconditioning FFTSolver with {eps =}:") + self._column[0] *= 1.0 + eps # -------------------------------------- # Abstract interface @@ -934,20 +1300,11 @@ def solve(self, rhs, out=None, transposed=False): assert rhs.T.shape[0] == self._column.size if out is None: - out = solve_circulant(self._column, rhs.T).T - - else: - assert out.shape == rhs.shape - assert out.dtype == rhs.dtype - - try: - out[:] = solve_circulant(self._column, rhs.T).T - except xp.linalg.LinAlgError: - eps = 1e-4 - logger.info(f"Stabilizing singular preconditioning FFTSolver with {eps =}:") - self._column[0] *= 1.0 + eps - out[:] = solve_circulant(self._column, rhs.T).T + return solve_circulant(self._column, rhs.T).T + assert out.shape == rhs.shape + assert out.dtype == rhs.dtype + out[:] = solve_circulant(self._column, rhs.T).T return out @@ -970,12 +1327,7 @@ def is_circulant(mat): assert len(mat.shape) == 2 assert mat.shape[0] == mat.shape[1] - if mat.shape[0] > 1: - for i in range(mat.shape[0] - 1): - circulant = xp.allclose(mat[i, :], xp.roll(mat[i + 1, :], -1)) - if not circulant: - return circulant - else: - circulant = True - - return circulant + # circulant: every row is the previous one shifted by one, i.e. mat[i, j] = mat[0, (j - i) % n] + n = mat.shape[0] + idx = (xp.arange(n)[None, :] - xp.arange(n)[:, None]) % n + return bool(xp.allclose(mat, mat[0][idx])) diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index e1a1beeeb..7a85e2b9d 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -2129,7 +2129,7 @@ def __init__(self, projector, transposed=False, apply_bc=False): # temporary vectors for dot product if self._is_composed: tmp_vectors = [] - for op in self._I.multiplicants[1:]: + for op in self._I.multiplicands[1:]: tmp_vectors.append(op.codomain.zeros()) self._tmp_vectors = tuple(tmp_vectors) @@ -2197,7 +2197,7 @@ def solve(self, rhs, out=None): x = rhs for i in range(len(self._tmp_vectors)): y = self._tmp_vectors[-1 - i] - A = self._I.multiplicants[-1 - i] + A = self._I.multiplicands[-1 - i] if isinstance(A, (StencilMatrix, KroneckerStencilMatrix, BlockLinearOperator)): self.solver.dot(x, out=y) else: @@ -2205,7 +2205,7 @@ def solve(self, rhs, out=None): x = y # last operator - A = self._I.multiplicants[0] + A = self._I.multiplicands[0] if out is None: out = A.dot(x) else: diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index 3f304c3f7..23af76b82 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -1595,7 +1595,8 @@ def test_average_operator_subcomm(): @pytest.mark.parametrize("dim_reduce", [0, 1, 2]) -def test_mass_preconditioner_array_weights_mpi(dim_reduce): +@pytest.mark.parametrize("weight_reduction", ["midpoint", "average"]) +def test_mass_preconditioner_array_weights_mpi(dim_reduce, weight_reduction): """Preconditioner with array weights must not depend on the MPI decomposition (num_elements not divisible by the number of processes).""" @@ -1630,13 +1631,61 @@ def test_mass_preconditioner_array_weights_mpi(dim_reduce): M.assemble() _, v = create_equal_random_arrays(derham.V0fem, seed=1234) - out += [MassMatrixPreconditioner(M, dim_reduce=dim_reduce).dot(v)] + out += [MassMatrixPreconditioner(M, dim_reduce=dim_reduce, weight_reduction=weight_reduction).dot(v)] s, e = out[0].space.starts, out[0].space.ends sl = tuple(slice(si, ei + 1) for si, ei in zip(s, e)) assert xp.allclose(out[0][sl], out[1][sl], rtol=1e-12, atol=1e-14) +@pytest.mark.parametrize("d", [0, 1, 2]) +def test_reduced_weight_1d(d): + """1d weights of the MassMatrixPreconditioner (mid point and mean of the 3d weight) + for callable and array weights, against analytic values.""" + + import cunumpy as xp + + from struphy.feec.preconditioner import _reduced_weight_1d + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham(TensorProductGrid(num_elements=[7, 5, 4]), DerhamOptions(degree=[2, 2, 1], bcs=(None, None, None))) + + def w(e1, e2, e3): + return 1.0 + e1 + 2.0 * e2**2 + 3.0 * e3**3 + e1 * e2 * e3 + + # analytic mid point and mean over the other two directions + midpoint = { + 0: lambda x: w(x, 0.5, 0.5), + 1: lambda x: w(0.5, x, 0.5), + 2: lambda x: w(0.5, 0.5, x), + } + average = { + 0: lambda x: 1.0 + x + 2.0 / 3.0 + 3.0 / 4.0 + x / 4.0, + 1: lambda x: 1.5 + 2.0 * x**2 + 3.0 / 4.0 + x / 4.0, + 2: lambda x: 1.5 + 2.0 / 3.0 + 3.0 * x**3 + x / 4.0, + } + + pts = [xp.ravel(p) for p in derham.spline_attributes["H1"].quad_grid_pts[0]] + arr = w(*xp.meshgrid(*pts, indexing="ij")) + x = pts[d] + + for reduction, exact in (("average", average[d]), ("midpoint", None)): + from_callable = _reduced_weight_1d(w, 0, d, derham, reduction)(x) + from_array = _reduced_weight_1d(arr, 0, d, derham, reduction) + if reduction == "average": + assert xp.allclose(from_callable, exact(x), rtol=1e-13, atol=1e-13) + assert xp.allclose(from_array, exact(x), rtol=1e-13, atol=1e-13) + else: + assert xp.allclose(from_callable, midpoint[d](x), rtol=1e-13, atol=1e-13) + # the array weight is only known at the quadrature points: it is taken at the + # global quadrature point with index (num_elements * nquads) // 2 + mid = [(n * q) // 2 for n, q in zip(derham.num_elements, derham.nquads)] + cut = tuple(slice(None) if k == d else mid[k] for k in range(3)) + assert xp.allclose(from_array, arr[cut], rtol=1e-13, atol=1e-13) + + def test_transpose_and_copy(): """WeightedMassOperator.T and .copy() must work without a name, keep spline weights and transpose the data of accumulation-type (symm/asym) matrices (#501).""" diff --git a/src/struphy/feec/tests/test_preconditioner_transpose.py b/src/struphy/feec/tests/test_preconditioner_transpose.py index e5e74d0a6..f0b49eaf6 100644 --- a/src/struphy/feec/tests/test_preconditioner_transpose.py +++ b/src/struphy/feec/tests/test_preconditioner_transpose.py @@ -2,7 +2,9 @@ @pytest.mark.parametrize("dim_reduce", [0, 2]) -def test_mass_preconditioner_transpose(dim_reduce): +@pytest.mark.parametrize("weight_reduction", ["midpoint", "average"]) +@pytest.mark.parametrize("diagonal_scaling", [False, True]) +def test_mass_preconditioner_transpose(dim_reduce, weight_reduction, diagonal_scaling): """The transposes of the mass-matrix preconditioners keep their class and options.""" from struphy import domains @@ -16,23 +18,32 @@ def test_mass_preconditioner_transpose(dim_reduce): derham = Derham(grid, DerhamOptions(degree=(1, 1, 1), bcs=(None, None, None))) mass_ops = WeightedMassOperators(derham, domains.Colella()) - pc = MassMatrixPreconditioner(mass_ops.M1, apply_bc=False, dim_reduce=dim_reduce) + pc = MassMatrixPreconditioner( + mass_ops.M1, + apply_bc=False, + dim_reduce=dim_reduce, + weight_reduction=weight_reduction, + diagonal_scaling=diagonal_scaling, + ) pc_T = pc.transpose() assert type(pc_T) is MassMatrixPreconditioner assert pc_T._apply_bc is False assert pc_T._dim_reduce == dim_reduce + assert pc_T.weight_reduction == weight_reduction + assert pc_T.diagonal_scaling == diagonal_scaling pc_diag_T = MassMatrixDiagonalPreconditioner(mass_ops.M1, apply_bc=False).transpose() assert type(pc_diag_T) is MassMatrixDiagonalPreconditioner assert pc_diag_T._apply_bc is False -def test_mass_diagonal_preconditioner_assembles_diagonal_blocks_only(monkeypatch): - """The logical mass matrix of MassMatrixDiagonalPreconditioner needs only the diagonal blocks.""" +def test_mass_diagonal_preconditioner_assembles_1d_matrices_only(monkeypatch): + """MassMatrixDiagonalPreconditioner assembles only 1d mass matrices (no 3d logical mass matrix); + the diagonal of the logical mass matrix is the one of its Kronecker approximation.""" from struphy import domains from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators - from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner + from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, MassMatrixPreconditioner from struphy.feec.psydac_derham import Derham from struphy.io.options import DerhamOptions from struphy.topology.grids import TensorProductGrid @@ -42,18 +53,78 @@ def test_mass_diagonal_preconditioner_assembles_diagonal_blocks_only(monkeypatch mass_ops = WeightedMassOperators(derham, domains.Colella()) M1 = mass_ops.M1 - weights_infos = [] + created = [] init = WeightedMassOperator.__init__ - def recording_init(self, *args, weights_info=None, **kwargs): - weights_infos.append(weights_info) - init(self, *args, weights_info=weights_info, **kwargs) + def recording_init(self, derham, V, W, *args, **kwargs): + created.append(V.ldim) + init(self, derham, V, W, *args, **kwargs) monkeypatch.setattr(WeightedMassOperator, "__init__", recording_init) - MassMatrixDiagonalPreconditioner(M1) + pc = MassMatrixDiagonalPreconditioner(M1) + + assert isinstance(pc, MassMatrixPreconditioner) + assert pc.dim_reduce is None and pc.diagonal_scaling + assert created and all(ldim == 1 for ldim in created) + + +def test_mass_diagonal_preconditioner_zero_diagonal_block(): + """A mass operator with a zero (None) diagonal block, e.g. after reassembly with a zero weight: + construction and update_mass_operator must work; the scaling of that block is 1.""" + + import cunumpy as xp + + from struphy import domains + from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators + from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + grid = TensorProductGrid(num_elements=(3, 3, 2)) + derham = Derham(grid, DerhamOptions(degree=(1, 1, 1), bcs=(None, None, None))) + mass_ops = WeightedMassOperators(derham, domains.Colella()) + + def one(e1, e2, e3): + return xp.ones_like(e1, dtype=float) + + weights = [[one if i == j and i != 2 else None for j in range(3)] for i in range(3)] + M = WeightedMassOperator(derham, derham.V1fem, derham.V1fem, weights_info=weights) + M.assemble() + assert M.matrix[2, 2] is None + + pc = MassMatrixDiagonalPreconditioner(mass_ops.M1) + pc.update_mass_operator(M) + scaling = pc._scaling[0] + assert xp.allclose(scaling[2, 2]._data, 1.0) + + _, x = create_equal_random_arrays(derham.V1fem, seed=1) + assert all(xp.all(xp.isfinite(b.toarray())) for b in pc.dot(x).blocks) + + +def test_mass_diagonal_preconditioner_small_periodic_direction(): + """Periodic direction with fewer points than the stencil width (1 element, degree 1): the main + diagonal of the 1d factors must be at offset 0, otherwise the diagonal scaling vanishes.""" + + import cunumpy as xp + from feectools.linalg.solvers import inverse + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, _local_diagonal + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham(grid=TensorProductGrid(num_elements=(4, 4, 1)), options=DerhamOptions(degree=(2, 2, 1))) + mass_ops = WeightedMassOperators(derham=derham, domain=domains.Cuboid()) + + pc = MassMatrixDiagonalPreconditioner(mass_ops.Mv) + assert all(float(xp.min(d)) > 0.0 for d in _local_diagonal(pc.matrix)) - # the logical mass matrix is the only 3x3 weighted mass operator created - (fun,) = [w for w in weights_infos if isinstance(w, list) and len(w) == 3] - for i in range(3): - for j in range(3): - assert callable(fun[i][j]) == (i == j) + _, x = create_equal_random_arrays(derham.Vvfem, seed=1) + inv = inverse(mass_ops.Mv, "pcg", pc=pc, tol=1e-12, maxiter=100) + inv.dot(x) + assert inv._info["success"] diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py new file mode 100644 index 000000000..6305551ca --- /dev/null +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -0,0 +1,272 @@ +import pytest + +FORMS = {"grad": "0", "curl": "1", "div": "2"} + + +def _blocks(v): + from feectools.linalg.block import BlockVector + + return v.blocks if isinstance(v, BlockVector) else (v,) + + +def _max_abs(arrays, comm): + import cunumpy as xp + from maybempi import MPI + + m = max(float(xp.max(xp.abs(a))) for a in arrays) + return comm.allreduce(m, op=MPI.MAX) if comm is not None and comm.Get_size() > 1 else m + + +@pytest.mark.parametrize("derivative", ["grad", "curl", "div"]) +@pytest.mark.parametrize("sigma", [0.0, 0.7]) +def test_stiffness_approximation_unit_cube(derivative, sigma): + """On the unit cube (unit weights), the diagonal blocks of the Kronecker approximation equal those of + the stiffness operator, the solver inverts the approximation, and probing gives its exact diagonal.""" + + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import StiffnessPreconditioner, _local_diagonal, _probe_diagonal + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + comm = MPI.COMM_WORLD + derham = Derham( + TensorProductGrid(num_elements=[6, 5, 4]), DerhamOptions(degree=[2, 3, 2], bcs=(None, None, None)), comm=comm + ) + mass_ops = WeightedMassOperators(derham, domains.Cuboid()) + + P = StiffnessPreconditioner( + mass_ops, + derivative, + sigma=sigma, + apply_bc=False, + diagonal_scaling=False, + kernel_correction=False, + weights="unit", + ) + A = P.core_operator + x = create_equal_random_arrays(derham.fem_spaces[FORMS[derivative]], seed=1)[1] + + # diagonal blocks: apply to one component at a time + for c in range(len(_blocks(x))): + xc = x.copy() + for k, b in enumerate(_blocks(xc)): + if k != c: + b._data[:] = 0.0 + r = A.dot(xc) - P.matrix.dot(xc) + ref = A.dot(xc) + assert _max_abs([_blocks(r)[c].toarray()], comm) < 1e-12 * _max_abs([_blocks(ref)[c].toarray()], comm) + + if sigma > 0: + r = P.solver.dot(P.matrix.dot(x)) - x + assert _max_abs([b.toarray() for b in _blocks(r)], comm) < 1e-12 + + # probing the stiffness operator gives its diagonal, which equals the one of the approximation + diffs = [a - b for a, b in zip(_probe_diagonal(A, P.matrix.domain), _local_diagonal(P.matrix))] + assert _max_abs(diffs, comm) < 1e-12 * _max_abs(_local_diagonal(P.matrix), comm) + + +@pytest.mark.parametrize("mapping", ["Cuboid", "HollowCylinder"]) +@pytest.mark.parametrize("sigma", [0.0, 1.0]) +def test_stiffness_preconditioner_grad(mapping, sigma): + """PCG for G^T M1 G + sigma M0: exact on the unit cube, fewer iterations than without preconditioner otherwise.""" + + from feectools.linalg.solvers import inverse + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import StiffnessPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham( + TensorProductGrid(num_elements=[8, 10, 4]), + DerhamOptions(degree=[2, 3, 1], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=MPI.COMM_WORLD, + ) + domain = domains.Cuboid() if mapping == "Cuboid" else domains.HollowCylinder(a1=0.1, a2=1.0, Lz=3.0) + mass_ops = WeightedMassOperators(derham, domain) + + A = derham.grad.T @ mass_ops.M1 @ derham.grad + if sigma: + A = A + sigma * mass_ops.M0 + b = derham.boundary_ops["0"].dot(create_equal_random_arrays(derham.V0fem, seed=2, flattened=True)[1]) + + niter = {} + for label, P in ( + ("none", None), + ("unit", StiffnessPreconditioner(mass_ops, "grad", sigma=sigma, weights="unit", diagonal_scaling=False)), + ("scaled", StiffnessPreconditioner(mass_ops, "grad", sigma=sigma, weights="unit")), + ): + inv = inverse(A, "pcg", pc=P, tol=1e-10, maxiter=2000) if P else inverse(A, "cg", tol=1e-10, maxiter=2000) + inv.dot(b) + assert inv._info["success"] + niter[label] = inv._info["niter"] + + if mapping == "Cuboid": + assert niter["unit"] <= 2 + assert niter["scaled"] <= 2 + else: + assert niter["unit"] < niter["none"] + assert niter["scaled"] <= niter["unit"] + + +@pytest.mark.parametrize("derivative", ["curl", "div"]) +@pytest.mark.parametrize("mapping", ["Cuboid", "HollowCylinder"]) +def test_stiffness_preconditioner_kernel_correction(derivative, mapping): + """PCG for C^T M2 C + sigma M1 and D^T M3 D + sigma M2: with kernel correction, fewer iterations + than with block Jacobi alone and than with the mass-matrix preconditioner.""" + + from feectools.linalg.solvers import inverse + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import MassMatrixPreconditioner, StiffnessPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham( + TensorProductGrid(num_elements=[8, 10, 4]), + DerhamOptions(degree=[2, 3, 1], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=MPI.COMM_WORLD, + ) + domain = domains.Cuboid() if mapping == "Cuboid" else domains.HollowCylinder(a1=0.1, a2=1.0, Lz=3.0) + mass_ops = WeightedMassOperators(derham, domain) + + k = int(FORMS[derivative]) + d = getattr(derham, derivative) + M_k, M_k1 = getattr(mass_ops, f"M{k}"), getattr(mass_ops, f"M{k + 1}") + sigma = 1.0 + A = d.T @ M_k1 @ d + sigma * M_k + b = derham.boundary_ops[FORMS[derivative]].dot( + create_equal_random_arrays(derham.fem_spaces[FORMS[derivative]], seed=2, flattened=True)[1] + ) + + P = StiffnessPreconditioner(mass_ops, derivative, sigma=sigma) + assert P.kernel_correction is not None + niter = {} + for label, pc in ( + ("mass", MassMatrixPreconditioner(M_k, diagonal_scaling=True)), + ("block Jacobi", StiffnessPreconditioner(mass_ops, derivative, sigma=sigma, kernel_correction=False)), + ("kernel correction", P), + ): + inv = inverse(A, "pcg", pc=pc, tol=1e-10, maxiter=3000) + inv.dot(b) + assert inv._info["success"] + niter[label] = inv._info["niter"] + + assert niter["kernel correction"] < niter["block Jacobi"] + assert niter["kernel correction"] < niter["mass"] + + with pytest.raises(ValueError): + StiffnessPreconditioner(mass_ops, derivative, sigma=0.0) + + +@pytest.mark.parametrize("derivative", ["grad", "curl", "div"]) +def test_stiffness_preconditioner_average_weights(derivative): + """On a stretched cuboid (constant, anisotropic weights), weights="average" is exact for grad + and better than unit weights for curl and div.""" + + from feectools.linalg.solvers import inverse + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import StiffnessPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham( + TensorProductGrid(num_elements=[8, 10, 4]), + DerhamOptions(degree=[2, 3, 1], bcs=(("dirichlet", "dirichlet"), None, None)), + comm=MPI.COMM_WORLD, + ) + mass_ops = WeightedMassOperators(derham, domains.Cuboid(l1=0.0, r1=2.0, l2=0.0, r2=0.5, l3=0.0, r3=1.0)) + + k = int(FORMS[derivative]) + d = getattr(derham, derivative) + sigma = 0.0 if derivative == "grad" else 1.0 + A = d.T @ getattr(mass_ops, f"M{k + 1}") @ d + if sigma: + A = A + sigma * getattr(mass_ops, f"M{k}") + b = derham.boundary_ops[FORMS[derivative]].dot( + create_equal_random_arrays(derham.fem_spaces[FORMS[derivative]], seed=2, flattened=True)[1] + ) + + niter = {} + for weights in ("unit", "average"): + P = StiffnessPreconditioner(mass_ops, derivative, sigma=sigma, weights=weights) + assert P.weights == weights + inv = inverse(A, "pcg", pc=P, tol=1e-10, maxiter=3000) + inv.dot(b) + assert inv._info["success"] + niter[weights] = inv._info["niter"] + + if derivative == "grad": + assert niter["average"] <= 2 + assert niter["average"] < niter["unit"] + + +@pytest.mark.parametrize("derivative", ["grad", "curl", "div"]) +@pytest.mark.parametrize("weights", ["unit", "average"]) +def test_stiffness_preconditioner_mpi(derivative, weights): + """The preconditioner must not depend on the MPI decomposition.""" + + import cunumpy as xp + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import StiffnessPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.feec.utilities import create_equal_random_arrays + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + out = [] + for comm in (MPI.COMM_WORLD, None): + derham = Derham( + TensorProductGrid(num_elements=[7, 5, 4]), + DerhamOptions(degree=[2, 2, 1], bcs=(("dirichlet", "free"), None, None)), + comm=comm, + ) + mass_ops = WeightedMassOperators(derham, domains.HollowCylinder(a1=0.1, a2=1.0, Lz=3.0)) + _, v = create_equal_random_arrays(derham.fem_spaces[FORMS[derivative]], seed=1234) + out += [StiffnessPreconditioner(mass_ops, derivative, sigma=1.0, weights=weights).dot(v)] + + for a, b in zip(_blocks(out[0]), _blocks(out[1])): + sl = tuple(slice(si, ei + 1) for si, ei in zip(a.space.starts, a.space.ends)) + assert xp.allclose(a[sl], b[sl], rtol=1e-12, atol=1e-14) + + +def test_stiffness_preconditioner_options(): + """Option checks, transpose and polar splines.""" + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import StiffnessPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + derham = Derham(TensorProductGrid(num_elements=[4, 4, 2]), DerhamOptions(degree=[1, 1, 1], bcs=(None, None, None))) + mass_ops = WeightedMassOperators(derham, domains.Cuboid()) + P = StiffnessPreconditioner(mass_ops, "grad", sigma=0.5) + assert P.transpose() is P + assert P.derivative == "grad" and P.sigma == 0.5 and P.diagonal_scaling and P.weights == "average" + with pytest.raises(AssertionError): + StiffnessPreconditioner(mass_ops, "rot") + with pytest.raises(AssertionError): + StiffnessPreconditioner(mass_ops, "grad", sigma=-1.0) diff --git a/src/struphy/linear_algebra/multigrid/coarsen.py b/src/struphy/linear_algebra/multigrid/coarsen.py index b33bce84e..fc9dd4084 100644 --- a/src/struphy/linear_algebra/multigrid/coarsen.py +++ b/src/struphy/linear_algebra/multigrid/coarsen.py @@ -111,7 +111,7 @@ def __call__(self, A: LinearOperator) -> LinearOperator: return SumLinearOperator(self.space(A.domain), self.space(A.codomain), *addends) if isinstance(A, ComposedLinearOperator): - factors = [self(a) for a in A.multiplicants] + factors = [self(a) for a in A.multiplicands] return ComposedLinearOperator(self.space(A.domain), self.space(A.codomain), *factors) if isinstance(A, PowerLinearOperator): diff --git a/src/struphy/propagators/variational_resistivity.py b/src/struphy/propagators/variational_resistivity.py index e0a4a06fc..758a654be 100644 --- a/src/struphy/propagators/variational_resistivity.py +++ b/src/struphy/propagators/variational_resistivity.py @@ -249,7 +249,7 @@ def __call_newton(self, dt): total_resistivity = self._update_artificial_resistivity(bn, dt) self._scaled_stiffness._scalar = dt * self._eta - # self.evol_op._multiplicants[1]._addends[0]._scalar = -dt*self._eta/2. + # self.evol_op._multiplicands[1]._addends[0]._scalar = -dt*self._eta/2. if self._linearize_current: bn1 = self.evol_op.dot( bn diff --git a/src/struphy/propagators/variational_viscosity.py b/src/struphy/propagators/variational_viscosity.py index 0c8d496b6..a9b8f11c7 100644 --- a/src/struphy/propagators/variational_viscosity.py +++ b/src/struphy/propagators/variational_viscosity.py @@ -272,7 +272,7 @@ def __call_newton(self, dt): self._scaled_stiffness._scalar = dt * self._mu # /2. self._scaled_Mv._scalar = dt * self._alpha - # self.evol_op._multiplicants[1]._addends[0]._scalar = - dt*self._mu/2. + # self.evol_op._multiplicands[1]._addends[0]._scalar = - dt*self._mu/2. un1 = self.evol_op.dot(un, out=self._tmp_un1) if self._info: logger.info(f"information on the linear solver : {self.inv_lop._info}")