Kronecker preconditioners: StiffnessPreconditioner, refactor of the mass preconditioners; bump feectools

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "cunumpy >= 0.6.1, <= 0.6.1",
     "pyccel>=2.2.0, <=2.2.3",
     "maybempi>=0.1.2",
-    "feectools>=0.6.0, <=0.6.0",
+    "feectools>=0.7.0, <=0.7.0",
     "scipy<=1.18.0",
     "h5py<=3.16.0",
     "h5netcdf<=1.8.1",

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):

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,15 +2197,15 @@ 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:
                     A.dot(x, out=y)
                 x = y
 
             # last operator
-            A = self._I.multiplicants[0]
+            A = self._I.multiplicands[0]
             if out is None:
                 out = A.dot(x)
             else:

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)."""

src/struphy/feec/tests/test_preconditioner_transpose.py

@@ -2,7 +2,9 @@
 
 
 @pytest.mark.parametrize("dim_reduce", [0, 2])
-def test_mass_preconditioner_transpose(dim_reduce):
+@pytest.mark.parametrize("weight_reduction", ["midpoint", "average"])
+@pytest.mark.parametrize("diagonal_scaling", [False, True])
+def test_mass_preconditioner_transpose(dim_reduce, weight_reduction, diagonal_scaling):
     """The transposes of the mass-matrix preconditioners keep their class and options."""
 
     from struphy import domains
@@ -16,23 +18,32 @@ def test_mass_preconditioner_transpose(dim_reduce):
     derham = Derham(grid, DerhamOptions(degree=(1, 1, 1), bcs=(None, None, None)))
     mass_ops = WeightedMassOperators(derham, domains.Colella())
 
-    pc = MassMatrixPreconditioner(mass_ops.M1, apply_bc=False, dim_reduce=dim_reduce)
+    pc = MassMatrixPreconditioner(
+        mass_ops.M1,
+        apply_bc=False,
+        dim_reduce=dim_reduce,
+        weight_reduction=weight_reduction,
+        diagonal_scaling=diagonal_scaling,
+    )
     pc_T = pc.transpose()
     assert type(pc_T) is MassMatrixPreconditioner
     assert pc_T._apply_bc is False
     assert pc_T._dim_reduce == dim_reduce
+    assert pc_T.weight_reduction == weight_reduction
+    assert pc_T.diagonal_scaling == diagonal_scaling
 
     pc_diag_T = MassMatrixDiagonalPreconditioner(mass_ops.M1, apply_bc=False).transpose()
     assert type(pc_diag_T) is MassMatrixDiagonalPreconditioner
     assert pc_diag_T._apply_bc is False
 
 
-def test_mass_diagonal_preconditioner_assembles_diagonal_blocks_only(monkeypatch):
-    """The logical mass matrix of MassMatrixDiagonalPreconditioner needs only the diagonal blocks."""
+def test_mass_diagonal_preconditioner_assembles_1d_matrices_only(monkeypatch):
+    """MassMatrixDiagonalPreconditioner assembles only 1d mass matrices (no 3d logical mass matrix);
+    the diagonal of the logical mass matrix is the one of its Kronecker approximation."""
 
     from struphy import domains
     from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators
-    from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner
+    from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, MassMatrixPreconditioner
     from struphy.feec.psydac_derham import Derham
     from struphy.io.options import DerhamOptions
     from struphy.topology.grids import TensorProductGrid
@@ -42,18 +53,78 @@ def test_mass_diagonal_preconditioner_assembles_diagonal_blocks_only(monkeypatch
     mass_ops = WeightedMassOperators(derham, domains.Colella())
 
     M1 = mass_ops.M1
-    weights_infos = []
+    created = []
     init = WeightedMassOperator.__init__
 
-    def recording_init(self, *args, weights_info=None, **kwargs):
-        weights_infos.append(weights_info)
-        init(self, *args, weights_info=weights_info, **kwargs)
+    def recording_init(self, derham, V, W, *args, **kwargs):
+        created.append(V.ldim)
+        init(self, derham, V, W, *args, **kwargs)
 
     monkeypatch.setattr(WeightedMassOperator, "__init__", recording_init)
-    MassMatrixDiagonalPreconditioner(M1)
+    pc = MassMatrixDiagonalPreconditioner(M1)
+
+    assert isinstance(pc, MassMatrixPreconditioner)
+    assert pc.dim_reduce is None and pc.diagonal_scaling
+    assert created and all(ldim == 1 for ldim in created)
+
+
+def test_mass_diagonal_preconditioner_zero_diagonal_block():
+    """A mass operator with a zero (None) diagonal block, e.g. after reassembly with a zero weight:
+    construction and update_mass_operator must work; the scaling of that block is 1."""
+
+    import cunumpy as xp
+
+    from struphy import domains
+    from struphy.feec.mass import WeightedMassOperator, WeightedMassOperators
+    from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner
+    from struphy.feec.psydac_derham import Derham
+    from struphy.feec.utilities import create_equal_random_arrays
+    from struphy.io.options import DerhamOptions
+    from struphy.topology.grids import TensorProductGrid
+
+    grid = TensorProductGrid(num_elements=(3, 3, 2))
+    derham = Derham(grid, DerhamOptions(degree=(1, 1, 1), bcs=(None, None, None)))
+    mass_ops = WeightedMassOperators(derham, domains.Colella())
+
+    def one(e1, e2, e3):
+        return xp.ones_like(e1, dtype=float)
+
+    weights = [[one if i == j and i != 2 else None for j in range(3)] for i in range(3)]
+    M = WeightedMassOperator(derham, derham.V1fem, derham.V1fem, weights_info=weights)
+    M.assemble()
+    assert M.matrix[2, 2] is None
+
+    pc = MassMatrixDiagonalPreconditioner(mass_ops.M1)
+    pc.update_mass_operator(M)
+    scaling = pc._scaling[0]
+    assert xp.allclose(scaling[2, 2]._data, 1.0)
+
+    _, x = create_equal_random_arrays(derham.V1fem, seed=1)
+    assert all(xp.all(xp.isfinite(b.toarray())) for b in pc.dot(x).blocks)
+
+
+def test_mass_diagonal_preconditioner_small_periodic_direction():
+    """Periodic direction with fewer points than the stencil width (1 element, degree 1): the main
+    diagonal of the 1d factors must be at offset 0, otherwise the diagonal scaling vanishes."""
+
+    import cunumpy as xp
+    from feectools.linalg.solvers import inverse
+
+    from struphy import domains
+    from struphy.feec.mass import WeightedMassOperators
+    from struphy.feec.preconditioner import MassMatrixDiagonalPreconditioner, _local_diagonal
+    from struphy.feec.psydac_derham import Derham
+    from struphy.feec.utilities import create_equal_random_arrays
+    from struphy.io.options import DerhamOptions
+    from struphy.topology.grids import TensorProductGrid
+
+    derham = Derham(grid=TensorProductGrid(num_elements=(4, 4, 1)), options=DerhamOptions(degree=(2, 2, 1)))
+    mass_ops = WeightedMassOperators(derham=derham, domain=domains.Cuboid())
+
+    pc = MassMatrixDiagonalPreconditioner(mass_ops.Mv)
+    assert all(float(xp.min(d)) > 0.0 for d in _local_diagonal(pc.matrix))
 
-    # the logical mass matrix is the only 3x3 weighted mass operator created
-    (fun,) = [w for w in weights_infos if isinstance(w, list) and len(w) == 3]
-    for i in range(3):
-        for j in range(3):
-            assert callable(fun[i][j]) == (i == j)
+    _, x = create_equal_random_arrays(derham.Vvfem, seed=1)
+    inv = inverse(mass_ops.Mv, "pcg", pc=pc, tol=1e-12, maxiter=100)
+    inv.dot(x)
+    assert inv._info["success"]