KroneckerSumSolver: stiffness coefficients and additional Kronecker terms

feectools/linalg/kron.py

@@ -1373,21 +1373,30 @@ class KroneckerSumSolver(LinearOperator):
 
     .. math::
 
-        A = \sum_{d=1}^n M_1 \otimes \dots \otimes M_{d-1} \otimes S_d \otimes M_{d+1} \otimes \dots \otimes M_n
-            + \sigma \, M_1 \otimes \dots \otimes M_n \,,
+        A = \sum_{d=1}^n a_d \, M_1 \otimes \dots \otimes M_{d-1} \otimes S_d \otimes M_{d+1} \otimes \dots \otimes M_n
+            + \sigma \, M_1 \otimes \dots \otimes M_n
+            + \sum_t c_t \, F_{t,1} \otimes \dots \otimes F_{t,n} \,,
 
     with symmetric 1d matrices $S_d$ and symmetric positive definite 1d matrices $M_d$
-    (e.g. stiffness and mass matrices of a Laplacian on a tensor-product grid).
+    (e.g. stiffness and mass matrices of a Laplacian on a tensor-product grid), coefficients
+    $a_d$ (default 1) and optional additional terms with symmetric 1d factors $F_{t,d}$.
 
     In each direction, the generalized eigenproblem $S_d U_d = M_d U_d \Lambda_d$ is solved,
     with $U_d^T M_d U_d = I$. Then $U^T A U = \Lambda$ with $U = U_1 \otimes \dots \otimes U_n$
-    and the diagonal $\Lambda = \Lambda_1 \oplus \dots \oplus \Lambda_n + \sigma$ (entries
-    $\lambda_{1,i_1} + \dots + \lambda_{n,i_n} + \sigma$), hence
+    and the diagonal $\Lambda = a_1 \Lambda_1 \oplus \dots \oplus a_n \Lambda_n + \sigma + \sum_t c_t \,
+    D_{t,1} \otimes \dots \otimes D_{t,n}$ with $D_{t,d} = \mathrm{diag}(U_d^T F_{t,d} U_d)$, hence
 
     .. math::
 
         A^{-1} = U \, \Lambda^{-1} \, U^T \,.
 
+    This is exact if the $U_d$ diagonalize the $F_{t,d}$, e.g. for $F_{t,d} = M_d$, or for
+    $F_{t,d} = c_d c_d^T / (\mathbb 1^T c_d)$ with $c_d = M_d \mathbb 1$ when the constants are a
+    (simple) eigenvector, i.e. $S_d \mathbb 1 = 0$: then $U_d^T F_{t,d} U_d = e_0 e_0^T$, and
+    $M_d - F_{t,d}$ removes the mean (e.g. in a periodic direction). Otherwise the additional terms
+    are replaced by their diagonal in the eigenbasis (an approximation).
+    The stiffness $S_d$ defines the eigenbasis also for $a_d = 0$.
+
     The dense matrices $U_d^T$ and $U_d$ are applied with KroneckerLinearSolver (also along
     distributed axes), $\Lambda^{-1}$ locally. Entries of $\Lambda$ that vanish (relative to its
     largest entry, e.g. constants for a periodic Laplacian with $\sigma = 0$) are skipped, which
@@ -1409,6 +1418,16 @@ class KroneckerSumSolver(LinearOperator):
 
     rtol : float
         Entries of $\Lambda$ with ``|lambda| <= rtol * max|Lambda|`` are treated as zero (pseudo-inverse).
+        Without ``stiffness_coeffs`` and ``mass_terms``, ``max|Lambda|`` is estimated by
+        ``n * max|lambda_d| + |sigma|`` (as before); otherwise it is the maximum of the assembled
+        diagonal (the terms may cancel).
+
+    stiffness_coeffs : sequence of float, optional
+        Coefficients $a_d$ of the stiffness terms (default 1).
+
+    mass_terms : sequence of (float, sequence of array | None), optional
+        Additional terms $c_t \, F_{t,1} \otimes \dots \otimes F_{t,n}$ as pairs ``(c_t, factors)``
+        with dense global 1d matrices $F_{t,d}$; None stands for $M_d$.
     """
 
     class _DenseApply(LinearSolver):
@@ -1437,12 +1456,22 @@ def __init__(self,
                  stiffness: Sequence,
                  mass: Sequence,
                  sigma: float = 0.0,
-                 rtol: float = 1e-12):
+                 rtol: float = 1e-12,
+                 *,
+                 stiffness_coeffs: Sequence[float] | None = None,
+                 mass_terms: Sequence[tuple] | None = None):
 
         from scipy.linalg import eigh  # deferred: scipy.linalg is slow to import
 
         assert isinstance(V, StencilVectorSpace)
         assert len(stiffness) == len(mass) == V.ndim
+        stiffness_coeffs_given = stiffness_coeffs is not None
+        if stiffness_coeffs is None:
+            stiffness_coeffs = [1.0] * V.ndim
+        assert len(stiffness_coeffs) == V.ndim
+        mass_terms = list(mass_terms or [])
+        for coeff, factors in mass_terms:
+            assert len(factors) == V.ndim, 'Each mass term needs one factor (or None) per direction.'
 
         U = []
         lam = []
@@ -1463,12 +1492,33 @@ def __init__(self,
                 lam.append(l_d)
 
         # eigenvalues of A on the rows owned by this process
-        local_lam = [l[s:e+1] for l, s, e in zip(lam, V.starts, V.ends)]
+        local_shape = tuple(e - s + 1 for s, e in zip(V.starts, V.ends))
+        local_lam = [a * l[s:e+1] for a, l, s, e in zip(stiffness_coeffs, lam, V.starts, V.ends)]
         Lam = reduce(np.add.outer, local_lam) + sigma if V.ndim > 1 else local_lam[0] + sigma
-        Lam = np.asarray(Lam).reshape(tuple(e - s + 1 for s, e in zip(V.starts, V.ends)))
-
-        # pseudo-inverse: skip (numerically) vanishing eigenvalues
-        lam_max = max(abs(l).max() for l in lam) * V.ndim + abs(sigma)
+        Lam = np.asarray(Lam).reshape(local_shape)
+
+        # additional terms: their diagonal in the eigenbasis, diag(U_d^T F U_d) (1 for F = M_d)
+        generalized = stiffness_coeffs_given or bool(mass_terms)
+        for coeff, factors in mass_terms:
+            diags = []
+            for d, (F, U_d) in enumerate(zip(factors, U)):
+                if F is None:
+                    diags.append(np.ones(V.npts[d]))
+                else:
+                    F = np.asarray(xp.to_numpy(xp.asarray(F)), dtype=float)
+                    assert F.shape == (V.npts[d], V.npts[d]), f'Factor of direction {d} has shape {F.shape}.'
+                    diags.append(np.einsum('ij,ik,kj->j', U_d, F, U_d))
+            local_diags = [g[s:e+1] for g, s, e in zip(diags, V.starts, V.ends)]
+            Lam = Lam + coeff * np.asarray(reduce(np.multiply.outer, local_diags)).reshape(local_shape)
+
+        # pseudo-inverse: skip (numerically) vanishing eigenvalues, relative to max|Lambda|
+        if generalized:
+            # signed terms may cancel: maximum of the assembled diagonal (over all processes)
+            lam_max = float(np.abs(Lam).max()) if Lam.size else 0.0
+            if V.parallel:
+                lam_max = max(V.cart.comm.allgather(lam_max))
+        else:
+            lam_max = max(abs(l).max() for l in lam) * V.ndim + abs(sigma)
         inv_Lam = np.zeros_like(Lam)
         nonzero = np.abs(Lam) > rtol * lam_max
         inv_Lam[nonzero] = 1.0 / Lam[nonzero]

feectools/linalg/tests/test_kron_stencil_matrix.py

@@ -475,6 +475,71 @@ def test_kron_sum_solver_par(sigma, with_none):
     check_sum_solver(MPI.COMM_WORLD, sigma, with_none)
 
 
+def check_sum_solver_terms(comm, sigma):
+    """Stiffness coefficients (one of them zero) and a term removing the mean in the periodic
+    directions 0 and 2: exact, since the constants are eigenvectors there."""
+    npts, periods = [6, 7, 8], [True, False, True]
+    W = make_space(comm, npts, [2, 2, 3], periods)
+    S, M = zip(*(laplace_1d(n, P, seed=d) for d, (n, P) in enumerate(zip(npts, periods))))
+    coeffs = [1.0, 0.5, 0.0]
+
+    def mean_projection(d):
+        c = M[d] @ np.ones(npts[d])
+        return np.outer(c, c) / c.sum()
+
+    factors = [mean_projection(0), None, mean_projection(2)]
+
+    def kron_term(d):
+        return reduce(np.kron, [M[e] if e != d else S[d] for e in range(3)])
+
+    A = sum(a * kron_term(d) for d, a in enumerate(coeffs)) + sigma * reduce(np.kron, M)
+    A = A - sigma * reduce(np.kron, [M[d] if f is None else f for d, f in enumerate(factors)])
+    solver = KroneckerSumSolver(W, S, M, sigma=sigma, stiffness_coeffs=coeffs, mass_terms=[(-sigma, factors)])
+
+    # A is singular (e.g. the constants): pseudo-inverse, exact for right-hand sides in the range of A
+    bglob = A @ np.random.default_rng(6).random(A.shape[0])
+    x = solver.dot(StencilVector_from(W, bglob))
+    # relative to the largest entry (entries close to zero next to large ones)
+    assert np.abs(A @ local_to_global(W, x) - bglob).max() < 1e-12 * np.abs(bglob).max()
+
+
+def check_sum_solver_cancellation(comm, eps):
+    """sigma = 1 and a term (eps - 1) * M_1 x M_2 x M_3 cancel to A = eps * M_1 x M_2 x M_3: the
+    pseudo-inverse cutoff relative to the assembled diagonal keeps all eigenvalues."""
+    npts, periods = [6, 7, 8], [True, False, True]
+    W = make_space(comm, npts, [2, 2, 3], periods)
+    _, M = zip(*(laplace_1d(n, P, seed=d) for d, (n, P) in enumerate(zip(npts, periods))))
+    A = eps * reduce(np.kron, M)
+    solver = KroneckerSumSolver(W, [None] * 3, M, sigma=1.0, mass_terms=[(eps - 1.0, [None] * 3)])
+
+    bglob = A @ np.random.default_rng(7).random(A.shape[0])
+    x = solver.dot(StencilVector_from(W, bglob))
+    # 1 + (eps - 1) carries the round-off of 1, i.e. a relative error ~ 1e-16 / eps; with discarded
+    # eigenvalues (cutoff relative to the sum of the terms) x would vanish and the residual be |b|
+    tol = 100 * np.finfo(float).eps / eps
+    assert np.abs(A @ local_to_global(W, x) - bglob).max() < tol * np.abs(bglob).max()
+
+
+@pytest.mark.parametrize('eps', [1e-6, 1e-13])
+def test_kron_sum_solver_cancellation_ser(eps):
+    check_sum_solver_cancellation(None, eps)
+
+@pytest.mark.mpi
+@pytest.mark.parametrize('eps', [1e-6, 1e-13])
+def test_kron_sum_solver_cancellation_par(eps):
+    check_sum_solver_cancellation(MPI.COMM_WORLD, eps)
+
+
+@pytest.mark.parametrize('sigma', [0.7, 3.0])
+def test_kron_sum_solver_terms_ser(sigma):
+    check_sum_solver_terms(None, sigma)
+
+@pytest.mark.mpi
+@pytest.mark.parametrize('sigma', [0.7, 3.0])
+def test_kron_sum_solver_terms_par(sigma):
+    check_sum_solver_terms(MPI.COMM_WORLD, sigma)
+
+
 #===============================================================================
 # Row layout of the factors, duplicate entries of periodic factors
 #===============================================================================

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "feectools"
-version         = "0.7.0"
+version         = "0.8.0"
 description     = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies."
 readme          = "README.md"
 requires-python = ">= 3.10"