From 39202a32777c72cbce4fc4b3a41e253555b6bfcb Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 18 Aug 2026 16:46:04 +0200 Subject: [PATCH 1/7] Added DirectSolver(InverseLinearOperator) --- feectools/feec/derivatives.py | 22 ++++- feectools/linalg/direct_solvers.py | 15 ++- feectools/linalg/solvers.py | 147 +++++++++++++++++++++++++++++ feectools/linalg/stencil.py | 13 --- 4 files changed, 174 insertions(+), 23 deletions(-) diff --git a/feectools/feec/derivatives.py b/feectools/feec/derivatives.py index f50ce3c64..6b12a7a43 100644 --- a/feectools/feec/derivatives.py +++ b/feectools/feec/derivatives.py @@ -290,8 +290,13 @@ def tosparse(self, **kwargs): with_pads = kwargs.pop('with_pads', False) - # avoid this case (no pads, but parallel) - assert not (self.domain.parallel and not with_pads) + # avoid this case (no pads, but genuinely decomposed across more than one rank): + # `.parallel` only means "an MPI communicator is attached", true even at 1 rank + # (e.g. under `srun -n 1`), where the no-pads local range already *is* the full + # global range and this restriction does not apply -- so check the rank count + # (`cart.nprocs`) directly rather than `.parallel`. + if self.domain.parallel: + assert with_pads or all(n == 1 for n in self._spaceV.cart.nprocs) # begin with a 1×1 matrix matrix = spa.identity(1, format='coo') @@ -315,13 +320,20 @@ def tosparse(self, **kwargs): directional_matrix = spa.coo_array((codomain_local, domain_local)) else: - maindiag = xp.ones(domain_local) * (-sign) - adddiag = xp.ones(domain_local) * sign + # Plain NumPy, not xp: scipy.sparse.diags is host-only and rejects a + # CuPy array outright (unlike an implicit numpy->cupy conversion, + # cupy->numpy needs an explicit .get()/xp.to_numpy()) -- these + # diagonals are tiny and only ever feed this one-time host-side + # sparse assembly, never a device computation. + import numpy as np + + maindiag = np.ones(domain_local) * (-sign) + adddiag = np.ones(domain_local) * sign # handle special case with not self.domain.parallel and not with_pads and periodic if self.domain.periods[d] and not self.domain.parallel and not with_pads: # then: add element to other side of the array - adddiagcirc = xp.array([sign]) + adddiagcirc = np.array([sign]) offsets = (-codomain_local+1, 0, 1) diags = (adddiagcirc, maindiag, adddiag) else: diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index 1b6096d12..0baa0c8b9 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -234,12 +234,17 @@ def solve(self, rhs, out=None): assert out.shape == rhs.shape assert out.dtype == rhs.dtype - # currently no in-place solve exposed - if array_backend.backend == "numpy": - out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T - else: - rhs_cpu = rhs.get() + # currently no in-place solve exposed. Branch on whether `rhs` itself is a + # device array (not the global `array_backend.backend` flag): the LU + # factorization always lives on the host regardless of backend, and a caller + # may deliberately pass an already-host `rhs`/`out` pair even while the + # active backend is CuPy (see feectools.linalg.solvers.DirectSolver), in + # which case `.get()`-ing a plain NumPy array would fail outright. + if xp.is_gpu(rhs): + rhs_cpu = xp.to_numpy(rhs) result_cpu = self._splu.solve(rhs_cpu.T, trans='T' if transposed else 'N').T out[:] = xp.asarray(result_cpu) + else: + out[:] = self._splu.solve(rhs.T, trans='T' if transposed else 'N').T return out diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index d2e673a5f..2d2650085 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -4,6 +4,7 @@ """ import cunumpy as xp +import numpy as np from math import sqrt, inf from feectools.utilities.utils import is_real @@ -17,6 +18,7 @@ 'inverse', 'ConjugateGradient', 'PConjugateGradient', + 'DirectSolver', 'BiConjugateGradient', 'BiConjugateGradientStabilized', 'PBiConjugateGradientStabilized', @@ -60,6 +62,7 @@ def inverse(A, solver, **kwargs): solvers_dict = { 'cg' : ConjugateGradient, 'pcg' : PConjugateGradient, + 'direct' : DirectSolver, 'bicg' : BiConjugateGradient, 'bicgstab' : BiConjugateGradientStabilized, 'pbicgstab': PBiConjugateGradientStabilized, @@ -411,6 +414,150 @@ def solve(self, b, out=None): def dot(self, b, out=None): return self.solve(b, out=out) +#=============================================================================== +class DirectSolver(InverseLinearOperator): + """ + Exact sparse-direct solve, for linear systems whose left-hand-side operator A does + not actually change across repeated `solve()` calls -- e.g. a time-independent field + operator solved once per time step with only the right-hand side changing (see + `struphy.propagators.implicit_diffusion.ImplicitDiffusion`, whose LHS is constant + whenever `divide_by_dt=False`). A single sparse LU factorization + (`feectools.linalg.direct_solvers.SparseSolver`) then serves every call, instead of + an iterative method repeating (in the worst case, all the way to `maxiter`) every + single call. + + The factorization is built lazily, on the first `solve()` call, and then reused by + every later call without ever re-examining `A` again -- including through a `.linop` + reassignment, e.g. `ImplicitDiffusion.__call__` unconditionally reassigns `.linop` to + a freshly *built* operator every step, regardless of whether its *values* actually + changed. This is a deliberate, cheap-by-construction design, not a value comparison: + `A.tosparse()` is not assumed to be cheap (composed operators can include a + basis-vector sweep, see e.g. `AverageOperator.tosparse`/`BoundaryOperator.tosparse` + in `struphy.feec.mass`/`struphy.feec.linear_operators`), so re-deriving and comparing + it on every call would undo most of the point of factorizing once. The caller is + therefore responsible for knowing that `A`'s *values* are actually constant across + calls (true whenever `ImplicitDiffusion.divide_by_dt=False`, since neither `epsilon` + nor `Z` change during a run); call `invalidate()` explicitly if `A` does change and + the factorization must be rebuilt on the next `solve()`. + + Only supports a serial (non-MPI-parallel) `A`/domain/codomain: a distributed + sparse-direct solve would need its own implementation, which + `feectools.linalg.direct_solvers.SparseSolver` (and therefore this class) does not + attempt. + + Parameters + ---------- + A : feectools.linalg.basic.LinearOperator + Left-hand-side matrix A of the linear system. Must support `.tosparse()` and + have a serial (non-parallel) domain/codomain. + + pc, tol, maxiter, verbose : ignored + Accepted only so this class is a drop-in alternative to the iterative solvers + behind the same `solvers.inverse(A, solver, ...)` call site; a direct solve has + no preconditioner, iteration count, or convergence tolerance. + + x0 : feectools.linalg.basic.Vector, optional + Ignored for solving (a direct solve needs no initial guess); if `recycle=True`, + still receives a copy of each solution, for interface consistency with the + iterative solvers (some callers read `x0` back out directly). + + recycle : bool + If True, a copy of the output is stored in x0, as the iterative solvers do. + """ + + def __init__(self, A, *, pc=None, x0=None, tol=None, maxiter=None, verbose=False, recycle=False): + + self._options = {"x0": x0, "pc": pc, "tol": tol, "maxiter": maxiter, "verbose": verbose, "recycle": recycle} + + super().__init__(A, **self._options) + + # `.parallel` only means "an MPI communicator is attached", true even at 1 rank + # (e.g. under `srun -n 1`); what actually matters for a local sparse-direct + # solve is the rank *count*, so check `cart.nprocs` (per-direction process + # counts) directly rather than `.parallel`. + if self.domain.parallel: + assert all(n == 1 for n in self.domain.cart.nprocs), \ + "DirectSolver only supports a single MPI rank; SparseSolver has no distributed factorization." + + self._sparse_solver = None + self._info = None + + def _check_options(self, **kwargs): + # tol/maxiter/verbose are meaningless for a direct solve (see class docstring); + # only x0, if given, is worth the base class's type/space check. + x0 = kwargs.get("x0") + if x0 is not None: + assert isinstance(x0, Vector), "x0 must be a Vector or None" + assert x0.space == self.codomain, "x0 belongs to the wrong VectorSpace" + + def invalidate(self): + """Force the next `solve()` call to rebuild the factorization from `A`. + + Call this after actually changing `A` (in place, or via the `.linop` setter with + a numerically different operator) -- see the class docstring for why this is not + detected automatically. + """ + self._sparse_solver = None + + def _ensure_factorized(self): + if self._sparse_solver is None: + from feectools.linalg.direct_solvers import SparseSolver + + self._sparse_solver = SparseSolver(self._A.tosparse().tocsc()) + + def solve(self, b, out=None): + """ + Solve A x = b exactly via the cached sparse LU factorization. + + Parameters + ---------- + b : feectools.linalg.stencil.StencilVector + Right-hand-side vector of the linear system. + + out : feectools.linalg.basic.Vector | NoneType + The output vector, or None (optional). + + Returns + ------- + x : feectools.linalg.basic.Vector + The exact (up to factorization round-off) solution of the linear system. + """ + assert isinstance(b, Vector) + assert b.space is self.domain + + self._ensure_factorized() + + # SparseSolver's factorization always lives on the host (scipy splu); the + # host round trip here is one flat vector of the field-solve's DOF count, not + # the particle arrays, so it is cheap relative to the iterations it replaces. + b_flat = xp.to_numpy(b.toarray()) + x_flat = np.empty_like(b_flat) + self._sparse_solver.solve(b_flat, out=x_flat) + + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + + # Same local/no-pad interior slice StencilVector.toarray_local() reads from, + # see feectools.linalg.stencil.StencilVector.toarray_local. + idx = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(out.pads, out.space.shifts) + ) + out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) + + self._info = {'niter': 1, 'success': True, 'res_norm': 0.0} + + if self._options.get("recycle") and self._options.get("x0") is not None: + out.copy(out=self._options["x0"]) + + return out + + def dot(self, b, out=None): + return self.solve(b, out=out) + #=============================================================================== class BiConjugateGradient(InverseLinearOperator): """ diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py index 4848595c3..2dd9f9148 100644 --- a/feectools/linalg/stencil.py +++ b/feectools/linalg/stencil.py @@ -1749,20 +1749,7 @@ def _tocoo_no_pads(self , order='C'): data[:ind] = cp.asarray(data_np[:ind]) rows[:ind] = cp.asarray(rows_np[:ind]) cols[:ind] = cp.asarray(cols_np[:ind]) - nrl = [_np.int64(e-s+1) for s,e in zip(self.codomain.starts, self.codomain.ends)] - ncl = [_np.int64(i) for i in self._data.shape[nd:]] - ss = [_np.int64(i) for i in ss] - nr = [_np.int64(i) for i in nr] - nc = [_np.int64(i) for i in nc] - dm = [_np.int64(i) for i in dm] - cm = [_np.int64(i) for i in cm] - cpads = [_np.int64(i) for i in cpads] - pp = [_np.int64(i) for i in pp] - stencil2coo = kernels['stencil2coo'][order][nd] - ind = stencil2coo(self._data, data, rows, cols, *nrl, *ncl, *ss, *nr, *nc, *dm, *cm, *cpads, *pp) - - if array_backend.backend == "cupy": M = coo_matrix( (data[:ind].get(), (rows[:ind].get(), cols[:ind].get())), From e2d6ce6eec784a87eb3d995ada8516cfc0438dfe Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Tue, 18 Aug 2026 16:47:25 +0200 Subject: [PATCH 2/7] Updated version number --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 0794c0d0f..ac462afe6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "feectools" -version = "0.1.10" +version = "0.1.11" description = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies." readme = "README.md" requires-python = ">= 3.10" From da50cf7b98f25d1e2fe7d1d47b912a9e59be10f0 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 13:10:08 +0200 Subject: [PATCH 3/7] Fixes for nprocs > 1 --- feectools/linalg/solvers.py | 110 ++++++++++---- feectools/linalg/tests/test_solvers.py | 61 +++++++- feectools/linalg/utilities.py | 190 ++++++++++++++++++++++++- 3 files changed, 333 insertions(+), 28 deletions(-) diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index 2d2650085..883c05e3e 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -7,6 +7,8 @@ import numpy as np from math import sqrt, inf +from feectools.ddm.mpi import MockComm +from feectools.ddm.mpi import mpi as MPI from feectools.utilities.utils import is_real from feectools.linalg.utilities import _sym_ortho from feectools.linalg.basic import (Vector, LinearOperator, @@ -440,16 +442,25 @@ class DirectSolver(InverseLinearOperator): nor `Z` change during a run); call `invalidate()` explicitly if `A` does change and the factorization must be rebuilt on the next `solve()`. - Only supports a serial (non-MPI-parallel) `A`/domain/codomain: a distributed - sparse-direct solve would need its own implementation, which - `feectools.linalg.direct_solvers.SparseSolver` (and therefore this class) does not - attempt. + At `nprocs > 1`, this factorizes a *replicated* copy of the full global matrix on + every rank (assembled once via `feectools.linalg.utilities.tosparse_via_matvec`, + which applies `A` to every global unit vector through its own -- already + MPI-correct -- `.dot()`, since `A.tosparse()` itself is only correct in serial for + several composed/derivative operators), rather than attempting an actual + distributed factorization. Every rank redundantly solves the same full system and + keeps only its own slice of the result -- correct and simple, but each rank does + `O(A.domain.dimension)` work per solve instead of `O(A.domain.dimension / nprocs)`, + and the one-time assembly is `O(A.domain.dimension)` *collective* `.dot()` calls that + do not get cheaper with more ranks. This trade only makes sense for problems small + enough that `splu` and this redundant work stay cheap (e.g. the few-thousand-DOF + field solves this class targets); a genuinely distributed sparse-direct solve (e.g. + via PETSc/MUMPS) would need its own implementation. Parameters ---------- A : feectools.linalg.basic.LinearOperator - Left-hand-side matrix A of the linear system. Must support `.tosparse()` and - have a serial (non-parallel) domain/codomain. + Left-hand-side matrix A of the linear system. Must support `.tosparse()` (serial) + or `.dot()` (parallel, via `tosparse_via_matvec`). pc, tol, maxiter, verbose : ignored Accepted only so this class is a drop-in alternative to the iterative solvers @@ -472,12 +483,12 @@ def __init__(self, A, *, pc=None, x0=None, tol=None, maxiter=None, verbose=False super().__init__(A, **self._options) # `.parallel` only means "an MPI communicator is attached", true even at 1 rank - # (e.g. under `srun -n 1`); what actually matters for a local sparse-direct - # solve is the rank *count*, so check `cart.nprocs` (per-direction process - # counts) directly rather than `.parallel`. - if self.domain.parallel: - assert all(n == 1 for n in self.domain.cart.nprocs), \ - "DirectSolver only supports a single MPI rank; SparseSolver has no distributed factorization." + # (e.g. under `srun -n 1`), where the serial `.tosparse()` path is already + # correct (local range == global range) and faster than the replicated-assembly + # path -- so check the rank *count* (`cart.nprocs`) directly. + cart = self.domain.spaces[0].cart if isinstance(self.domain, BlockVectorSpace) else self.domain.cart + self._parallel = self.domain.parallel and any(n != 1 for n in cart.nprocs) + self._comm = cart.comm if self._parallel else None self._sparse_solver = None self._info = None @@ -503,7 +514,30 @@ def _ensure_factorized(self): if self._sparse_solver is None: from feectools.linalg.direct_solvers import SparseSolver - self._sparse_solver = SparseSolver(self._A.tosparse().tocsc()) + if self._parallel: + from feectools.linalg.utilities import tosparse_via_matvec + + mat = tosparse_via_matvec(self._A, format="csr") + else: + mat = self._A.tosparse().tocsr() + + # `A` can be exactly singular at essential-BC-masked DOFs: an operator + # built through a BoundaryOperator zero-masks both the input and output at + # those rows by design (struphy.feec.linear_operators.BoundaryOperator.dot, + # via apply_essential_bc_to_array) -- fine for an iterative solver, which + # never inverts A directly, as long as `b` is masked the same way (true for + # every caller here: e.g. ImplicitDiffusion.__call__ builds `rhs` via the + # same BoundaryOperator-wrapped `.dot()`, so `b` is already 0 at these rows + # too). A direct factorization needs those rows regularized to identity so + # `x = 1^{-1} * 0 = 0` comes out right there instead of `splu` raising + # "Factor is exactly singular" -- a zero row is unsolvable on its own even + # though the underlying (masked) system is perfectly well posed. + zero_rows = np.flatnonzero(mat.getnnz(axis=1) == 0) + if zero_rows.size: + mat = mat.tolil() + mat[zero_rows, zero_rows] = 1.0 + + self._sparse_solver = SparseSolver(mat.tocsc()) def solve(self, b, out=None): """ @@ -531,22 +565,48 @@ def solve(self, b, out=None): # host round trip here is one flat vector of the field-solve's DOF count, not # the particle arrays, so it is cheap relative to the iterations it replaces. b_flat = xp.to_numpy(b.toarray()) + + if self._parallel: + # `b.toarray()` in parallel already returns the full global-shape array with + # only this rank's own (disjoint) entries filled in -- see + # `StencilVector._toarray_parallel_no_pads` -- so summing every rank's copy + # assembles the true global right-hand side. + if isinstance(self._comm, MockComm): + b_global = b_flat + else: + b_global = np.empty_like(b_flat) + self._comm.Allreduce(b_flat, b_global, op=MPI.SUM) + b_flat = b_global + x_flat = np.empty_like(b_flat) self._sparse_solver.solve(b_flat, out=x_flat) - if out is None: - out = self.codomain.zeros() + if self._parallel: + from feectools.linalg.utilities import array_to_psydac + + # x_flat should be a numpy array since SparseSolver's factorization + # is on host + x_vec = array_to_psydac(xp.asarray(x_flat), self.codomain) + if out is None: + out = x_vec + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + x_vec.copy(out=out) else: - assert isinstance(out, Vector) - assert out.space is self.codomain - - # Same local/no-pad interior slice StencilVector.toarray_local() reads from, - # see feectools.linalg.stencil.StencilVector.toarray_local. - idx = tuple( - slice(m * p, -m * p) if p != 0 else slice(0, None) - for p, m in zip(out.pads, out.space.shifts) - ) - out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, Vector) + assert out.space is self.codomain + + # Same local/no-pad interior slice StencilVector.toarray_local() reads from, + # see feectools.linalg.stencil.StencilVector.toarray_local. + idx = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(out.pads, out.space.shifts) + ) + out._data[idx] = xp.asarray(x_flat.reshape(out._data[idx].shape, order='C')) self._info = {'niter': 1, 'success': True, 'res_norm': 0.0} diff --git a/feectools/linalg/tests/test_solvers.py b/feectools/linalg/tests/test_solvers.py index e877888f4..02bc39159 100644 --- a/feectools/linalg/tests/test_solvers.py +++ b/feectools/linalg/tests/test_solvers.py @@ -1,10 +1,11 @@ import cunumpy as xp import pytest -from feectools.linalg.solvers import inverse +from feectools.linalg.solvers import inverse, DirectSolver from feectools.linalg.stencil import StencilVectorSpace, StencilMatrix, StencilVector from feectools.linalg.basic import LinearSolver from feectools.ddm.cart import DomainDecomposition, CartDecomposition +from feectools.ddm.mpi import mpi as MPI def define_data_hermitian(n, p, dtype=float): @@ -204,6 +205,64 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False): assert errh_norm < tol assert solver == 'pcg' or errc_norm < tol +#=============================================================================== +def _compute_global_starts_ends(domain_decomposition, npts): + # Same as feectools.linalg.tests.test_block.compute_global_starts_ends. + global_starts = [None] * len(npts) + global_ends = [None] * len(npts) + for axis in range(len(npts)): + ee = domain_decomposition.global_element_ends[axis] + global_ends[axis] = ee.copy() + global_ends[axis][-1] = npts[axis] - 1 + global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) + return global_starts, global_ends + + +@pytest.mark.parametrize('n1', [8, 16]) +@pytest.mark.parametrize('n2', [8, 12]) +@pytest.mark.parametrize('p1', [1, 2]) +@pytest.mark.parallel +def test_direct_solver_parallel(n1, n2, p1, verbose=False): + """`DirectSolver` at nprocs > 1 must recover the exact solution, same as serial.""" + p2 = 1 + + comm = MPI.COMM_WORLD + D = DomainDecomposition([n1, n2], periods=[False, False], comm=comm) + npts = [n1, n2] + global_starts, global_ends = _compute_global_starts_ends(D, npts) + cart = CartDecomposition(D, npts, global_starts, global_ends, pads=[p1, p2], shifts=[1, 1]) + + V = StencilVectorSpace(cart, dtype=float) + A = StencilMatrix(V, V) + + # Diagonally dominant (hence nonsingular) stencil: -1 on every off-diagonal, enough + # on the main diagonal to dominate the row sum -- same style as define_data_hermitian + # above, extended to 2D. + n_offdiag = (2 * p1 + 1) * (2 * p2 + 1) - 1 + for k1 in range(-p1, p1 + 1): + for k2 in range(-p2, p2 + 1): + A[:, :, k1, k2] = 0.0 if (k1 == 0 and k2 == 0) else -1.0 + A[:, :, 0, 0] = n_offdiag + 1.0 + A.remove_spurious_entries() + + s1, s2 = V.starts + e1, e2 = V.ends + xe = StencilVector(V) + for i1 in range(s1, e1 + 1): + for i2 in range(s2, e2 + 1): + xe[i1, i2] = xp.random.random() + xe.update_ghost_regions() + + be = A @ xe + + solv = DirectSolver(A) + x = solv.solve(be) + + err_norm = xp.linalg.norm((x - xe).toarray()) + if verbose: + print(f"n1={n1} n2={n2} p1={p1} p2={p2} nprocs={comm.Get_size()} err_norm={err_norm:.2e}") + assert err_norm < 1e-9 + # =============================================================================== # SCRIPT FUNCTIONALITY #=============================================================================== diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 42d5636b7..e7d47a233 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -1,8 +1,14 @@ # coding: utf-8 -import cunumpy as xp +import itertools from math import sqrt +import cunumpy as xp +import numpy as np +from scipy import sparse + +from feectools.ddm.mpi import MockComm +from feectools.ddm.mpi import mpi as MPI from feectools.linalg.basic import Vector from feectools.linalg.stencil import StencilVector, StencilVectorSpace from feectools.linalg.block import BlockVector, BlockVectorSpace @@ -11,6 +17,7 @@ __all__ = ( 'array_to_psydac', 'petsc_to_psydac', + 'tosparse_via_matvec', '_sym_ortho', ) @@ -72,7 +79,186 @@ def _array_to_psydac_recursive(x, u): else: raise NotImplementedError(f'Can only handle StencilVector or BlockVector spaces, got {type(V)} instead') - + +#============================================================================== +def tosparse_via_matvec(op, format="csc"): + """ + Assemble the full global sparse matrix of a `LinearOperator` by applying it to every + global unit vector via `.dot()`, rather than via `.tosparse()`. + + Every operator's `.dot()` is already exercised (and therefore correct, including + cross-rank ghost/boundary coupling) every time it is actually used, unlike + `.tosparse()`, which several composed/derivative operators only implement correctly + in serial (see e.g. `feectools.feec.derivatives.DirectionalDerivativeOperator.tosparse`). + This is a port of `struphy.feec.linear_operators.LinOpWithTransp.toarray_struphy`'s + `is_sparse=True` branch into feectools (which `DirectSolver` -- the caller this exists + for -- must not import struphy from): same Allgather-starts/ends plus + unit-vector-`dot()` plus gather/broadcast-triplets algorithm, so every rank ends up + with an identical copy of the full global matrix (a "replicated" assembly, not a + distributed one -- deliberate, see `feectools.linalg.solvers.DirectSolver`). + + Cost: O(N) collective `.dot()` calls, N = `op.domain.dimension` -- does not shrink + with rank count (every call needs every rank's participation), so this is only + appropriate as a one-time, cached setup cost, not something to call every step. + + Parameters + ---------- + op : feectools.linalg.basic.LinearOperator + Operator to assemble. `op.domain`/`op.codomain` must each be a + `StencilVectorSpace` or `BlockVectorSpace`. + + format : str + scipy.sparse matrix format of the result ("csr", "csc", "coo", ...). + + Returns + ------- + out : scipy.sparse matrix + The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on + every rank. + """ + v = op.domain.zeros() + tmp2 = op.codomain.zeros() + + if isinstance(op.domain, BlockVectorSpace): + comm = op.domain.spaces[0].cart.comm + elif isinstance(op.domain, StencilVectorSpace): + comm = op.domain.cart.comm + else: + raise NotImplementedError( + f'tosparse_via_matvec only supports StencilVectorSpace/BlockVectorSpace domains, got {type(op.domain)}', + ) + + if comm is None or isinstance(comm, MockComm): + rank = 0 + size = 1 + else: + rank = comm.Get_rank() + size = comm.Get_size() + + numrows = op.codomain.dimension + numcols = op.domain.dimension + data, row, col = [], [], [] + + if isinstance(op.domain, BlockVectorSpace): + starts = [vi.starts for vi in v] + ends = [vi.ends for vi in v] + npts = [sp.npts for sp in op.domain.spaces] + nsp = len(op.domain.spaces) + ndim = [sp.ndim for sp in op.domain.spaces] + + # Plain NumPy throughout: this is tiny host-side index bookkeeping (rank + # starts/ends, a running column count), never device compute -- `xp.array` + # under the CuPy backend would produce 0-d CuPy scalars that `range()` (and + # plain Python int arithmetic below) cannot consume, the same class of + # NumPy-vs-CuPy scalar-typing trap documented for `AdhocTorus`/`xp.sqrt`. + startsarr = np.array([starts[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + allstarts = np.empty(size * len(startsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allstarts = startsarr + else: + comm.Allgather(startsarr, allstarts) + allstarts = allstarts.reshape((size, len(startsarr))) + + endsarr = np.array([ends[i][j] for i in range(nsp) for j in range(ndim[i])], dtype=int) + allends = np.empty(size * len(endsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allends = endsarr + else: + comm.Allgather(endsarr, allends) + allends = allends.reshape((size, len(endsarr))) + + for currentrank in range(size): + spoint = 0 + npredim = 0 + for h in range(nsp): + iterables = [ + range(int(allstarts[currentrank][i + npredim]), int(allends[currentrank][i + npredim]) + 1) + for i in range(ndim[h]) + ] + for i in itertools.product(*iterables): + if rank == currentrank: + v[h][i] = 1.0 + v[h].update_ghost_regions() + tmp2 *= 0.0 + op.dot(v, out=tmp2) + c = spoint + int(np.ravel_multi_index(i, npts[h])) + aux = xp.to_numpy(tmp2.toarray()) + for r in np.nonzero(aux)[0]: + data.append(aux[r]) + col.append(c) + row.append(int(r)) + if rank == currentrank: + v[h][i] = 0.0 + v[h].update_ghost_regions() + cumulative = 1 + for i in range(ndim[h]): + cumulative *= npts[h][i] + spoint += cumulative + npredim += ndim[h] + + else: + starts = v.starts + ends = v.ends + npts = op.domain.npts + ndim = op.domain.ndim + + # Plain NumPy, same reasoning as the BlockVectorSpace branch above. + startsarr = np.array([starts[j] for j in range(ndim)], dtype=int) + allstarts = np.empty(size * len(startsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allstarts = startsarr + else: + comm.Allgather(startsarr, allstarts) + allstarts = allstarts.reshape((size, len(startsarr))) + + endsarr = np.array([ends[j] for j in range(ndim)], dtype=int) + allends = np.empty(size * len(endsarr), dtype=int) + if comm is None or isinstance(comm, MockComm): + allends = endsarr + else: + comm.Allgather(endsarr, allends) + allends = allends.reshape((size, len(endsarr))) + + for currentrank in range(size): + iterables = [ + range(int(allstarts[currentrank][i]), int(allends[currentrank][i]) + 1) for i in range(ndim) + ] + for i in itertools.product(*iterables): + if rank == currentrank: + v[i] = 1.0 + v.update_ghost_regions() + op.dot(v, out=tmp2) + c = int(np.ravel_multi_index(i, npts)) + aux = xp.to_numpy(tmp2.toarray()) + for r in np.nonzero(aux)[0]: + data.append(aux[r]) + col.append(c) + row.append(int(r)) + if rank == currentrank: + v[i] = 0.0 + v.update_ghost_regions() + + if comm is None or isinstance(comm, MockComm): + all_rows, all_cols, all_data = row, col, data + else: + gathered_rows = comm.gather(row, root=0) + gathered_cols = comm.gather(col, root=0) + gathered_data = comm.gather(data, root=0) + if rank == 0: + all_rows = [item for sublist in gathered_rows for item in sublist] + all_cols = [item for sublist in gathered_cols for item in sublist] + all_data = [item for sublist in gathered_data for item in sublist] + comm.bcast(all_rows, root=0) + comm.bcast(all_cols, root=0) + comm.bcast(all_data, root=0) + else: + all_rows = comm.bcast(None, root=0) + all_cols = comm.bcast(None, root=0) + all_data = comm.bcast(None, root=0) + + mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) + return mat.asformat(format) + #============================================================================== def petsc_to_psydac(x, Xh, out=None): """ From d972105bb68b9511b5dcc00f7c3226497aef542e Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 19:06:51 +0200 Subject: [PATCH 4/7] Added some utilities --- feectools/linalg/solvers.py | 16 +- feectools/linalg/tests/test_utilities.py | 428 +++++++++++++++++++++ feectools/linalg/utilities.py | 453 +++++++++++++++++++++++ 3 files changed, 894 insertions(+), 3 deletions(-) create mode 100644 feectools/linalg/tests/test_utilities.py diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py index 883c05e3e..32a8821ee 100644 --- a/feectools/linalg/solvers.py +++ b/feectools/linalg/solvers.py @@ -515,9 +515,19 @@ def _ensure_factorized(self): from feectools.linalg.direct_solvers import SparseSolver if self._parallel: - from feectools.linalg.utilities import tosparse_via_matvec - - mat = tosparse_via_matvec(self._A, format="csr") + from feectools.linalg.utilities import FastAssemblyUnavailable, parallel_tosparse, tosparse_via_matvec + + # `parallel_tosparse` assembles in O(1) collective rounds (one per leaf + # operator) instead of `tosparse_via_matvec`'s O(A.domain.dimension) + # rounds (one per global DOF) -- a difference of several orders of + # magnitude for a field-solve-sized system (see its docstring for how). + # It only recognizes a subset of operator shapes, self-verified against + # `A`'s own `.dot()`; fall back to the always-correct (if much slower) + # sweep when it can't. + try: + mat = parallel_tosparse(self._A, self._comm, format="csr") + except FastAssemblyUnavailable: + mat = tosparse_via_matvec(self._A, format="csr") else: mat = self._A.tosparse().tocsr() diff --git a/feectools/linalg/tests/test_utilities.py b/feectools/linalg/tests/test_utilities.py new file mode 100644 index 000000000..94e511c5b --- /dev/null +++ b/feectools/linalg/tests/test_utilities.py @@ -0,0 +1,428 @@ +import cunumpy as xp +import numpy as np +import pytest +from scipy import sparse + +from feectools.ddm.cart import CartDecomposition, DomainDecomposition +from feectools.ddm.mpi import mpi as MPI +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ComposedLinearOperator, IdentityOperator, LinearOperator, ScaledLinearOperator, SumLinearOperator +from feectools.linalg.block import BlockLinearOperator, BlockVectorSpace +from feectools.linalg.stencil import StencilMatrix, StencilVectorSpace +from feectools.linalg.utilities import ( + FastAssemblyUnavailable, + _get_entry, + _local_flat_entries, + _set_entry, + parallel_tosparse, + tosparse_via_matvec, +) + + +def compute_global_starts_ends(domain_decomposition, npts): + global_starts = [None] * len(npts) + global_ends = [None] * len(npts) + for axis in range(len(npts)): + ee = domain_decomposition.global_element_ends[axis] + global_ends[axis] = ee.copy() + global_ends[axis][-1] = npts[axis] - 1 + global_starts[axis] = xp.array([0] + (global_ends[axis][:-1] + 1).tolist()) + return global_starts, global_ends + + +def make_space(n1, n2, p1, p2, comm, periodic=True): + D = DomainDecomposition([n1, n2], periods=[periodic, False], comm=comm) + npts = [n1, n2] + gs, ge = compute_global_starts_ends(D, npts) + cart = CartDecomposition(D, npts, gs, ge, pads=[p1, p2], shifts=[1, 1]) + return StencilVectorSpace(cart, dtype=float) + + +def make_stencil_matrix(V, p1, p2, scale=1.0): + A = StencilMatrix(V, V) + n_offdiag = (2 * p1 + 1) * (2 * p2 + 1) - 1 + for k1 in range(-p1, p1 + 1): + for k2 in range(-p2, p2 + 1): + A[:, :, k1, k2] = 0.0 if (k1 == 0 and k2 == 0) else -1.0 * scale + A[:, :, 0, 0] = (n_offdiag + 1.0) * scale + A.remove_spurious_entries() + return A + + +class FakeBoundaryOperator(LinearOperator): + """Mimics struphy.feec.linear_operators.BoundaryOperator well enough to exercise + `parallel_tosparse`'s diagonal-probe fallback: a diagonal 0/1 mask (zeroes the + first local DOF of each block on each rank), with a `.tosparse()` that is only + valid in serial (returns a small *locally*-indexed diagonal, exactly as struphy's + real implementation does -- see `struphy.feec.linear_operators. + BoundaryOperator.tosparse`) -- so the fast path must detect the mismatch via + validation and retry with the probe strategy instead of trusting `.tosparse()` + blindly. Works for both a plain StencilVectorSpace (V) and a BlockVectorSpace + (e.g. Hcurl, the codomain a real `grad` is wrapped in) domain, matching either + shape BoundaryOperator actually appears in. + """ + + def __init__(self, V): + self._V = V + self._entries = _local_flat_entries(V) + + @property + def domain(self): + return self._V + + @property + def codomain(self): + return self._V + + @property + def dtype(self): + return float + + def tosparse(self): + # Deliberately wrong at nprocs > 1: local indices, not global ones. + n = len(self._entries) + diag = np.ones(n) + diag[0] = 0.0 + return sparse.diags(diag, format="csr") + + def toarray(self): + return self.tosparse().toarray() + + def transpose(self, conjugate=False): + return self + + def dot(self, v, out=None): + if out is None: + out = self.codomain.zeros() + else: + out *= 0.0 + for k, (setter, _) in enumerate(self._entries): + _set_entry(out, setter, 0.0 if k == 0 else _get_entry(v, setter)) + out.update_ghost_regions() + return out + + +class UnsupportedOperator(LinearOperator): + """A leaf `parallel_tosparse` cannot possibly handle: no `.tosparse()`, and + domain is not codomain (so the diagonal-probe strategy doesn't apply either).""" + + def __init__(self, V_in, V_out): + self._Vin = V_in + self._Vout = V_out + + @property + def domain(self): + return self._Vin + + @property + def codomain(self): + return self._Vout + + @property + def dtype(self): + return float + + def tosparse(self): + raise NotImplementedError + + def toarray(self): + raise NotImplementedError + + def transpose(self, conjugate=False): + return UnsupportedOperator(self._Vout, self._Vin) + + def dot(self, v, out=None): + if out is None: + out = self.codomain.zeros() + return out + + +class FakeWeightedMassOperator(LinearOperator): + """Mimics struphy.feec.mass.WeightedMassOperator closely enough to exercise + `parallel_tosparse`'s duck-typed `._mat`-composition unwrap: wraps an inner + `._mat` behind identity extraction ops (trivial, as struphy's own serial + `.tosparse()` requires) but *non-trivial* boundary ops (masks the first and last + local DOF on each rank) -- i.e. `.dot()` is genuinely not the same as `._mat.dot()` + alone, which is exactly the bug this class was written to catch (found via a real + Struphy run, not anticipated up front -- see the git history of + `parallel_tosparse`'s `._mat`-unwrap branch). + """ + + def __init__(self, V, mat, mask_first=True, mask_last=True): + self._V = V + self._mat = mat + self._V_extraction_op = IdentityOperator(V) + self._W_extraction_op = IdentityOperator(V) + self._V_boundary_op = _MaskBoundaryOperator(V, mask_first, mask_last) + self._W_boundary_op = _MaskBoundaryOperator(V, mask_first, mask_last) + self._transposed = False + + @property + def domain(self): + return self._V + + @property + def codomain(self): + return self._V + + @property + def dtype(self): + return float + + def tosparse(self): + # Deliberately unusable at nprocs > 1, exactly like struphy's real + # WeightedMassOperator.tosparse() when boundary masking is actually active + # (it asserts outright there); here it just raises, which + # `parallel_tosparse` must also handle gracefully. + raise NotImplementedError + + def toarray(self): + raise NotImplementedError + + def transpose(self, conjugate=False): + raise NotImplementedError + + def dot(self, v, out=None): + tmp = self._V_boundary_op.transpose().dot(v) + tmp = self._V_extraction_op.transpose().dot(tmp) + tmp = self._mat.dot(tmp) + tmp = self._W_extraction_op.dot(tmp) + return self._W_boundary_op.dot(tmp, out=out) + + +class _MaskBoundaryOperator(LinearOperator): + """A minimal BoundaryOperator stand-in: zeroes the first and/or last local DOF on + each rank (self-adjoint, so `.transpose()` returns itself).""" + + def __init__(self, V, mask_first, mask_last): + self._V = V + self._entries = _local_flat_entries(V) + self._mask_first = mask_first + self._mask_last = mask_last + + @property + def domain(self): + return self._V + + @property + def codomain(self): + return self._V + + @property + def dtype(self): + return float + + def tosparse(self): + raise NotImplementedError + + def toarray(self): + raise NotImplementedError + + def transpose(self, conjugate=False): + return self + + def dot(self, v, out=None): + if out is None: + out = self.codomain.zeros() + else: + out *= 0.0 + n = len(self._entries) + for k, (setter, _) in enumerate(self._entries): + masked = (self._mask_first and k == 0) or (self._mask_last and k == n - 1) + _set_entry(out, setter, 0.0 if masked else _get_entry(v, setter)) + out.update_ghost_regions() + return out + + +@pytest.mark.parametrize('n1', [8, 16]) +@pytest.mark.parametrize('p1', [1, 2]) +@pytest.mark.parallel +def test_parallel_tosparse_matches_matvec_stencil_matrix(n1, p1, verbose=False): + """A plain StencilMatrix (no wrapper) must assemble identically via the fast + (`parallel_tosparse`) and slow (`tosparse_via_matvec`) paths.""" + n2, p2 = 8, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm) + A = make_stencil_matrix(V, p1, p2) + + fast = parallel_tosparse(A, comm) + slow = tosparse_via_matvec(A, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"n1={n1} p1={p1} nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-10 + + +@pytest.mark.parallel +def test_parallel_tosparse_composed_and_sum(verbose=False): + """A Sum-of-Scaled-and-Composed operator tree (the shape ImplicitDiffusion's + left-hand side actually takes: sigma*M + G^T @ D @ G) must also match the slow + reference path.""" + n1, n2, p1, p2 = 10, 8, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm) + M = make_stencil_matrix(V, p1, p2, scale=1.0) + G = make_stencil_matrix(V, p1, p2, scale=0.5) + D = make_stencil_matrix(V, p1, p2, scale=2.0) + + composed = ComposedLinearOperator(V, V, G, D, G) + scaled = ScaledLinearOperator(V, V, c=3.0, A=M) + total = SumLinearOperator(V, V, scaled, composed) + + fast = parallel_tosparse(total, comm) + slow = tosparse_via_matvec(total, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_block_vector_space(verbose=False): + """A BlockLinearOperator (grad's actual shape, e.g. H1 -> Hcurl's 3 stacked + components) must also assemble identically via the fast and slow paths.""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + W = BlockVectorSpace(V, V) + B = BlockLinearOperator(W, W) + B[0, 0] = make_stencil_matrix(V, p1, p2, scale=1.0) + B[1, 1] = make_stencil_matrix(V, p1, p2, scale=2.0) + B[0, 1] = make_stencil_matrix(V, p1, p2, scale=0.3) + + fast = parallel_tosparse(B, comm) + slow = tosparse_via_matvec(B, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_diagonal_probe_fallback_block_vector_space(verbose=False): + """The same diagonal-mask-on-both-sides shape as + `test_parallel_tosparse_diagonal_probe_fallback`, but over a BlockVectorSpace -- + the actual shape `BoundaryOperator ∘ grad ∘ BoundaryOperator` takes in the real + Poisson benchmark (grad's codomain, Hcurl, has 3 stacked components).""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + W = BlockVectorSpace(V, V) + B = BlockLinearOperator(W, W) + B[0, 0] = make_stencil_matrix(V, p1, p2, scale=1.0) + B[1, 1] = make_stencil_matrix(V, p1, p2, scale=2.0) + mask = FakeBoundaryOperator(W) + + total = ComposedLinearOperator(W, W, mask, B, mask) + + fast = parallel_tosparse(total, comm) + slow = tosparse_via_matvec(total, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_diagonal_probe_fallback(verbose=False): + """A BoundaryOperator-like diagonal mask, wired in on both sides of a + StencilMatrix (the actual shape struphy's BC-wrapped operators take), must be + correctly recovered by the diagonal-probe fallback -- not silently misassembled + from its serial-only `.tosparse()`.""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + A = make_stencil_matrix(V, p1, p2) + mask = FakeBoundaryOperator(V) + + total = ComposedLinearOperator(V, V, mask, A, mask) + + fast = parallel_tosparse(total, comm) + slow = tosparse_via_matvec(total, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 + + +@pytest.mark.parallel +def test_parallel_tosparse_raises_for_unsupported_operator(verbose=False): + """An operator this module genuinely cannot handle (no `.tosparse()`, not + diagonal-shaped) must raise `FastAssemblyUnavailable` -- identically on every + rank, so callers can fall back to `tosparse_via_matvec` without any risk of a + partial/divergent collective-call sequence.""" + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + Vin = make_space(n1, n2, p1, p2, comm, periodic=False) + Vout = make_space(n1, n2, p1, p2, comm, periodic=False) + op = UnsupportedOperator(Vin, Vout) + + with pytest.raises(FastAssemblyUnavailable): + parallel_tosparse(op, comm) + + +def make_deriv_space(n1, n2, p1, p2, comm, periodic): + # dim 0 always periodic (matches make_space's default); dim 1 (the + # differentiation direction in test_parallel_tosparse_directional_derivative) + # uses `periodic`, since that's the one whose periodicity actually changes the + # matrix structure (wraparound coupling vs. a one-fewer-point boundary). + D = DomainDecomposition([n1, n2], periods=[True, periodic], comm=comm) + npts = [n1, n2] + gs, ge = compute_global_starts_ends(D, npts) + cart = CartDecomposition(D, npts, gs, ge, pads=[p1, p2], shifts=[1, 1]) + return StencilVectorSpace(cart, dtype=float) + + +@pytest.mark.parametrize('diffdir_periodic', [False, True]) +@pytest.mark.parametrize('transposed', [False, True]) +@pytest.mark.parallel +def test_parallel_tosparse_directional_derivative(diffdir_periodic, transposed, verbose=False): + """`DirectionalDerivativeOperator`'s default `.tosparse()` asserts outright at + nprocs > 1 (see `_directional_derivative_triples`'s docstring) -- the closed-form + reconstruction it falls back to instead must match the slow reference path, for + both a periodic and a non-periodic differentiation direction, transposed or not. + """ + p1, p2 = 1, 1 + comm = MPI.COMM_WORLD + n1, n2 = 8, 8 + V = make_deriv_space(n1, n2, p1, p2, comm, periodic=diffdir_periodic) + W = make_deriv_space(n1, n2 if diffdir_periodic else n2 - 1, p1, p2, comm, periodic=diffdir_periodic) + + op = DirectionalDerivativeOperator(V, W, diffdir=1, negative=False, transposed=transposed) + + fast = parallel_tosparse(op, comm) + slow = tosparse_via_matvec(op, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"periodic={diffdir_periodic} transposed={transposed} nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-10 + + +@pytest.mark.parallel +def test_parallel_tosparse_weighted_mass_operator_boundary_composition(verbose=False): + """`FakeWeightedMassOperator` wraps a StencilMatrix behind trivial extraction ops + but *non-trivial* boundary masking (`.dot()` != `._mat.dot()` alone) -- exactly + the shape that caused a real, silent ~8% numeric mismatch against the naive + "just unwrap `._mat`" shortcut on an actual Struphy run (before + `parallel_tosparse`'s `._mat`-composition unwrap rebuilt the *whole* + boundary/extraction chain instead). Must match the slow reference path. + """ + n1, n2, p1, p2 = 8, 6, 1, 1 + comm = MPI.COMM_WORLD + V = make_space(n1, n2, p1, p2, comm, periodic=False) + inner = make_stencil_matrix(V, p1, p2) + op = FakeWeightedMassOperator(V, inner) + + fast = parallel_tosparse(op, comm) + slow = tosparse_via_matvec(op, format="csr") + + diff = abs(fast - slow).max() + if verbose: + print(f"nprocs={comm.Get_size()} diff={diff:.2e}") + assert diff < 1e-8 diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index e7d47a233..79dd35121 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -18,6 +18,8 @@ 'array_to_psydac', 'petsc_to_psydac', 'tosparse_via_matvec', + 'parallel_tosparse', + 'FastAssemblyUnavailable', '_sym_ortho', ) @@ -259,6 +261,457 @@ def tosparse_via_matvec(op, format="csc"): mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) return mat.asformat(format) +#============================================================================== +class FastAssemblyUnavailable(Exception): + """Raised by `parallel_tosparse` when the operator tree could not be assembled via + the fast (O(1)-communication-round) path -- see its docstring. Callers should catch + this and fall back to `tosparse_via_matvec`.""" + + +def _local_flat_entries(V): + """ + List of (setter, global_flat_index) pairs, one per DOF *owned* by this rank (no + ghost/pad region), for a StencilVectorSpace or BlockVectorSpace V. + + `setter` is `('b', block_index, multi_index)` (usable as `vec[h][idx] = ...` for a + BlockVector) or `('s', multi_index)` (usable as `vec[idx] = ...` otherwise) -- + tagged rather than inferred from shape, since a StencilVectorSpace's own + `multi_index` can itself start with an int indistinguishable from a block index. + `global_flat_index` uses the same block-major, + `numpy.ravel_multi_index`-against-global-`npts` convention as + `tosparse_via_matvec` and `StencilMatrix.tosparse()` (`_tocoo_no_pads`) -- the same + one `Vector.toarray()` flattens to, which every caller of this module (e.g. + `DirectSolver.solve`'s `b.toarray()`/`x_flat.reshape(...)` round trip) already + relies on. All three MUST agree, since results from this function are combined + with plain `StencilMatrix`/`BlockLinearOperator` sparse matrices in the same + right-hand-side/solution vectors. + """ + if isinstance(V, BlockVectorSpace): + entries = [] + spoint = 0 + for h, sp in enumerate(V.spaces): + npts = sp.npts + iterables = [range(s, e + 1) for s, e in zip(sp.starts, sp.ends)] + for idx in itertools.product(*iterables): + flat = spoint + int(np.ravel_multi_index(idx, npts)) + entries.append((('b', h, idx), flat)) + spoint += int(np.prod(npts)) + return entries + elif isinstance(V, StencilVectorSpace): + npts = V.npts + iterables = [range(s, e + 1) for s, e in zip(V.starts, V.ends)] + return [(('s', idx), int(np.ravel_multi_index(idx, npts))) for idx in itertools.product(*iterables)] + else: + raise FastAssemblyUnavailable( + f'_local_flat_entries only supports StencilVectorSpace/BlockVectorSpace, got {type(V)}', + ) + + +def _set_entry(vec, setter, value): + if setter[0] == 'b': + _, h, idx = setter + vec[h][idx] = value + else: + _, idx = setter + vec[idx] = value + + +def _get_entry(vec, setter): + if setter[0] == 'b': + _, h, idx = setter + return vec[h][idx] + else: + _, idx = setter + return vec[idx] + + +def _replicate_triples(rows, cols, vals, shape, comm, dtype): + """Gather (rows, cols, vals) COO triples -- assumed *local* to this rank -- from + every rank and sum-combine them (matching duplicates, e.g. periodic wraparound, + exactly as `scipy.sparse.coo_matrix` does on `.tocsr()`) into one matrix identical + on every rank. One collective round, regardless of `shape`. + + `rows`/`cols`/`vals` may be plain sequences or arrays; always gathered and + concatenated as numpy arrays (`comm.allgather` pickles a numpy array through + mpi4py's out-of-band buffer protocol, and `np.concatenate` is vectorized) rather + than as Python lists -- for a leaf with real FEM bandwidth (tens to hundreds of + thousands of local nonzeros, e.g. a 3D mass matrix), converting through + element-by-element Python lists first was the dominant cost, dwarfing the O(1) + round-count win this function exists for. + """ + rows = np.asarray(rows, dtype=np.int64) + cols = np.asarray(cols, dtype=np.int64) + vals = np.asarray(vals, dtype=dtype) + + if comm is None or isinstance(comm, MockComm): + all_rows, all_cols, all_vals = rows, cols, vals + else: + gathered = comm.allgather((rows, cols, vals)) + all_rows = np.concatenate([g[0] for g in gathered]) + all_cols = np.concatenate([g[1] for g in gathered]) + all_vals = np.concatenate([g[2] for g in gathered]) + return sparse.coo_matrix((all_vals, (all_rows, all_cols)), shape=shape, dtype=dtype).tocsr() + + +def _probe_vector(V, entries, offset): + """A deterministic, reproducible-across-ranks probe Vector of space V: each owned + DOF gets a distinct nonzero value derived from its global flat index (never 0, and + never equal across two different `offset`s), so an operator's actual coupling + structure is very unlikely to accidentally look diagonal/masking by coincidence.""" + v = V.zeros() + for setter, flat in entries: + _set_entry(v, setter, 1.0 + 0.618033988749895 * ((flat + offset) % 104729)) + v.update_ghost_regions() + return v + + +def _validate_against_dot(node, candidate, comm, entries_domain, entries_codomain, seed): + """Check `candidate @ p == node.dot(p)` (this rank's owned output entries only) for + one probe vector `p`built from `_probe_vector`. `entries_domain` values already + carry each entry's global flat column index; `entries_codomain` likewise for rows. + Returns a local bool -- the caller combines these across ranks (see + `parallel_tosparse`) before trusting `candidate`.""" + V_domain = node.domain + p = _probe_vector(V_domain, entries_domain, seed) + p_flat_local = {flat: _get_entry(p, setter) for setter, flat in entries_domain} + + if comm is None or isinstance(comm, MockComm): + p_flat_full = dict(p_flat_local) + else: + gathered = comm.allgather(p_flat_local) + p_flat_full = {} + for d in gathered: + p_flat_full.update(d) + + p_full = np.zeros(candidate.shape[1], dtype=candidate.dtype) + for flat, val in p_flat_full.items(): + p_full[flat] = val + + q_candidate = candidate @ p_full + q_true = node.dot(p) + + # An aggregate (L2-norm) check, not a per-entry one: a wide-bandwidth FEM operator + # (e.g. a mass matrix with a degree-3 spline direction) sums many terms per row, + # in a different order than `candidate`'s (scipy's own summation order for the + # sparse matvec) -- individual output entries can then legitimately differ by much + # more than a tight per-entry relative tolerance even when `candidate` is exactly + # right, especially where terms partially cancel. Comparing the whole local output + # vector's norm to the whole error vector's norm is robust to that per-entry + # cancellation while still easily catching a genuinely wrong `candidate` (which + # differs at O(1) relative scale, not at rounding-error scale). + true_local = np.fromiter((_get_entry(q_true, setter) for setter, _ in entries_codomain), dtype=float) + cand_local = np.fromiter((q_candidate[flat] for _, flat in entries_codomain), dtype=float) + err = float(np.linalg.norm(true_local - cand_local)) + scale = float(np.linalg.norm(true_local)) + return err <= 1e-8 * scale + 1e-10 + + +def _directional_derivative_triples(op): + """Closed-form local (row, col, value) triples for a + `feectools.feec.derivatives.DirectionalDerivativeOperator` -- `.tosparse()`'s + default (no-pads) form isn't valid at nprocs > 1 for this operator (it asserts), + and its `with_pads=True` form returns a small *local* matrix in a totally + different (ghost-inclusive, non-globally-indexed) convention this module's + block-major global indexing can't reuse -- so this reconstructs the same bidiagonal + difference-operator matrix its serial `.tosparse()` builds, directly from the + operator's own definition (`out[i] = sign * (in[i + e_d] - in[i])` along direction + `d = op._diffdir`, `e_d` wrapped modulo `V.npts[d]` when periodic), one local + (globally-indexed) row at a time -- no basis-vector sweep, no padding subtleties. + Built in the "V -> W" (non-transposed) sense regardless of `op._transposed`; + `parallel_tosparse` transposes the result back if needed, exactly as the operator's + own serial `.tosparse()` does. + """ + V, W, d = op._spaceV, op._spaceW, op._diffdir + sign = -1.0 if op._negative else 1.0 + periodic = V.periods[d] + + if V.npts[d] == 1 and W.npts[d] == 1 and periodic: + return [], [], [] # degenerate single-cell-periodic case: the zero matrix + + rows, cols, vals = [], [], [] + for idx, row_flat in _local_flat_entries(W): + _, ii = idx + jj = ii + jj_next = list(ii) + jj_next[d] = (ii[d] + 1) % V.npts[d] if periodic else ii[d] + 1 + col_flat = int(np.ravel_multi_index(jj, V.npts)) + rows.append(row_flat) + cols.append(col_flat) + vals.append(-sign) + if periodic or jj_next[d] < V.npts[d]: + col_next_flat = int(np.ravel_multi_index(jj_next, V.npts)) + rows.append(row_flat) + cols.append(col_next_flat) + vals.append(sign) + return rows, cols, vals + + +def parallel_tosparse(op, comm, format="csr"): + """ + Assemble the full global sparse matrix of a `LinearOperator` tree using O(1) + collective-communication rounds (one per leaf node, roughly), instead of + `tosparse_via_matvec`'s O(`op.domain.dimension`) rounds (one basis vector per + global DOF) -- for the same replicated-on-every-rank result. + + Walks the operator tree using only types `feectools` itself defines + (`SumLinearOperator`, `ScaledLinearOperator`, `ComposedLinearOperator`, + `IdentityOperator`, `ZeroOperator`, `BlockLinearOperator`, + `DirectionalDerivativeOperator`): the first five compose exactly the way + `.tosparse()` already does in serial, just with the *leaves* below assembled + without a per-DOF basis-vector sweep; `BlockLinearOperator` is recursed into + block-by-block (not treated as one leaf) since a real Derham `grad`/`grad.T` can + have `DirectionalDerivativeOperator` blocks, whose own `.tosparse()` is unusable + in parallel (see `_directional_derivative_triples`, which reconstructs it in + closed form instead). For anything else -- most of it defined outside feectools + (`struphy.feec.mass.WeightedMassOperator`, `struphy.feec.linear_operators. + BoundaryOperator`, ...), which this module must not import -- one of three + O(1)-round strategies applies, tried in order: + + 1. Duck-typed unwrap: if the node has a `._mat` plus the same + `._V_extraction_op`/`._W_extraction_op`/`._V_boundary_op`/`._W_boundary_op`/ + `._transposed` attributes `struphy.feec.mass.WeightedMassOperator` has, + rebuild the exact composition its own `.dot()` applies (boundary and + extraction maps included, not just `._mat` alone -- an earlier version of + this function assumed trivial extraction ops meant `._mat` alone was enough, + which a real Struphy run showed is false whenever the boundary masks are + non-trivial) from parts each recursed into via `build()` in turn. + + 2. If the leaf has its own `.tosparse()` (true of `StencilMatrix` and + `BlockLinearOperator`): call it *locally* (no communication -- the same call + `A.tosparse()` already makes in serial, just once per rank instead of once + globally) and `allgather`-sum the local fragments into the replicated global + matrix. + + 3. Otherwise, if `leaf.domain is leaf.codomain` (a necessary condition to act as + a diagonal map): probe it with a value-tagged vector and check whether the + output is consistent with a per-DOF diagonal scaling (this is exactly what + essential-BC masking operators like `BoundaryOperator` are). Two `.dot()` + calls plus one `allgather`. + + Every leaf's result is *always* cross-checked against the operator's own `.dot()` + on a probe vector (strategies 1 and 2 both feed their candidate through the same + `_validate_against_dot` check that strategy 3 uses to detect diagonality in the + first place); an unrecognized leaf type (none of the three strategies applicable + or valid) is treated the same as a failed check. Whether any of this happens is a + pure function of operator *types*, identical on every rank by construction (same + model, same run) -- so every rank always issues the same sequence of collective + calls regardless of any individual check's pass/fail outcome; only *after* the + full tree is walked does one final `allreduce(MPI.LAND)` combine every check + across every rank into a single decision, so a data-dependent failure on one rank + cannot leave another rank waiting on a collective call that rank never issues (no + deadlock risk from divergent control flow). If that combined decision is False -- + or the tree contains a node type this function does not know how to handle at all + (e.g. `MatrixFreeLinearOperator`, always a deterministic, type-only decision, so + still consistent across ranks) -- `FastAssemblyUnavailable` is raised (on every + rank, identically) and the caller should fall back to `tosparse_via_matvec`. + + Parameters + ---------- + op : feectools.linalg.basic.LinearOperator + Operator to assemble. + + comm : MPI.Comm | feectools.ddm.mpi.MockComm | None + Communicator spanning every rank that owns a piece of `op`. + + format : str + scipy.sparse matrix format of the result. + + Returns + ------- + out : scipy.sparse matrix + The full `(op.codomain.dimension, op.domain.dimension)` matrix, identical on + every rank. + """ + # Imported here, not at module scope: these are feectools types this function + # checks via isinstance, kept local to make the "only feectools composite types + # are special-cased" contract easy to audit at a glance. + from feectools.linalg.basic import ComposedLinearOperator, IdentityOperator, ScaledLinearOperator, SumLinearOperator, ZeroOperator + from feectools.linalg.block import BlockLinearOperator + from feectools.feec.derivatives import DirectionalDerivativeOperator + + checks = [] + probe_seed = [1000003] # mutable cell; a fresh seed per leaf keeps probes independent + + def build(node): + if isinstance(node, ScaledLinearOperator): + return node._scalar * build(node._operator) + if isinstance(node, SumLinearOperator): + mats = [build(a) for a in node._addends] + out = mats[0] + for m in mats[1:]: + out = out + m + return out + if isinstance(node, ComposedLinearOperator): + mats = [build(m) for m in node._multiplicants] + out = mats[0] + for m in mats[1:]: + out = out @ m + return out + if isinstance(node, IdentityOperator): + return sparse.identity(node.domain.dimension, format="csr", dtype=node.dtype or float) + if isinstance(node, ZeroOperator): + return sparse.csr_matrix(node.shape, dtype=node.dtype or float) + if isinstance(node, BlockLinearOperator): + # Recurse into each block individually rather than calling + # `node.tosparse()` on the whole thing: a real Derham `grad`/`grad.T` is a + # BlockLinearOperator whose blocks can themselves be + # `DirectionalDerivativeOperator`s, whose *own* default `.tosparse()` + # asserts outright at nprocs > 1 (see `_directional_derivative_triples`) + # -- one such block would otherwise make the whole (possibly mostly + # StencilMatrix) BlockLinearOperator's `.tosparse()` raise. + nrows, ncols = node.n_block_rows, node.n_block_cols + block_domain = (lambda j: node.domain[j]) if ncols > 1 else (lambda j: node.domain) + block_codomain = (lambda i: node.codomain[i]) if nrows > 1 else (lambda i: node.codomain) + grid = [[None for _ in range(ncols)] for _ in range(nrows)] + for i in range(nrows): + for j in range(ncols): + if (i, j) in node._blocks: + grid[i][j] = build(node._blocks[i, j]) + else: + grid[i][j] = sparse.csr_matrix((block_codomain(i).dimension, block_domain(j).dimension)) + return sparse.bmat(grid, format="csr") + if isinstance(node, DirectionalDerivativeOperator): + # No generic strategy below applies (not diagonal-shaped in general, and + # its own `.tosparse()` is unusable here -- see + # `_directional_derivative_triples`); its structure is simple and fixed + # enough to reconstruct in closed form directly, still validated below + # like everything else. + V, W = node._spaceV, node._spaceW + rows, cols, vals = _directional_derivative_triples(node) + mat_vw = _replicate_triples(rows, cols, vals, (W.dimension, V.dimension), comm, node.dtype or float) + candidate = mat_vw.T.tocsr() if node._transposed else mat_vw + entries_domain = _local_flat_entries(node.domain) + entries_codomain = _local_flat_entries(node.codomain) + probe_seed[0] += 97 + ok = _validate_against_dot(node, candidate, comm, entries_domain, entries_codomain, probe_seed[0]) + checks.append(ok) + return candidate + + # Leaf: not one of the composite types above. Every strategy applicable to + # this leaf's *type* is always attempted, on every rank, regardless of any + # other rank's or strategy's data-dependent validation outcome -- see the + # docstring's "same collective calls on every rank" invariant. Only the first + # strategy that actually validates is kept. + entries_domain = _local_flat_entries(node.domain) + entries_codomain = _local_flat_entries(node.codomain) + shape = (node.codomain.dimension, node.domain.dimension) + dtype = node.dtype or float + probe_seed[0] += 97 + + # Duck-typed unwrap: struphy's `WeightedMassOperator` (M0, M1, ...) computes + # `V_boundary_op @ V_extraction_op @ _mat @ W_extraction_op.T @ W_boundary_op.T` + # (or the mirrored order when `._transposed`) on every `.dot()` call, by + # default with the boundary masks actually applied (`apply_bc=True`) -- *not* + # just `._mat` alone, even when both extraction ops are the identity (its own + # boundary masks can still be non-trivial, e.g. Dirichlet-BC-adjacent DOFs on + # a component of an Hcurl mass matrix -- discovered by this function's own + # validation rejecting the naive "just `._mat`" shortcut on exactly such a + # case, not by inspecting struphy's BC configuration). Rebuilding that same + # composition from its parts -- each recursed into via `build()`, so a + # boundary mask that itself needs the diagonal-probe strategy below still + # gets it -- is exact when every part is present (duck-typed by attribute, + # not `isinstance`, since none of these types live in feectools); still + # validated below regardless, as insurance against this composition itself + # being incomplete for some other struphy wrapper shaped differently. + parts = [getattr(node, name, "missing") for name in ( + "_mat", "_V_extraction_op", "_W_extraction_op", "_V_boundary_op", "_W_boundary_op", + )] + if "missing" not in parts: + inner_mat, v_ext, w_ext, v_bnd, w_bnd = parts + transposed = bool(getattr(node, "_transposed", False)) + try: + # Matches struphy.feec.mass.WeightedMassOperator.dot's own step + # sequence exactly (apply_bc=True, its default): non-transposed + # applies V_boundary_op.T, then V_extraction_op.T, then `._mat`, then + # W_extraction_op, then W_boundary_op, in that order (v -> out); the + # composed *matrix* is those same maps in reverse (rightmost applied + # first). `._transposed` mirrors V and W throughout. + if not transposed: + order = [w_bnd, w_ext, inner_mat, v_ext.transpose(), v_bnd.transpose()] + else: + order = [v_bnd, v_ext, inner_mat, w_ext.transpose(), w_bnd.transpose()] + mats = [build(m) for m in order] + candidate_inner = mats[0] + for m in mats[1:]: + candidate_inner = candidate_inner @ m + except Exception: + candidate_inner = None + if candidate_inner is not None and _validate_against_dot( + node, candidate_inner, comm, entries_domain, entries_codomain, probe_seed[0], + ): + checks.append(True) + return candidate_inner + + candidate_tosparse = None + try: + local_coo = node.tosparse().tocoo() + candidate_tosparse = _replicate_triples( + local_coo.row, local_coo.col, local_coo.data, + shape, comm, dtype, + ) + except Exception: + pass # this leaf's .tosparse() -- if it has one -- doesn't work here (e.g. + # raises, or -- as for struphy's BoundaryOperator -- succeeds but is + # documented serial-only and produces locally- rather than + # globally-indexed rows/cols at nprocs > 1); validation below (or, failing + # that, the diagonal-probe strategy) is what actually decides trust, not + # whether this call happened to raise. + + if candidate_tosparse is not None and _validate_against_dot( + node, candidate_tosparse, comm, entries_domain, entries_codomain, probe_seed[0], + ): + checks.append(True) + return candidate_tosparse + + if node.domain is not node.codomain: + # Not diagonal-shaped, and the .tosparse() attempt above (if any) didn't + # validate: nothing left to try for this leaf. + checks.append(False) + return candidate_tosparse if candidate_tosparse is not None else sparse.csr_matrix(shape, dtype=dtype) + + # Diagonal-probe strategy: two independent value-tagged probes; a true + # diagonal map reproduces (scaled by a per-DOF constant) or zeroes each one, + # consistently between the two -- exactly what an essential-BC mask does. + p1 = _probe_vector(node.domain, entries_domain, probe_seed[0]) + p2 = _probe_vector(node.domain, entries_domain, probe_seed[0] + 50000) + try: + o1 = node.dot(p1) + o2 = node.dot(p2) + except Exception: + checks.append(False) + return candidate_tosparse if candidate_tosparse is not None else sparse.csr_matrix(shape, dtype=dtype) + + rows, cols, vals = [], [], [] + ok = True + for setter, flat in entries_domain: + v1 = _get_entry(p1, setter) + v2 = _get_entry(p2, setter) + a1 = _get_entry(o1, setter) + a2 = _get_entry(o2, setter) + is_zero = abs(a1) < 1e-300 and abs(a2) < 1e-300 + if is_zero: + continue + d1 = a1 / v1 + d2 = a2 / v2 + if abs(d1 - d2) > 1e-8 * max(1.0, abs(d1)): + ok = False + break + rows.append(flat) + cols.append(flat) + vals.append(d1) + checks.append(ok) + return _replicate_triples(rows, cols, vals, shape, comm, dtype) + + result = build(op) + + all_ok = all(checks) + if comm is not None and not isinstance(comm, MockComm): + all_ok = comm.allreduce(all_ok, op=MPI.LAND) + if not all_ok: + raise FastAssemblyUnavailable('one or more leaf operators could not be verified') + + return result.asformat(format) + #============================================================================== def petsc_to_psydac(x, Xh, out=None): """ From 8b5ccf2cf97e0fc0eb2d04e6164f3d169f319c23 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 19 Aug 2026 19:18:17 +0200 Subject: [PATCH 5/7] Added codomain_local_nonzero_rows utility --- feectools/linalg/direct_solvers.py | 6 +--- feectools/linalg/utilities.py | 54 +++++++++++++++++++++++------- 2 files changed, 43 insertions(+), 17 deletions(-) diff --git a/feectools/linalg/direct_solvers.py b/feectools/linalg/direct_solvers.py index 0baa0c8b9..785ce0107 100644 --- a/feectools/linalg/direct_solvers.py +++ b/feectools/linalg/direct_solvers.py @@ -69,11 +69,7 @@ def __init__(self, u, l, bmat, transposed=False): else: msg = f'Cannot create a BandedSolver for bmat.dtype = {bmat.dtype}' raise NotImplementedError(msg) - # print(f"{bmat = } {type(bmat) = }") - if hasattr(bmat, "get"): # CuPy array - bmat = bmat.get() - else: - bmat = xp.asanyarray(bmat) + bmat = xp.to_numpy(bmat) self._bmat, self._ipiv, self._finfo = self._factor_function(bmat, l, u) self._sinfo = None diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 79dd35121..41407d9c5 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -141,6 +141,33 @@ def tosparse_via_matvec(op, format="csc"): numcols = op.domain.dimension data, row, col = [], [], [] + def local_nonzero_rows(stencil_vec, row_offset): + """(global_row_indices, values) for `stencil_vec`'s local interior data.""" + space = stencil_vec.space + idx_local = tuple( + slice(m * p, -m * p) if p != 0 else slice(0, None) + for p, m in zip(stencil_vec.pads, space.shifts) + ) + local_data = xp.to_numpy(stencil_vec._data[idx_local]) + nz = np.nonzero(local_data) + starts = space.starts + global_multi = tuple(nz[d] + int(starts[d]) for d in range(len(nz))) + rows = row_offset + np.ravel_multi_index(global_multi, space.npts) + return rows, local_data[nz] + + def codomain_local_nonzero_rows(vec): + """`local_nonzero_rows`, dispatched over `op.codomain`'s type.""" + if isinstance(op.codomain, BlockVectorSpace): + all_rows, all_vals = [], [] + row_offset = 0 + for b, sp in enumerate(op.codomain.spaces): + r, val = local_nonzero_rows(vec[b], row_offset) + all_rows.append(r) + all_vals.append(val) + row_offset += sp.dimension + return np.concatenate(all_rows), np.concatenate(all_vals) + return local_nonzero_rows(vec, 0) + if isinstance(op.domain, BlockVectorSpace): starts = [vi.starts for vi in v] ends = [vi.ends for vi in v] @@ -184,14 +211,12 @@ def tosparse_via_matvec(op, format="csc"): tmp2 *= 0.0 op.dot(v, out=tmp2) c = spoint + int(np.ravel_multi_index(i, npts[h])) - aux = xp.to_numpy(tmp2.toarray()) - for r in np.nonzero(aux)[0]: - data.append(aux[r]) - col.append(c) - row.append(int(r)) + rs, vals = codomain_local_nonzero_rows(tmp2) + row.append(rs) + col.append(np.full(rs.shape, c)) + data.append(vals) if rank == currentrank: v[h][i] = 0.0 - v[h].update_ghost_regions() cumulative = 1 for i in range(ndim[h]): cumulative *= npts[h][i] @@ -231,14 +256,12 @@ def tosparse_via_matvec(op, format="csc"): v.update_ghost_regions() op.dot(v, out=tmp2) c = int(np.ravel_multi_index(i, npts)) - aux = xp.to_numpy(tmp2.toarray()) - for r in np.nonzero(aux)[0]: - data.append(aux[r]) - col.append(c) - row.append(int(r)) + rs, vals = codomain_local_nonzero_rows(tmp2) + row.append(rs) + col.append(np.full(rs.shape, c)) + data.append(vals) if rank == currentrank: v[i] = 0.0 - v.update_ghost_regions() if comm is None or isinstance(comm, MockComm): all_rows, all_cols, all_data = row, col, data @@ -258,6 +281,13 @@ def tosparse_via_matvec(op, format="csc"): all_cols = comm.bcast(None, root=0) all_data = comm.bcast(None, root=0) + if all_rows: + all_rows = np.concatenate(all_rows) + all_cols = np.concatenate(all_cols) + all_data = np.concatenate(all_data) + else: + all_rows = all_cols = all_data = np.empty(0, dtype=int) + mat = sparse.coo_matrix((all_data, (all_rows, all_cols)), shape=(numrows, numcols), dtype=op.dtype) return mat.asformat(format) From 06f95fcc62f095e43b1cd2d46eb46452c287318f Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 20 Aug 2026 12:53:22 +0200 Subject: [PATCH 6/7] bugfix --- feectools/linalg/utilities.py | 36 +++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 41407d9c5..8dcc54221 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -589,6 +589,7 @@ def build(node): # asserts outright at nprocs > 1 (see `_directional_derivative_triples`) # -- one such block would otherwise make the whole (possibly mostly # StencilMatrix) BlockLinearOperator's `.tosparse()` raise. + checks_before = len(checks) nrows, ncols = node.n_block_rows, node.n_block_cols block_domain = (lambda j: node.domain[j]) if ncols > 1 else (lambda j: node.domain) block_codomain = (lambda i: node.codomain[i]) if nrows > 1 else (lambda i: node.codomain) @@ -599,7 +600,34 @@ def build(node): grid[i][j] = build(node._blocks[i, j]) else: grid[i][j] = sparse.csr_matrix((block_codomain(i).dimension, block_domain(j).dimension)) - return sparse.bmat(grid, format="csr") + candidate = sparse.bmat(grid, format="csr") + + children_ok = all(checks[checks_before:]) + if comm is not None and not isinstance(comm, MockComm): + children_ok = comm.allreduce(children_ok, op=MPI.LAND) + block_ok = False + if children_ok: + entries_domain = _local_flat_entries(node.domain) + entries_codomain = _local_flat_entries(node.codomain) + probe_seed[0] += 97 + block_ok = _validate_against_dot( + node, candidate, comm, entries_domain, entries_codomain, probe_seed[0], + ) + if comm is not None and not isinstance(comm, MockComm): + block_ok = comm.allreduce(block_ok, op=MPI.LAND) + if block_ok: + checks.append(True) + return candidate + + try: + fallback = tosparse_via_matvec(node, format="csr") + except Exception: + checks.append(False) + return candidate + + del checks[checks_before:] + checks.append(True) + return fallback if isinstance(node, DirectionalDerivativeOperator): # No generic strategy below applies (not diagonal-shaped in general, and # its own `.tosparse()` is unusable here -- see @@ -729,8 +757,12 @@ def build(node): rows.append(flat) cols.append(flat) vals.append(d1) + candidate_diag = _replicate_triples(rows, cols, vals, shape, comm, dtype) + ok = ok and _validate_against_dot( + node, candidate_diag, comm, entries_domain, entries_codomain, probe_seed[0], + ) checks.append(ok) - return _replicate_triples(rows, cols, vals, shape, comm, dtype) + return candidate_diag result = build(op) From cefae919ade2bd8b339d74a11b5d59e1c88fa900 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 20 Aug 2026 13:33:17 +0200 Subject: [PATCH 7/7] Added v[h].update_ghost_regions() --- feectools/linalg/utilities.py | 1 + 1 file changed, 1 insertion(+) diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index 8dcc54221..bd5556141 100644 --- a/feectools/linalg/utilities.py +++ b/feectools/linalg/utilities.py @@ -217,6 +217,7 @@ def codomain_local_nonzero_rows(vec): data.append(vals) if rank == currentrank: v[h][i] = 0.0 + v[h].update_ghost_regions() cumulative = 1 for i in range(ndim[h]): cumulative *= npts[h][i]