From 1264e5281738c6fa0f487259ffc72c2bb0afcc7a Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 08:52:05 +0200 Subject: [PATCH 01/18] Bump feectools: KroneckerStencilMatrix matmul and grouped factors Points the feectools submodule to the kron-stencil-matmul branch (KroneckerStencilMatrix @, factors on groups of axes, factor_ndims in KroneckerLinearSolver; feectools 0.6.0). Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index f97ea5af7..6ee8f082c 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit f97ea5af7c9a21808f7d5adfc95c907c338d74de +Subproject commit 6ee8f082c5e1f9cdfd25ed726f6f99f4b5ecd97b From 500a0561924806c3cd8f126c6136612fa6679d62 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 10:16:17 +0200 Subject: [PATCH 02/18] Refactor MassMatrixPreconditioner with helper functions - Split the construction of the Kronecker approximation into module-level helpers (boundary conditions, MPI weight gathering, 1d weight, 1d mass matrix + bcs, 1d solver, process-local 1d factor, block-diagonal assembly, operator to invert); MassMatrixDiagonalPreconditioner reuses them instead of a copy of the same code. - Vectorized copy into the process-local 1d stencil matrix; block-diagonal operators no longer hard-coded to 3 components. - Docstrings, comments and type annotations. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 965 +++++++++++++---------------- 1 file changed, 429 insertions(+), 536 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index ac413ba59..1d96f3f5c 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -1,11 +1,15 @@ +from __future__ import annotations + import logging +from collections.abc import Callable +from dataclasses import dataclass import cunumpy as xp from feectools.api.essential_bc import apply_essential_bc_stencil from feectools.ddm.cart import CartDecomposition, DomainDecomposition 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, 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 @@ -20,28 +24,37 @@ class MassMatrixPreconditioner(LinearOperator): - """ + r""" Preconditioner for inverting 3d weighted mass matrices. 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. + (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: + the diagonal block ``(c, c)`` of the 3d weight, evaluated at the mid point + (0.5) of the other two directions. + * 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 :meth:`solve`). Parameters ---------- - mass_operator : struphy.feec.mass.WeightedMassOperator + 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 - Along which axis to take the approximate value of the weight + Axis along which the weight is kept (in the other axes it is taken at 0.5). """ - def __init__(self, mass_operator, apply_bc=True, dim_reduce=0): + def __init__(self, mass_operator: WeightedMassOperator, apply_bc: bool = True, dim_reduce: int = 0): assert isinstance(mass_operator, WeightedMassOperator) assert mass_operator.domain == mass_operator.codomain, "Only square mass matrices can be inverted!" @@ -54,327 +67,57 @@ def __init__(self, mass_operator, apply_bc=True, dim_reduce=0): self._apply_bc = apply_bc self._dim_reduce = dim_reduce - # 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 - - n_comps = len(femspaces) n_dims = self._femspace.ldim - assert n_dims == 3 # other dims not yet implemented assert dim_reduce < n_dims - # 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 + # boundary conditions are only imposed if the mass operator has a BoundaryOperator + bc = _boundary_conditions(mass_operator, apply_bc) + apply_bc = 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, - ) + # MPI setup for gathering array-valued weights along dim_reduce + gather = _WeightGather.setup(derham, dim_reduce) - sblocks = [ - [solverblocks[0], None, None], - [None, solverblocks[1], None], - [None, None, solverblocks[2]], - ] + def weight_1d(c: int, d: int) -> Callable | xp.ndarray: + if d == dim_reduce: + return _reduced_weight_1d(mass_operator.weights[c][c], c, d, derham, gather) + return _ones_1d - self._solver = BlockLinearOperator( - self._femspace.coeff_space, - self._femspace.coeff_space, - blocks=sblocks, - ) + # Kronecker approximation of the mass matrix and its exact inverse + self._matrix, self._solver = _kronecker_approximation(mass_operator, bc, weight_1d) - # save mass operator to be inverted (needed in solve method) - if apply_bc: - self._M = mass_operator.M0 - else: - self._M = mass_operator.M - - self._is_composed = isinstance(self._M, ComposedLinearOperator) - - # temporary vectors for dot product + # mass operator to be inverted (with boundary operators B, E if apply_bc), needed in solve + self._M, self._is_composed, tmp = _operator_to_invert(mass_operator, apply_bc) if self._is_composed: - tmp_vectors = [] - for op in self._M.multiplicants[1:]: - tmp_vectors.append(op.codomain.zeros()) - - self._tmp_vectors = tuple(tmp_vectors) + self._tmp_vectors = tmp else: - self._tmp_vector = self._M.codomain.zeros() + self._tmp_vector = tmp @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) -> KroneckerStencilMatrix | BlockLinearOperator: + """Approximation of the input mass matrix as KroneckerStencilMatrix (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) -> KroneckerLinearSolver | BlockLinearOperator: + """KroneckerLinearSolver (block diagonal for vector-valued spaces) for exactly inverting the approximate mass matrix 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,17 +125,20 @@ def codomain(self): def dtype(self): return self._dtype - def transpose(self, conjugate=False): + def transpose(self, conjugate: bool = False) -> "MassMatrixPreconditioner": """ Returns the transposed operator. """ return MassMatrixPreconditioner(self._mass_operator.transpose(), self._apply_bc, self._dim_reduce) @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. + The operators of the composition are applied from right to left; the mass + matrix itself is replaced by the Kronecker solver. + Parameters ---------- rhs : feectools.linalg.basic.Vector @@ -410,35 +156,34 @@ def solve(self, rhs, out=None): assert isinstance(rhs, Vector) assert rhs.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 not self._is_composed: if out is None: - out = A.dot(x) + out = self._tmp_vector.copy() + self.solver.dot(rhs, out=out) + return out + + # successive dot products with all but the first (left-most) operator + x = rhs + for A, y in zip(reversed(self._M.multiplicants[1:]), reversed(self._tmp_vectors)): + if isinstance(A, (StencilMatrix, BlockLinearOperator)): + # the mass matrix: apply the approximate inverse instead + self.solver.dot(x, out=y) else: - assert isinstance(out, Vector) - assert out.space == self._space - A.dot(x, out=out) + A.dot(x, out=y) + x = y + # first operator + A = self._M.multiplicants[0] + if out is None: + out = A.dot(x) else: - if out is None: - out = self._tmp_vector.copy() - self.solver.dot(rhs, out=out) + assert isinstance(out, Vector) + assert out.space == self._space + A.dot(x, out=out) return 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) @@ -493,223 +238,22 @@ def __init__(self, mass_operator, apply_bc=True): self._domain = mass_operator.domain self._apply_bc = 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 - - n_comps = len(femspaces) n_dims = self._femspace.ldim - assert n_dims == 3 # other dims not yet implemented - # 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], - ) - - 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]], - ] - self._solver = BlockLinearOperator( - self._femspace.coeff_space, - self._femspace.coeff_space, - blocks=sblocks, - ) - - # save mass operator to be inverted (needed in solve method) - if apply_bc: - self._M = mass_operator.M0 - else: - self._M = mass_operator.M + # boundary conditions are only imposed if the mass operator has a BoundaryOperator + bc = _boundary_conditions(mass_operator, apply_bc) + apply_bc = bc is not None - self._is_composed = isinstance(self._M, ComposedLinearOperator) + # mass matrix on the logical domain (unit weights) as Kronecker product, and its exact inverse + self._matrix, self._solver = _kronecker_approximation(mass_operator, bc, lambda c, d: _ones_1d) - # temporary vectors for dot product + # mass operator to be inverted (with boundary operators B, E if apply_bc), needed in solve + self._M, self._is_composed, tmp = _operator_to_invert(mass_operator, apply_bc) if self._is_composed: - tmp_vectors = [] - for op in self._M.multiplicants[1:]: - tmp_vectors.append(op.codomain.zeros()) - - self._tmp_vectors = tuple(tmp_vectors) + self._tmp_vectors = tmp else: - self._tmp_vector = self._M.codomain.zeros() + self._tmp_vector = tmp # Need to assemble the logical mass matrix to extract the coefficients fun = [ @@ -884,6 +428,355 @@ def dot(self, v, out=None): return out +# -------------------------------------------------------------------------------------- +# 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) + + +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``. + + 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.multiplicants[-1] + if not isinstance(last, BoundaryOperator): + return None + return last.bc + + +def _operator_to_invert( + mass_operator: WeightedMassOperator, apply_bc: bool +) -> tuple[LinearOperator, bool, tuple[Vector, ...] | Vector]: + """ + The operator approximately inverted by the preconditioners: ``M0`` (with boundary + operators) if ``apply_bc``, else ``M``. + + Returns + ------- + M : LinearOperator + The operator. + + is_composed : bool + Whether ``M`` is a ComposedLinearOperator. + + tmp : tuple[Vector, ...] | Vector + Temporary vectors for ``solve``: the codomains of all but the first factor of + ``M`` if it is composed, else one vector in the codomain of ``M``. + """ + M = mass_operator.M0 if apply_bc else mass_operator.M + is_composed = isinstance(M, ComposedLinearOperator) + if is_composed: + tmp = tuple(op.codomain.zeros() for op in M.multiplicants[1:]) + else: + tmp = M.codomain.zeros() + return M, is_composed, tmp + + +@dataclass +class _WeightGather: + """ + MPI setup for gathering a 1d cut of an array-valued weight along ``dim_reduce``. + + In the directions other than ``dim_reduce``, the weight is taken at the global mid + point; the ranks owning the mid element there (exactly one rank per slab along + ``dim_reduce``) form ``subcomm``. ``root`` is the lowest of these ranks in ``comm``. + """ + + comm: object + subcomm: object + root: int + + @classmethod + def setup(cls, derham, dim_reduce: int) -> "_WeightGather | None": + """Collective on ``derham.comm``; returns None in serial runs.""" + comm = derham.comm + if isinstance(comm, (SerialComm, type(None))): + return None + + dom_dec = derham.domain_decomposition + rank = comm.Get_rank() + is_selected = all( + dom_dec.starts[i] <= derham.num_elements[i] // 2 <= dom_dec.ends[i] + for i in range(len(derham.num_elements)) + 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}") + return cls(comm, subcomm, root) + + +def _reduced_weight_1d(weight, c: int, d: int, derham, gather: _WeightGather | None) -> Callable | xp.ndarray: + """ + 1d weight along direction ``d``: the 3d weight at the mid point (0.5) of the other directions. + + 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 (only used for logging and for the callable case). + + derham : Derham + Discrete de Rham sequence of the mass operator. + + gather : _WeightGather | None + MPI setup from ``_WeightGather.setup``; None in serial runs. + + Returns + ------- + Callable | xp.ndarray + A function of the 1d quadrature points, or the weight at all global 1d + quadrature points (collective on ``gather.comm`` in parallel runs). + """ + n_dims = 3 + + if weight is None: + return _ones_1d + + if callable(weight): + + def fun(e): + # evaluate the 3d weight on the "meshgrid" (0.5, ..., e, ..., 0.5) + 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( + weight( + *[xp.array(xp.full_like(f, 0.5)) if i != d else xp.array(f) for i in range(n_dims)], + ).squeeze(), + ) + + return fun + + if isinstance(weight, xp.ndarray): + s = weight.shape + logger.debug(f"{weight.shape = } for component {c} and direction {d}.") + dom_dec = derham.domain_decomposition + 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) + cut = tuple(slice(None) if i == d else mid[i] for i in range(n_dims)) + local_fun = xp.ascontiguousarray(weight[cut], dtype=float) + + logger.debug(f"{fun.size = } for component {c} and direction {d} before gathering on all processes.") + if gather is not None: + # local sizes differ if num_elements[d] is not divisible by the number of processes + if gather.subcomm != MPI.COMM_NULL: + counts = gather.subcomm.allgather(local_fun.size) + displs = [sum(counts[:j]) for j in range(len(counts))] + gather.subcomm.Allgatherv(local_fun, [fun, counts, displs, MPI.DOUBLE]) + gather.comm.Bcast(fun, root=gather.root) + else: + fun[:] = local_fun + logger.debug(f"{fun.shape = } for component {c} and direction {d} after gathering on all processes.") + return fun + + 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, diagonal (j - i + p) mod n + rows, cols = xp.nonzero(M_arr) + on_process = (rows >= s) & (rows <= e) + rows, cols = rows[on_process], cols[on_process] + M_local._data[rows - s + p, (cols + p - rows) % M_arr.shape[1]] = 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)], + ) + + class FFTSolver(BandedSolver): """ Solve the equation Ax = b for x, assuming A is a circulant matrix. From a19f5a2ad934f7f9bc448038e1cd299f9669b3be Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 10:41:50 +0200 Subject: [PATCH 03/18] MassMatrixPreconditioner: weight_reduction option, fixes - New option weight_reduction="midpoint" (default, unchanged results) or "average" (mean of the weight over the other two directions; Gauss quadrature of the derham for array weights, 16-point Gauss-Legendre for callables). Array weights are reduced with a single Allreduce instead of a subcommunicator + Bcast. - Locate the mass matrix in the composed operator by identity and apply the approximate inverse exactly there (shared _apply_composed for both preconditioners); previously a left-most mass matrix was applied instead of inverted. - FFTSolver: copy the column (it was a view of the matrix used for the stencil factor) and stabilize a singular matrix once at construction. - Vectorized is_circulant. - Tests: weight_reduction in the transpose and MPI array-weight tests, new test of the reduced 1d weights against analytic values. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 385 +++++++++--------- src/struphy/feec/tests/test_mass_matrices.py | 53 ++- .../tests/test_preconditioner_transpose.py | 6 +- 3 files changed, 249 insertions(+), 195 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 1d96f3f5c..ec493f43c 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -2,9 +2,9 @@ import logging from collections.abc import Callable -from dataclasses import dataclass 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.fem.tensor import TensorFemSpace @@ -31,9 +31,10 @@ class MassMatrixPreconditioner(LinearOperator): (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: - the diagonal block ``(c, c)`` of the 3d weight, evaluated at the mid point - (0.5) of the other two directions. + * 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``). * 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). @@ -51,10 +52,21 @@ class MassMatrixPreconditioner(LinearOperator): Whether to include boundary operators. dim_reduce : int - Axis along which the weight is kept (in the other axes it is taken at 0.5). + Axis along which the weight is kept. + + 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. """ - def __init__(self, mass_operator: WeightedMassOperator, apply_bc: bool = True, dim_reduce: int = 0): + def __init__( + self, + mass_operator: WeightedMassOperator, + apply_bc: bool = True, + dim_reduce: int = 0, + weight_reduction: str = "midpoint", + ): assert isinstance(mass_operator, WeightedMassOperator) assert mass_operator.domain == mass_operator.codomain, "Only square mass matrices can be inverted!" @@ -66,10 +78,12 @@ def __init__(self, mass_operator: WeightedMassOperator, apply_bc: bool = True, d self._domain = mass_operator.domain self._apply_bc = apply_bc self._dim_reduce = dim_reduce + self._weight_reduction = weight_reduction n_dims = self._femspace.ldim assert n_dims == 3 # other dims not yet implemented assert dim_reduce < n_dims + assert weight_reduction in ("midpoint", "average"), f"Unknown weight_reduction {weight_reduction!r}." # boundary conditions are only imposed if the mass operator has a BoundaryOperator bc = _boundary_conditions(mass_operator, apply_bc) @@ -78,23 +92,16 @@ def __init__(self, mass_operator: WeightedMassOperator, apply_bc: bool = True, d derham = mass_operator.derham logger.debug(f"{derham.num_elements = }, {derham.bcs = }, {derham.degree = }") - # MPI setup for gathering array-valued weights along dim_reduce - gather = _WeightGather.setup(derham, dim_reduce) - def weight_1d(c: int, d: int) -> Callable | xp.ndarray: if d == dim_reduce: - return _reduced_weight_1d(mass_operator.weights[c][c], c, d, derham, gather) + return _reduced_weight_1d(mass_operator.weights[c][c], c, d, derham, weight_reduction) return _ones_1d # Kronecker approximation of the mass matrix and its exact inverse self._matrix, self._solver = _kronecker_approximation(mass_operator, bc, weight_1d) # mass operator to be inverted (with boundary operators B, E if apply_bc), needed in solve - self._M, self._is_composed, tmp = _operator_to_invert(mass_operator, apply_bc) - if self._is_composed: - self._tmp_vectors = tmp - else: - self._tmp_vector = tmp + self._M, self._mass_index, self._tmp_vectors = _operator_to_invert(mass_operator, apply_bc) @property def space(self) -> VectorSpace: @@ -125,11 +132,21 @@ def codomain(self) -> VectorSpace: def dtype(self): return self._dtype + @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 transpose(self, conjugate: bool = False) -> "MassMatrixPreconditioner": """ Returns the transposed operator. """ - return MassMatrixPreconditioner(self._mass_operator.transpose(), self._apply_bc, self._dim_reduce) + return MassMatrixPreconditioner( + self._mass_operator.transpose(), + self._apply_bc, + self._dim_reduce, + self._weight_reduction, + ) @profile def solve(self, rhs: Vector, out: Vector | None = None) -> Vector: @@ -137,7 +154,7 @@ 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. The operators of the composition are applied from right to left; the mass - matrix itself is replaced by the Kronecker solver. + matrix M itself is replaced by the Kronecker solver. Parameters ---------- @@ -155,33 +172,11 @@ def solve(self, rhs: Vector, out: Vector | None = None) -> Vector: assert isinstance(rhs, Vector) assert rhs.space == self._space - - if not self._is_composed: - if out is None: - out = self._tmp_vector.copy() - self.solver.dot(rhs, out=out) - return out - - # successive dot products with all but the first (left-most) operator - x = rhs - for A, y in zip(reversed(self._M.multiplicants[1:]), reversed(self._tmp_vectors)): - if isinstance(A, (StencilMatrix, BlockLinearOperator)): - # the mass matrix: apply the approximate inverse instead - self.solver.dot(x, out=y) - else: - A.dot(x, out=y) - x = y - - # first operator - A = self._M.multiplicants[0] - if out is None: - out = A.dot(x) - else: + if out is not None: assert isinstance(out, Vector) assert out.space == self._space - A.dot(x, out=out) - return out + return _apply_composed(self._M, self._mass_index, self.solver.dot, self._tmp_vectors, rhs, out) def dot(self, v: Vector, out: Vector | None = None) -> Vector: """Apply linear operator to Vector v. Result is written to Vector out, if provided.""" @@ -249,11 +244,7 @@ def __init__(self, mass_operator, apply_bc=True): self._matrix, self._solver = _kronecker_approximation(mass_operator, bc, lambda c, d: _ones_1d) # mass operator to be inverted (with boundary operators B, E if apply_bc), needed in solve - self._M, self._is_composed, tmp = _operator_to_invert(mass_operator, apply_bc) - if self._is_composed: - self._tmp_vectors = tmp - else: - self._tmp_vector = tmp + self._M, self._mass_index, self._tmp_vectors = _operator_to_invert(mass_operator, apply_bc) # Need to assemble the logical mass matrix to extract the coefficients fun = [ @@ -306,18 +297,13 @@ def update_mass_operator(self, mass_operator): 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" - if self._is_composed: - if self._apply_bc: - assert isinstance(mass_operator.M0, ComposedLinearOperator) - else: - assert isinstance(mass_operator.M, ComposedLinearOperator) - self._mass_operator = mass_operator - if self._apply_bc: - self._M = mass_operator.M0 - else: - self._M = mass_operator.M + # the composition has the same structure as before, so the temporary vectors can be reused + apply_bc = _boundary_conditions(mass_operator, self._apply_bc) is not None + M, mass_index, _ = _operator_to_invert(mass_operator, apply_bc) + assert mass_index == self._mass_index, "The updated mass operator must have the same structure." + self._M = M self._M_invsrqt_diag = self._mass_operator.matrix.diagonal(inverse=True, sqrt=True, out=self._M_invsrqt_diag) def transpose(self, conjugate=False): @@ -380,34 +366,11 @@ def solve(self, rhs, out=None): 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._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) - - else: - if out is None: - out = self._tmp_vector.copy() - self._solve_no_bc(rhs, out=out) - - return out + return _apply_composed(self._M, self._mass_index, self._solve_no_bc, self._tmp_vectors, rhs, out) def dot(self, v, out=None): """Apply linear operator to Vector v. Result is written to Vector out, if provided.""" @@ -457,70 +420,80 @@ def _boundary_conditions(mass_operator: WeightedMassOperator, apply_bc: bool) -> def _operator_to_invert( mass_operator: WeightedMassOperator, apply_bc: bool -) -> tuple[LinearOperator, bool, tuple[Vector, ...] | Vector]: +) -> tuple[ComposedLinearOperator | LinearOperator, int | None, tuple[Vector, ...]]: """ The operator approximately inverted by the preconditioners: ``M0`` (with boundary operators) if ``apply_bc``, else ``M``. Returns ------- - M : LinearOperator + M : ComposedLinearOperator | LinearOperator The operator. - is_composed : bool - Whether ``M`` is a ComposedLinearOperator. + mass_index : int | None + Position of the mass matrix ``mass_operator.matrix`` in ``M.multiplicants``, or + None if ``M`` is not composed (then ``M`` is the mass matrix itself). - tmp : tuple[Vector, ...] | Vector - Temporary vectors for ``solve``: the codomains of all but the first factor of - ``M`` if it is composed, else one vector in the codomain of ``M``. + 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 - is_composed = isinstance(M, ComposedLinearOperator) - if is_composed: - tmp = tuple(op.codomain.zeros() for op in M.multiplicants[1:]) - else: - tmp = M.codomain.zeros() - return M, is_composed, tmp - - -@dataclass -class _WeightGather: + if not isinstance(M, ComposedLinearOperator): + return M, None, () + + mass_index = next((k for k, op in enumerate(M.multiplicants) 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.multiplicants[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: """ - MPI setup for gathering a 1d cut of an array-valued weight along ``dim_reduce``. + Apply the operator ``M`` (see ``_operator_to_invert``) to ``rhs``, with the mass matrix + replaced by ``apply_inverse(x, out)`` (an approximate inverse). - In the directions other than ``dim_reduce``, the weight is taken at the global mid - point; the ranks owning the mid element there (exactly one rank per slab along - ``dim_reduce``) form ``subcomm``. ``root`` is the lowest of these ranks in ``comm``. + 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.multiplicants + 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: + ops[k].dot(x, out=y) + x = y + return out - comm: object - subcomm: object - root: int - @classmethod - def setup(cls, derham, dim_reduce: int) -> "_WeightGather | None": - """Collective on ``derham.comm``; returns None in serial runs.""" - comm = derham.comm - if isinstance(comm, (SerialComm, type(None))): - return None +# number of Gauss-Legendre points per direction for averaging callable weights +_N_AVG_CALLABLE = 16 - dom_dec = derham.domain_decomposition - rank = comm.Get_rank() - is_selected = all( - dom_dec.starts[i] <= derham.num_elements[i] // 2 <= dom_dec.ends[i] - for i in range(len(derham.num_elements)) - 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}") - return cls(comm, subcomm, root) +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`: -def _reduced_weight_1d(weight, c: int, d: int, derham, gather: _WeightGather | None) -> Callable | xp.ndarray: - """ - 1d weight along direction ``d``: the 3d weight at the mid point (0.5) of the other directions. + * ``"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 ---------- @@ -529,70 +502,103 @@ def _reduced_weight_1d(weight, c: int, d: int, derham, gather: _WeightGather | N logical coordinates, its values at the local quadrature points, or None (unit weight). c, d : int - Component and direction (only used for logging and for the callable case). + Component and direction (``c`` is only used for logging). derham : Derham Discrete de Rham sequence of the mass operator. - gather : _WeightGather | None - MPI setup from ``_WeightGather.setup``; None in serial runs. + reduction : str + ``"midpoint"`` or ``"average"``. Returns ------- Callable | xp.ndarray - A function of the 1d quadrature points, or the weight at all global 1d - quadrature points (collective on ``gather.comm`` in parallel runs). + 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): - # evaluate the 3d weight on the "meshgrid" (0.5, ..., e, ..., 0.5) - 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( - weight( - *[xp.array(xp.full_like(f, 0.5)) if i != d else xp.array(f) for i in range(n_dims)], - ).squeeze(), - ) + # 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): - s = weight.shape logger.debug(f"{weight.shape = } for component {c} and direction {d}.") dom_dec = derham.domain_decomposition - 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) - cut = tuple(slice(None) if i == d else mid[i] for i in range(n_dims)) - local_fun = xp.ascontiguousarray(weight[cut], dtype=float) - - logger.debug(f"{fun.size = } for component {c} and direction {d} before gathering on all processes.") - if gather is not None: - # local sizes differ if num_elements[d] is not divisible by the number of processes - if gather.subcomm != MPI.COMM_NULL: - counts = gather.subcomm.allgather(local_fun.size) - displs = [sum(counts[:j]) for j in range(len(counts))] - gather.subcomm.Allgatherv(local_fun, [fun, counts, displs, MPI.DOUBLE]) - gather.comm.Bcast(fun, root=gather.root) + 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: - fun[:] = local_fun - logger.debug(f"{fun.shape = } for component {c} and direction {d} after gathering on all processes.") - return fun + # 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() + + 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)}") @@ -792,8 +798,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 @@ -827,20 +844,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 @@ -863,12 +871,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/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..6df252615 100644 --- a/src/struphy/feec/tests/test_preconditioner_transpose.py +++ b/src/struphy/feec/tests/test_preconditioner_transpose.py @@ -2,7 +2,8 @@ @pytest.mark.parametrize("dim_reduce", [0, 2]) -def test_mass_preconditioner_transpose(dim_reduce): +@pytest.mark.parametrize("weight_reduction", ["midpoint", "average"]) +def test_mass_preconditioner_transpose(dim_reduce, weight_reduction): """The transposes of the mass-matrix preconditioners keep their class and options.""" from struphy import domains @@ -16,11 +17,12 @@ 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) 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 pc_diag_T = MassMatrixDiagonalPreconditioner(mass_ops.M1, apply_bc=False).transpose() assert type(pc_diag_T) is MassMatrixDiagonalPreconditioner From c33b5541a43ae0eff41ac8aae77b4589aa8bf45d Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 11:02:39 +0200 Subject: [PATCH 04/18] Use multiplicands; bump feectools (ComposedKroneckerStencilMatrix) - Rename uses of ComposedLinearOperator.multiplicants to multiplicands (preconditioner, projectors, multigrid coarsening, comments). - linop_nbytes: sum over all child attributes, so that the factors of a ComposedKroneckerStencilMatrix (multiplicands and mats) are counted. - Bump feectools to the commit with ComposedKroneckerStencilMatrix. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- src/struphy/feec/memory.py | 7 ++++--- src/struphy/feec/preconditioner.py | 10 +++++----- src/struphy/feec/projectors.py | 6 +++--- src/struphy/linear_algebra/multigrid/coarsen.py | 2 +- src/struphy/propagators/variational_resistivity.py | 2 +- src/struphy/propagators/variational_viscosity.py | 2 +- 7 files changed, 16 insertions(+), 15 deletions(-) diff --git a/feectools b/feectools index 6ee8f082c..fa6fe9040 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 6ee8f082c5e1f9cdfd25ed726f6f99f4b5ecd97b +Subproject commit fa6fe90403cf34935709616712b6a199b2554036 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 ec493f43c..6e6f37333 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -412,7 +412,7 @@ def _boundary_conditions(mass_operator: WeightedMassOperator, apply_bc: bool) -> """ if not apply_bc or not isinstance(mass_operator.M0, ComposedLinearOperator): return None - last = mass_operator.M0.multiplicants[-1] + last = mass_operator.M0.multiplicands[-1] if not isinstance(last, BoundaryOperator): return None return last.bc @@ -431,7 +431,7 @@ def _operator_to_invert( The operator. mass_index : int | None - Position of the mass matrix ``mass_operator.matrix`` in ``M.multiplicants``, or + 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). tmp_vectors : tuple[Vector, ...] @@ -442,9 +442,9 @@ def _operator_to_invert( if not isinstance(M, ComposedLinearOperator): return M, None, () - mass_index = next((k for k, op in enumerate(M.multiplicants) if op is mass_operator.matrix), 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.multiplicants[1:]) + tmp_vectors = tuple(op.codomain.zeros() for op in M.multiplicands[1:]) return M, mass_index, tmp_vectors @@ -468,7 +468,7 @@ def _apply_composed( out = M.codomain.zeros() return apply_inverse(rhs, out) - ops = M.multiplicants + ops = M.multiplicands if out is None: out = ops[0].codomain.zeros() 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/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}") From fb2659d5cf9b7573f4d47ccfbeba696ec159b059 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 12:06:52 +0200 Subject: [PATCH 05/18] Bump feectools: KroneckerSumSolver, derivative properties Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index fa6fe9040..c774ecf09 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit fa6fe90403cf34935709616712b6a199b2554036 +Subproject commit c774ecf0985c23d8f6f851fb1606840eb9494726 From a30ca02508e8d39cf8325111f0583cf4610d96cc Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 12:11:01 +0200 Subject: [PATCH 06/18] KroneckerPreconditioner base class for the mass preconditioners - New base class KroneckerPreconditioner: P = B E S A_c~^{-1} S E^T B^T, where A_c~ approximates the core operator of the composition of a mass operator and S is an optional diagonal scaling (D_hat^{1/2} D^{-1/2}, Loli-Sangalli-Tani). It holds the common interface (matrix, solver, solve, dot, properties); subclasses build the approximation. - MassMatrixPreconditioner and MassMatrixDiagonalPreconditioner are subclasses; their results are unchanged (checked serial and on 2 ranks). - MassMatrixPreconditioner gets the option diagonal_scaling (default False), e.g. on HollowTorus the PCG iterations of M1 drop from 20 to 8. - Helpers for local diagonals of Kronecker/sum/block operators. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 471 ++++++++++-------- .../tests/test_preconditioner_transpose.py | 12 +- 2 files changed, 266 insertions(+), 217 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 6e6f37333..5fdb5dae3 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -8,11 +8,18 @@ from feectools.api.essential_bc import apply_essential_bc_stencil from feectools.ddm.cart import CartDecomposition, DomainDecomposition from feectools.fem.tensor import TensorFemSpace -from feectools.linalg.basic import ComposedLinearOperator, LinearOperator, Vector, VectorSpace +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.stencil import StencilDiagonalMatrix, StencilMatrix, StencilVectorSpace from line_profiler import profile from maybempi import MPI, SerialComm from scipy import sparse @@ -23,52 +30,49 @@ logger = logging.getLogger("struphy") -class MassMatrixPreconditioner(LinearOperator): +class KroneckerPreconditioner(LinearOperator): r""" - Preconditioner for inverting 3d weighted mass matrices. + Base class for preconditioners built from Kronecker approximations. - 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`: + 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 - * 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``). - * 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). + .. math:: - 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 :meth:`solve`). + 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 : WeightedMassOperator - The weighted mass operator for which the approximate inverse is needed. + 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 - Axis along which the weight is kept. - - 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 diagonal scaling :math:`S`. """ def __init__( self, mass_operator: WeightedMassOperator, apply_bc: bool = True, - dim_reduce: int = 0, - weight_reduction: str = "midpoint", + 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 @@ -77,45 +81,69 @@ def __init__( self._codomain = mass_operator.codomain self._domain = mass_operator.domain self._apply_bc = apply_bc - self._dim_reduce = dim_reduce - self._weight_reduction = weight_reduction + self._diagonal_scaling = diagonal_scaling - n_dims = self._femspace.ldim - assert n_dims == 3 # other dims not yet implemented - assert dim_reduce < n_dims - assert weight_reduction in ("midpoint", "average"), f"Unknown weight_reduction {weight_reduction!r}." + assert self._femspace.ldim == 3 # other dims not yet implemented # boundary conditions are only imposed if the mass operator has a BoundaryOperator - bc = _boundary_conditions(mass_operator, apply_bc) - apply_bc = bc is not None + self._bc = _boundary_conditions(mass_operator, apply_bc) - derham = mass_operator.derham - logger.debug(f"{derham.num_elements = }, {derham.bcs = }, {derham.degree = }") + # approximation of the core operator and its exact inverse + self._matrix, self._solver = self._build_approximation(self._bc) - def weight_1d(c: int, d: int) -> Callable | xp.ndarray: - if d == dim_reduce: - return _reduced_weight_1d(mass_operator.weights[c][c], c, d, derham, weight_reduction) - return _ones_1d + # 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 () - # Kronecker approximation of the mass matrix and its exact inverse - self._matrix, self._solver = _kronecker_approximation(mass_operator, bc, weight_1d) + # 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) + + # -------------------------------------- + # To be implemented by subclasses + # -------------------------------------- + def _build_approximation(self, bc: list | None) -> tuple[LinearOperator, LinearOperator]: + """ + Approximation of the core operator and its exact inverse. - # mass 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, apply_bc) + Parameters + ---------- + bc : list | None + Boundary conditions from ``_boundary_conditions``; None for no boundary conditions. + + Returns + ------- + matrix : LinearOperator + The approximation :math:`\\tilde A_c`. + + solver : LinearOperator + Its exact inverse. + """ + raise NotImplementedError + 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.") + + def transpose(self, conjugate: bool = False) -> KroneckerPreconditioner: + raise NotImplementedError + + # -------------------------------------- + # Common interface + # -------------------------------------- @property def space(self) -> VectorSpace: """Stencil-/BlockVectorSpace or PolarDerhamSpace.""" return self._space @property - def matrix(self) -> KroneckerStencilMatrix | BlockLinearOperator: - """Approximation of the input mass matrix as KroneckerStencilMatrix (block diagonal for vector-valued spaces).""" + 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 | BlockLinearOperator: - """KroneckerLinearSolver (block diagonal for vector-valued spaces) for exactly inverting the approximate mass matrix self.matrix.""" + def solver(self) -> LinearOperator: + """Exact inverse of the approximation self.matrix.""" return self._solver @property @@ -133,28 +161,42 @@ def dtype(self): return self._dtype @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 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, - ) + 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: 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 mass - matrix M itself is replaced by the Kronecker solver. + 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 ---------- @@ -167,37 +209,116 @@ def solve(self, rhs: Vector, out: Vector | None = None) -> Vector: 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 - return _apply_composed(self._M, self._mass_index, self.solver.dot, self._tmp_vectors, rhs, out) + return _apply_composed(self._M, self._mass_index, self._apply_inverse, self._tmp_vectors, rhs, out) 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``). + * 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). -class MassMatrixDiagonalPreconditioner(LinearOperator): + 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 + Axis along which the weight is kept. + + 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 = 0, + weight_reduction: str = "midpoint", + diagonal_scaling: bool = False, + ): + assert 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) + + 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 weight_reduction(self) -> str: + """How the weight is reduced in the directions other than ``dim_reduce``: ``"midpoint"`` or ``"average"``.""" + return self._weight_reduction + + 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(KroneckerPreconditioner): r""" Preconditioner for inverting 3d weighted mass matrices. The mass matrix is approximated by @@ -221,32 +342,15 @@ 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!" - - 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 - - n_dims = self._femspace.ldim - assert n_dims == 3 # other dims not yet implemented - - # boundary conditions are only imposed if the mass operator has a BoundaryOperator - bc = _boundary_conditions(mass_operator, apply_bc) - apply_bc = bc is not None + def __init__(self, mass_operator: WeightedMassOperator, apply_bc: bool = True): + super().__init__(mass_operator, apply_bc=apply_bc, diagonal_scaling=True) + def _build_approximation(self, bc): # mass matrix on the logical domain (unit weights) as Kronecker product, and its exact inverse - self._matrix, self._solver = _kronecker_approximation(mass_operator, bc, lambda c, d: _ones_1d) + return _kronecker_approximation(self._mass_operator, bc, lambda c, d: _ones_1d) - # mass 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, apply_bc) - - # Need to assemble the logical mass matrix to extract the coefficients + def _build_scaling(self): + # D^{-1/2} of the mass matrix and \hat D^{1/2} of the assembled logical mass matrix 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) ] @@ -259,39 +363,9 @@ def __init__(self, mass_operator, apply_bc=True): 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) + return (self._M_invsrqt_diag, self._logM_srqt_diag) - self._tmp_vector_no_bc = [self._mass_operator.matrix.codomain.zeros() for i in range(2)] - - @property - def space(self): - """Stencil-/BlockVectorSpace or PolarDerhamSpace.""" - return self._space - - @property - def matrix(self): - """Mass matrix on the logical domain as KroneckerStencilMatrix.""" - return self._matrix - - @property - def solver(self): - """KroneckerLinearSolver or BlockDiagonalSolver for exactly inverting the approximate mass matrix self.matrix.""" - return self._solver - - @property - def domain(self): - """The domain of the linear operator - an element of Vectorspace""" - return self._space - - @property - def codomain(self): - """The codomain of the linear operator - an element of Vectorspace""" - return self._codomain - - @property - def dtype(self): - return self._dtype - - def update_mass_operator(self, mass_operator): + def update_mass_operator(self, mass_operator: WeightedMassOperator) -> None: """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!" @@ -304,92 +378,15 @@ def update_mass_operator(self, mass_operator): M, mass_index, _ = _operator_to_invert(mass_operator, apply_bc) assert mass_index == self._mass_index, "The updated mass operator must have the same structure." self._M = M + # updated in place, so self._scaling stays valid self._M_invsrqt_diag = self._mass_operator.matrix.diagonal(inverse=True, sqrt=True, out=self._M_invsrqt_diag) - def transpose(self, conjugate=False): + def transpose(self, conjugate: bool = False) -> MassMatrixDiagonalPreconditioner: """ Returns the transposed operator. """ return MassMatrixDiagonalPreconditioner(self._mass_operator.transpose(), self._apply_bc) - 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. - - 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. - """ - - assert isinstance(rhs, Vector) - assert rhs.space == self._mass_operator.matrix.domain - - # 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 - - @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}`. - - Parameters - ---------- - rhs : Vector - The right-hand side vector. - - out : Vector, optional - If given, the output vector will be written into this vector in-place. - - Returns - ------- - out : Vector - The result of :math:`(B * E * M^{-1} * E^T * B^T) * rhs`. - """ - - assert isinstance(rhs, Vector) - assert rhs.space == self._space - if out is not None: - assert isinstance(out, Vector) - assert out.space == self._space - - return _apply_composed(self._M, self._mass_index, self._solve_no_bc, self._tmp_vectors, rhs, out) - - def dot(self, v, out=None): - """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: - assert isinstance(out, Vector) - assert out.space == self.codomain - self.solve(v, out=out) - - return out - # -------------------------------------------------------------------------------------- # Helper functions for building Kronecker approximations of mass matrices @@ -399,6 +396,50 @@ def _ones_1d(e: xp.ndarray) -> xp.ndarray: return xp.ones(e.size, dtype=float) +def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: + """ + Local (rows owned by this process) diagonal of ``op``, one array per block of its domain. + + Supported are StencilMatrix, StencilDiagonalMatrix, KroneckerStencilMatrix (1d factors), + sums, scalings and block-diagonal BlockLinearOperators of these. + """ + if isinstance(op, BlockLinearOperator): + return [a for i in range(op.n_block_rows) for a in _local_diagonal(op[i, i])] + 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." + off = A.codomain.pads[0] * A.codomain.shifts[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]) + + def _boundary_conditions(mass_operator: WeightedMassOperator, apply_bc: bool) -> list | None: """ Boundary conditions of the mass operator, taken from the BoundaryOperator diff --git a/src/struphy/feec/tests/test_preconditioner_transpose.py b/src/struphy/feec/tests/test_preconditioner_transpose.py index 6df252615..23a9296d5 100644 --- a/src/struphy/feec/tests/test_preconditioner_transpose.py +++ b/src/struphy/feec/tests/test_preconditioner_transpose.py @@ -3,7 +3,8 @@ @pytest.mark.parametrize("dim_reduce", [0, 2]) @pytest.mark.parametrize("weight_reduction", ["midpoint", "average"]) -def test_mass_preconditioner_transpose(dim_reduce, weight_reduction): +@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 @@ -17,12 +18,19 @@ def test_mass_preconditioner_transpose(dim_reduce, weight_reduction): 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, weight_reduction=weight_reduction) + 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 From 4d5e1bc5b5daee5a28ea8bbf1f7b4a64f981576e Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:02:24 +0200 Subject: [PATCH 07/18] StiffnessPreconditioner for d^T M_{k+1} d + sigma M_k (grad) New subclass of KroneckerPreconditioner for the stabilized stiffness operators G^T M1 G + sigma M0 (and C^T M2 C + sigma M1, D^T M3 D + sigma M2). With unit weights each diagonal block is a sum of Kronecker products of 1d stiffness and mass matrices, inverted exactly by fast diagonalization (KroneckerSumSolver). For grad this is exact on the logical cube; curl and div are block Jacobi for now (kernel correction follows). The geometry is included by diagonal scaling (default), with the diagonal of the stiffness operator computed by probing. Polar splines are not supported yet. Tests: approximation against the operator on the unit cube, PCG for grad (2 iterations on the unit cube, HollowCylinder 89 -> 48), independence of the MPI decomposition, options. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 289 +++++++++++++++++- .../tests/test_stiffness_preconditioner.py | 162 ++++++++++ 2 files changed, 449 insertions(+), 2 deletions(-) create mode 100644 src/struphy/feec/tests/test_stiffness_preconditioner.py diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 5fdb5dae3..0c4661c84 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -1,12 +1,15 @@ 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, @@ -18,14 +21,14 @@ ) 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.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") @@ -388,6 +391,164 @@ def transpose(self, conjugate: bool = False) -> MassMatrixDiagonalPreconditioner return MassMatrixDiagonalPreconditioner(self._mass_operator.transpose(), self._apply_bc) +class StiffnessPreconditioner(KroneckerPreconditioner): + r""" + Preconditioner for the stabilized stiffness operators + + .. math:: + + 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)\,, + + 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`. + + 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, + + .. math:: + + 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 \,, + + 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), so :math:`\sigma > 0` is required for a good preconditioner. + + The geometry is included by diagonal scaling (default, see :class:`KroneckerPreconditioner`). + 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. + """ + + _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, + ): + assert derivative in self._FORMS, f"derivative must be one of {tuple(self._FORMS)}, got {derivative!r}." + assert sigma >= 0.0 + + 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._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") + + if derivative != "grad" and sigma == 0.0: + logger.warning( + f"StiffnessPreconditioner for {derivative!r} with sigma = 0: the operator has a large kernel." + ) + + super().__init__(getattr(mass_ops, f"M{k}"), apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) + + @property + def derivative(self) -> str: + """``"grad"``, ``"curl"`` or ``"div"``.""" + return self._derivative + + @property + def sigma(self) -> float: + """Coefficient of the mass term.""" + return self._sigma + + @property + 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) + + 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, _ones_1d) + + 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, _ones_1d) + D_d = _difference_matrix_1d(M_d.shape[0], M_D.shape[0]) + S_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) + ) + + solverblocks.append(KroneckerSumSolver(coeff_space, stiffness, mass, sigma=self._sigma)) + matrixblocks.append(_kronecker_sum(coeff_space, local_stiffness, local_mass, self._sigma)) + + if is_scalar: + return matrixblocks[0], solverblocks[0] + return ( + _block_diagonal(femspace.coeff_space, matrixblocks), + _block_diagonal(femspace.coeff_space, solverblocks), + ) + + def _core_diagonal(self): + space = self._matrix.domain + return _diagonal_operator(space, _probe_diagonal(self.core_operator, space)) + + def transpose(self, conjugate: bool = False) -> StiffnessPreconditioner: + """The operator is symmetric, so is the preconditioner.""" + return self + + # -------------------------------------------------------------------------------------- # Helper functions for building Kronecker approximations of mass matrices # -------------------------------------------------------------------------------------- @@ -824,6 +985,130 @@ def _block_diagonal(space: BlockVectorSpace, blocks: list[LinearOperator]) -> Bl ) +# -------------------------------------------------------------------------------------- +# 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 _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): """ Solve the equation Ax = b for x, assuming A is a circulant matrix. 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..5d6c23b0f --- /dev/null +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -0,0 +1,162 @@ +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) + 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, diagonal_scaling=False)), + ("scaled", StiffnessPreconditioner(mass_ops, "grad", sigma=sigma)), + ): + 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", ["grad", "curl", "div"]) +def test_stiffness_preconditioner_mpi(derivative): + """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).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 + with pytest.raises(AssertionError): + StiffnessPreconditioner(mass_ops, "rot") + with pytest.raises(AssertionError): + StiffnessPreconditioner(mass_ops, "grad", sigma=-1.0) From 474e37a511090984d4ae5d51dfd6cbb8db244bde Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:03:49 +0200 Subject: [PATCH 08/18] StiffnessPreconditioner: kernel correction for curl Block Jacobi alone overestimates C^T M2 C + sigma M1 on gradients (where it equals sigma M1) and was never better than MassMatrixPreconditioner. With kernel_correction (default, needs sigma > 0) the preconditioner is P_BJ + sigma^{-1} G P_G G^T, with P_G the preconditioner of G^T M1 G, which is exact for gradients on the logical cube. PCG iterations (12x16x6, degree (2,3,2), sigma=1): Cuboid 432 -> 14 (mass preconditioner 144), HollowTorus 1947 -> 150 (mass preconditioner 716). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 48 ++++++++++++--- .../tests/test_stiffness_preconditioner.py | 58 ++++++++++++++++++- 2 files changed, 98 insertions(+), 8 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 0c4661c84..09d66e7a3 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -417,8 +417,19 @@ class StiffnessPreconditioner(KroneckerPreconditioner): :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), so :math:`\sigma > 0` is required for a good preconditioner. + inverse on the logical cube. For curl and div, the off-diagonal blocks are neglected + (block Jacobi, :math:`P_{BJ}`). + + Block Jacobi overestimates :math:`A` on the kernel of :math:`\mathbb d` (gradients for + curl), where :math:`A = \sigma \mathbb M_k`. With ``kernel_correction`` (default, needs + :math:`\sigma > 0`), the preconditioner for curl is + + .. math:: + + P = P_{BJ} + \sigma^{-1} \, \mathbb G \, P_{\mathbb G} \, \mathbb G^T \,, + + where :math:`P_{\mathbb G}` is the preconditioner for :math:`\mathbb G^T \mathbb M_1 \mathbb G`; + it is exact for gradients on the logical cube. The geometry is included by diagonal scaling (default, see :class:`KroneckerPreconditioner`). Polar splines are not supported yet. @@ -440,6 +451,10 @@ class StiffnessPreconditioner(KroneckerPreconditioner): diagonal_scaling : bool Whether to correct the approximation with the diagonals of :math:`A` and of its approximation. + + kernel_correction : bool + For curl: whether to add the correction on the kernel of the derivative (needs + ``sigma > 0``). Ignored for grad. """ _FORMS = {"grad": 0, "curl": 1, "div": 2} @@ -451,9 +466,13 @@ def __init__( sigma: float = 0.0, apply_bc: bool = True, diagonal_scaling: bool = True, + kernel_correction: bool = True, ): assert derivative in self._FORMS, f"derivative must be one of {tuple(self._FORMS)}, got {derivative!r}." assert sigma >= 0.0 + kernel_correction = kernel_correction and derivative == "curl" + 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: @@ -467,13 +486,16 @@ def __init__( # tensor-product derivative without boundary operators (no polar splines) self._d = getattr(derham, f"{derivative}_bcfree") - if derivative != "grad" and sigma == 0.0: - logger.warning( - f"StiffnessPreconditioner for {derivative!r} with sigma = 0: the operator has a large kernel." - ) - 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: + d_prev = derham.grad + P_prev = StiffnessPreconditioner(mass_ops, "grad", apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) + self._kernel_correction = (1.0 / sigma) * (d_prev @ P_prev @ d_prev.T) + self._tmp_kernel = self.codomain.zeros() + @property def derivative(self) -> str: """``"grad"``, ``"curl"`` or ``"div"``.""" @@ -484,6 +506,18 @@ def sigma(self) -> float: """Coefficient of the mass term.""" return self._sigma + @property + 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 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.""" diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py index 5d6c23b0f..50c5e9dfd 100644 --- a/src/struphy/feec/tests/test_stiffness_preconditioner.py +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -39,7 +39,9 @@ def test_stiffness_approximation_unit_cube(derivative, sigma): ) mass_ops = WeightedMassOperators(derham, domains.Cuboid()) - P = StiffnessPreconditioner(mass_ops, derivative, sigma=sigma, apply_bc=False, diagonal_scaling=False) + P = StiffnessPreconditioner( + mass_ops, derivative, sigma=sigma, apply_bc=False, diagonal_scaling=False, kernel_correction=False + ) A = P.core_operator x = create_equal_random_arrays(derham.fem_spaces[FORMS[derivative]], seed=1)[1] @@ -110,6 +112,60 @@ def test_stiffness_preconditioner_grad(mapping, sigma): assert niter["scaled"] <= niter["unit"] +@pytest.mark.parametrize("derivative", ["curl"]) +@pytest.mark.parametrize("mapping", ["Cuboid", "HollowCylinder"]) +def test_stiffness_preconditioner_kernel_correction(derivative, mapping): + """PCG for C^T M2 C + sigma M1: 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_mpi(derivative): """The preconditioner must not depend on the MPI decomposition.""" From d7b020a5e7cd6ac648923a7916762c592c434668 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:04:50 +0200 Subject: [PATCH 09/18] StiffnessPreconditioner: kernel correction for div Same idea as for curl: P_BJ + sigma^{-1} C P_C C^T, where P_C is the curl block Jacobi preconditioner (sigma = 0 is fine there: on the range of C^T the gradients do not matter, G^T C^T = 0). PCG iterations (12x16x6, degree (2,3,2), sigma=1): Cuboid 411 -> 18 (mass preconditioner 142), HollowTorus 2734 -> 175 (mass preconditioner 636). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 24 ++++++++++++------- .../tests/test_stiffness_preconditioner.py | 6 ++--- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 09d66e7a3..c972d6885 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -421,15 +421,18 @@ class StiffnessPreconditioner(KroneckerPreconditioner): (block Jacobi, :math:`P_{BJ}`). Block Jacobi overestimates :math:`A` on the kernel of :math:`\mathbb d` (gradients for - curl), where :math:`A = \sigma \mathbb M_k`. With ``kernel_correction`` (default, needs - :math:`\sigma > 0`), the preconditioner for curl is + curl, curls for div), where :math:`A = \sigma \mathbb M_k`. With ``kernel_correction`` + (default, needs :math:`\sigma > 0`), the preconditioner is .. math:: - P = P_{BJ} + \sigma^{-1} \, \mathbb G \, P_{\mathbb G} \, \mathbb G^T \,, + P = P_{BJ} + \sigma^{-1} \, \mathbb d_- \, P_- \, \mathbb d_-^T \,, - where :math:`P_{\mathbb G}` is the preconditioner for :math:`\mathbb G^T \mathbb M_1 \mathbb G`; - it is exact for gradients on the logical cube. + 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`. The geometry is included by diagonal scaling (default, see :class:`KroneckerPreconditioner`). Polar splines are not supported yet. @@ -453,7 +456,7 @@ class StiffnessPreconditioner(KroneckerPreconditioner): approximation. kernel_correction : bool - For curl: whether to add the correction on the kernel of the derivative (needs + For curl and div: whether to add the correction on the kernel of the derivative (needs ``sigma > 0``). Ignored for grad. """ @@ -470,7 +473,7 @@ def __init__( ): assert derivative in self._FORMS, f"derivative must be one of {tuple(self._FORMS)}, got {derivative!r}." assert sigma >= 0.0 - kernel_correction = kernel_correction and derivative == "curl" + 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).") @@ -491,8 +494,11 @@ def __init__( # correction on the kernel of the derivative: sigma^{-1} d_prev P_prev d_prev^T self._kernel_correction = None if kernel_correction: - d_prev = derham.grad - P_prev = StiffnessPreconditioner(mass_ops, "grad", apply_bc=apply_bc, diagonal_scaling=diagonal_scaling) + 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 + ) self._kernel_correction = (1.0 / sigma) * (d_prev @ P_prev @ d_prev.T) self._tmp_kernel = self.codomain.zeros() diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py index 50c5e9dfd..cc4912c8f 100644 --- a/src/struphy/feec/tests/test_stiffness_preconditioner.py +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -112,11 +112,11 @@ def test_stiffness_preconditioner_grad(mapping, sigma): assert niter["scaled"] <= niter["unit"] -@pytest.mark.parametrize("derivative", ["curl"]) +@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: with kernel correction, fewer iterations than with block Jacobi - alone and than with the mass-matrix preconditioner.""" + """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 4dfb9dde14686aa4185a63d5ae1e9faea337cc58 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:09:35 +0200 Subject: [PATCH 10/18] StiffnessPreconditioner: option weights="average" Fast diagonalization needs one 1d mass matrix per direction for all terms, so the weights are approximated by a common separable shape prod_d phi_d (mean over the other directions, averaged over the diagonal blocks of M_{k+1}) times the mean weight of each term. Exact for constant anisotropic weights. Default stays weights="unit" with diagonal scaling. PCG iterations (12x16x6, degree (2,3,2); unit+scaling -> average+scaling): stretched Cuboid grad 56 -> 2, HollowTorus grad 58 -> 41, curl (sigma=1) 150 -> 46, div (sigma=1) 175 -> 48. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 94 +++++++++++++++++-- .../tests/test_stiffness_preconditioner.py | 52 +++++++++- 2 files changed, 137 insertions(+), 9 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index c972d6885..73824a115 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -434,7 +434,15 @@ class StiffnessPreconditioner(KroneckerPreconditioner): 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`. - The geometry is included by diagonal scaling (default, see :class:`KroneckerPreconditioner`). + The geometry is included by diagonal scaling (default, see :class:`KroneckerPreconditioner`), + and/or with ``weights="average"``: 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 @@ -458,6 +466,10 @@ class StiffnessPreconditioner(KroneckerPreconditioner): 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 + ``"unit"`` (operator on the logical cube, default) or ``"average"`` (separable + approximation of the weights, see above). """ _FORMS = {"grad": 0, "curl": 1, "div": 2} @@ -470,9 +482,11 @@ def __init__( apply_bc: bool = True, diagonal_scaling: bool = True, kernel_correction: bool = True, + weights: str = "unit", ): 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).") @@ -485,6 +499,7 @@ def __init__( 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") @@ -497,7 +512,12 @@ def __init__( 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 + 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() @@ -512,6 +532,11 @@ def sigma(self) -> float: """Coefficient of the mass term.""" return self._sigma + @property + def weights(self) -> str: + """``"unit"`` or ``"average"``.""" + return self._weights + @property def kernel_correction(self) -> LinearOperator | None: """The correction on the kernel of the derivative, or None.""" @@ -540,6 +565,18 @@ def _build_approximation(self, bc): # 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 + matrixblocks = [] solverblocks = [] for c, comp in enumerate(comps): @@ -548,14 +585,14 @@ def _build_approximation(self, bc): 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, _ones_1d) + 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, _ones_1d) + 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 = D_d.T @ M_D @ D_d + 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]) @@ -570,8 +607,9 @@ def _build_approximation(self, bc): else _process_local_matrix_1d(xp.asarray(S_d), comp.coeff_space, d, domain_decomposition) ) - solverblocks.append(KroneckerSumSolver(coeff_space, stiffness, mass, sigma=self._sigma)) - matrixblocks.append(_kronecker_sum(coeff_space, local_stiffness, local_mass, self._sigma)) + 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)) if is_scalar: return matrixblocks[0], solverblocks[0] @@ -1051,6 +1089,48 @@ def _derivative_directions(derivative: LinearOperator) -> list[dict[int, int]]: 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]: diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py index cc4912c8f..6521099c4 100644 --- a/src/struphy/feec/tests/test_stiffness_preconditioner.py +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -167,7 +167,55 @@ def test_stiffness_preconditioner_kernel_correction(derivative, mapping): @pytest.mark.parametrize("derivative", ["grad", "curl", "div"]) -def test_stiffness_preconditioner_mpi(derivative): +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 @@ -190,7 +238,7 @@ def test_stiffness_preconditioner_mpi(derivative): ) 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).dot(v)] + 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)) From 079f672284457d1faf14c9da92d40b9ef3470192 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:18:43 +0200 Subject: [PATCH 11/18] StiffnessPreconditioner: weights="average" as default Average + diagonal scaling gave the fewest PCG iterations in all cases tested (e.g. HollowTorus curl 150 -> 46, div 175 -> 48). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 10 +++++----- .../feec/tests/test_stiffness_preconditioner.py | 14 ++++++++++---- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 73824a115..5785da683 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -434,8 +434,8 @@ class StiffnessPreconditioner(KroneckerPreconditioner): 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`. - The geometry is included by diagonal scaling (default, see :class:`KroneckerPreconditioner`), - and/or with ``weights="average"``: fast diagonalization needs the same 1d mass matrix in + 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 @@ -468,8 +468,8 @@ class StiffnessPreconditioner(KroneckerPreconditioner): ``sigma > 0``). Ignored for grad. weights : str - ``"unit"`` (operator on the logical cube, default) or ``"average"`` (separable - approximation of the weights, see above). + ``"average"`` (separable approximation of the weights, see above; default) or + ``"unit"`` (operator on the logical cube). """ _FORMS = {"grad": 0, "curl": 1, "div": 2} @@ -482,7 +482,7 @@ def __init__( apply_bc: bool = True, diagonal_scaling: bool = True, kernel_correction: bool = True, - weights: str = "unit", + weights: str = "average", ): assert derivative in self._FORMS, f"derivative must be one of {tuple(self._FORMS)}, got {derivative!r}." assert sigma >= 0.0 diff --git a/src/struphy/feec/tests/test_stiffness_preconditioner.py b/src/struphy/feec/tests/test_stiffness_preconditioner.py index 6521099c4..6305551ca 100644 --- a/src/struphy/feec/tests/test_stiffness_preconditioner.py +++ b/src/struphy/feec/tests/test_stiffness_preconditioner.py @@ -40,7 +40,13 @@ def test_stiffness_approximation_unit_cube(derivative, sigma): mass_ops = WeightedMassOperators(derham, domains.Cuboid()) P = StiffnessPreconditioner( - mass_ops, derivative, sigma=sigma, apply_bc=False, diagonal_scaling=False, kernel_correction=False + 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] @@ -96,8 +102,8 @@ def test_stiffness_preconditioner_grad(mapping, sigma): niter = {} for label, P in ( ("none", None), - ("unit", StiffnessPreconditioner(mass_ops, "grad", sigma=sigma, diagonal_scaling=False)), - ("scaled", StiffnessPreconditioner(mass_ops, "grad", sigma=sigma)), + ("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) @@ -259,7 +265,7 @@ def test_stiffness_preconditioner_options(): 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 + 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): From b28e206e420b39c10f89a6f6cb3673e720cdd97c Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:23:26 +0200 Subject: [PATCH 12/18] MassMatrixDiagonalPreconditioner as subclass of MassMatrixPreconditioner - MassMatrixPreconditioner: dim_reduce=None for unit weights in all directions; update_mass_operator (rebuilds the approximation only if it depends on the weights, updates the diagonal scaling). - MassMatrixDiagonalPreconditioner is now MassMatrixPreconditioner(dim_reduce=None, diagonal_scaling=True), kept for its name (solver options). The diagonal of the logical mass matrix is taken from the Kronecker approximation instead of assembling the 3d logical mass matrix; results are unchanged up to rounding (checked on HollowTorus and IGAPolarCylinder, incl. update_mass_operator, serial and on 2 ranks). - Test: only 1d mass matrices are assembled. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 82 +++++++++---------- .../tests/test_preconditioner_transpose.py | 25 +++--- 2 files changed, 51 insertions(+), 56 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 5785da683..938d38829 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -243,7 +243,8 @@ class MassMatrixPreconditioner(KroneckerPreconditioner): * 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``). + (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). @@ -260,8 +261,8 @@ class MassMatrixPreconditioner(KroneckerPreconditioner): apply_bc : bool Whether to include boundary operators. - dim_reduce : int - Axis along which the weight is kept. + 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) @@ -277,11 +278,11 @@ def __init__( self, mass_operator: WeightedMassOperator, apply_bc: bool = True, - dim_reduce: int = 0, + dim_reduce: int | None = 0, weight_reduction: str = "midpoint", diagonal_scaling: bool = False, ): - assert dim_reduce < 3 + 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 @@ -303,11 +304,39 @@ def weight_1d(c: int, d: int) -> Callable | xp.ndarray: 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. @@ -321,7 +350,7 @@ def transpose(self, conjugate: bool = False) -> MassMatrixPreconditioner: ) -class MassMatrixDiagonalPreconditioner(KroneckerPreconditioner): +class MassMatrixDiagonalPreconditioner(MassMatrixPreconditioner): r""" Preconditioner for inverting 3d weighted mass matrices. The mass matrix is approximated by @@ -331,6 +360,9 @@ class MassMatrixDiagonalPreconditioner(KroneckerPreconditioner): 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 ----- @@ -346,43 +378,7 @@ class MassMatrixDiagonalPreconditioner(KroneckerPreconditioner): """ def __init__(self, mass_operator: WeightedMassOperator, apply_bc: bool = True): - super().__init__(mass_operator, apply_bc=apply_bc, diagonal_scaling=True) - - def _build_approximation(self, bc): - # mass matrix on the logical domain (unit weights) as Kronecker product, and its exact inverse - return _kronecker_approximation(self._mass_operator, bc, lambda c, d: _ones_1d) - - def _build_scaling(self): - # D^{-1/2} of the mass matrix and \hat D^{1/2} of the assembled logical mass matrix - 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) - return (self._M_invsrqt_diag, self._logM_srqt_diag) - - def update_mass_operator(self, mass_operator: WeightedMassOperator) -> None: - """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" - - self._mass_operator = mass_operator - - # the composition has the same structure as before, so the temporary vectors can be reused - apply_bc = _boundary_conditions(mass_operator, self._apply_bc) is not None - M, mass_index, _ = _operator_to_invert(mass_operator, apply_bc) - assert mass_index == self._mass_index, "The updated mass operator must have the same structure." - self._M = M - # updated in place, so self._scaling stays valid - self._M_invsrqt_diag = self._mass_operator.matrix.diagonal(inverse=True, sqrt=True, out=self._M_invsrqt_diag) + super().__init__(mass_operator, apply_bc=apply_bc, dim_reduce=None, diagonal_scaling=True) def transpose(self, conjugate: bool = False) -> MassMatrixDiagonalPreconditioner: """ diff --git a/src/struphy/feec/tests/test_preconditioner_transpose.py b/src/struphy/feec/tests/test_preconditioner_transpose.py index 23a9296d5..4ca99f937 100644 --- a/src/struphy/feec/tests/test_preconditioner_transpose.py +++ b/src/struphy/feec/tests/test_preconditioner_transpose.py @@ -37,12 +37,13 @@ def test_mass_preconditioner_transpose(dim_reduce, weight_reduction, diagonal_sc 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 @@ -52,18 +53,16 @@ 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) - # 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) + assert isinstance(pc, MassMatrixPreconditioner) + assert pc.dim_reduce is None and pc.diagonal_scaling + assert created and all(ldim == 1 for ldim in created) From 97dd117360bc992dffc77c80e07d3fac3d1be3a1 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:43:01 +0200 Subject: [PATCH 13/18] Fix diagonal scaling for mass matrices with zero diagonal blocks WeightedMassOperator sets blocks with zero weight to None, also after reassembly (e.g. Mrho in VariationalDensityEvolve). _local_diagonal raised for such blocks, so update_mass_operator of MassMatrixDiagonalPreconditioner failed in the VariationalBarotropicFluid and VariationalPressurelessFluid model tests. A missing diagonal block now gives a zero diagonal (scaling 1 for that block); before the refactor, the stale diagonal of the previous assembly was kept. Checked: all models with VariationalDensityEvolve/Viscosity/Resistivity (11 models) and the variational propagator tests pass. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 11 ++++-- .../tests/test_preconditioner_transpose.py | 35 +++++++++++++++++++ 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 938d38829..0055b798d 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -636,10 +636,17 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: Local (rows owned by this process) diagonal of ``op``, one array per block of its domain. Supported are StencilMatrix, StencilDiagonalMatrix, KroneckerStencilMatrix (1d factors), - sums, scalings and block-diagonal BlockLinearOperators of these. + sums, scalings and BlockLinearOperators of these (missing diagonal blocks give zeros). """ if isinstance(op, BlockLinearOperator): - return [a for i in range(op.n_block_rows) for a in _local_diagonal(op[i, i])] + 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): diff --git a/src/struphy/feec/tests/test_preconditioner_transpose.py b/src/struphy/feec/tests/test_preconditioner_transpose.py index 4ca99f937..9d5ff52c3 100644 --- a/src/struphy/feec/tests/test_preconditioner_transpose.py +++ b/src/struphy/feec/tests/test_preconditioner_transpose.py @@ -66,3 +66,38 @@ def recording_init(self, derham, V, W, *args, **kwargs): 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) From 7839a62094fc0ba79e375708914b52550bbf05b5 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 13:56:38 +0200 Subject: [PATCH 14/18] Bump feectools (factor row layout, periodic duplicates); diagonal offset Local diagonal of a KroneckerStencilMatrix: use the global row offset, so that factors may own more rows than the codomain on this process. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- src/struphy/feec/preconditioner.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/feectools b/feectools index c774ecf09..e000fb438 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit c774ecf0985c23d8f6f851fb1606840eb9494726 +Subproject commit e000fb438bc3c03edd22e83bb63e281f79b5bf86 diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 0055b798d..ff503da39 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -656,7 +656,8 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: 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." - off = A.codomain.pads[0] * A.codomain.shifts[0] + # 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:]: From 8f207e467111447423df4839d3d626cf59f67638 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 14:16:40 +0200 Subject: [PATCH 15/18] Fix diagonal placement of process-local 1d factors (small periodic dirs) _process_local_matrix_1d stored entry (i, j) at diagonal (j - i + p) mod n. For periodic directions with 2p + 1 > n (e.g. 1 element, degree 1) this put the main diagonal at another equivalent offset; dot was unaffected, but the diagonal of the Kronecker approximation (used for diagonal scaling since MassMatrixDiagonalPreconditioner no longer assembles the 3d logical mass matrix) was zero, so PCG returned NaN in test_variational_mag_field_evolve. Periodic offsets are now taken nearest to 0; non-periodic ones are not wrapped (the modulo was also wrong there for 2p >= n). Regression test with a (4, 4, 1) grid. CI shard 2 passes locally. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/preconditioner.py | 10 +++++-- .../tests/test_preconditioner_transpose.py | 27 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index ff503da39..f671c39c2 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -982,11 +982,17 @@ def _process_local_matrix_1d( 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, diagonal (j - i + p) mod n + # 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] - M_local._data[rows - s + p, (cols + p - rows) % M_arr.shape[1]] = M_arr[rows, cols] + 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]) diff --git a/src/struphy/feec/tests/test_preconditioner_transpose.py b/src/struphy/feec/tests/test_preconditioner_transpose.py index 9d5ff52c3..f0b49eaf6 100644 --- a/src/struphy/feec/tests/test_preconditioner_transpose.py +++ b/src/struphy/feec/tests/test_preconditioner_transpose.py @@ -101,3 +101,30 @@ def one(e1, e2, e3): _, 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)) + + _, 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"] From 920492790f19e68f26bea88175764cfa1d3fac8a Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 17:09:05 +0200 Subject: [PATCH 16/18] Bump feectools (merge of devel-tiny, version 0.7.0) Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index e000fb438..9502142dd 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit e000fb438bc3c03edd22e83bb63e281f79b5bf86 +Subproject commit 9502142dde2bf8739e7b93c2e972a9d9221a7d13 From 45f2a81cdeaa9fda7a71e204fc4e2e8f375ba944 Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Thu, 8 Oct 2026 18:30:23 +0200 Subject: [PATCH 17/18] update feectools devel-tiny --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 9502142dd..5385af968 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 9502142dde2bf8739e7b93c2e972a9d9221a7d13 +Subproject commit 5385af968237c96f7d3ca643a5da4f4e4e31916d From c95939dbf870a0ca6979feaf98a7f1175df9fbea Mon Sep 17 00:00:00 2001 From: Stefan Possanner Date: Fri, 9 Oct 2026 08:39:11 +0200 Subject: [PATCH 18/18] require feectools 0.7.0 in pyproject --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 88547f445..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.5.0, <=0.5.0", + "feectools>=0.7.0, <=0.7.0", "scipy<=1.18.0", "h5py<=3.16.0", "h5netcdf<=1.8.1",