Geometric multigrid preconditioner (MPI, 3D coarsening, operator re-discretization)

pyproject.toml

@@ -26,7 +26,7 @@ dependencies = [
     "numpy<=2.5.0",
     "cunumpy>=0.2.0, <=0.3.0",
     "pyccel>=2.2.0, <=2.2.3",
-    "feectools>=0.1.11, <=0.2.0",
+    "feectools>=0.3.0, <=0.3.0",
     "scipy<=1.18.0",
     "h5py<=3.16.0",
     "h5netcdf<=1.8.1",

src/struphy/feec/basis_projection_ops.py

@@ -1851,6 +1851,56 @@ def dot(self, v, out=None, tol=1e-14, maxiter=1000):
 
         return out
 
+    @property
+    def is_reconstructible(self) -> bool:
+        """Whether the operator can be re-created from :meth:`to_dict` (e.g. on another Derham).
+
+        False if a weight is given as values at the projection points (an array bound to the current Derham).
+        """
+        return not any(isinstance(w, xp.ndarray) for row in self._weights for w in row)
+
+    def to_dict(self) -> dict:
+        """Recipe for re-creating the operator with :meth:`from_dict` (on any Derham).
+
+        Weights are stored as given (callables are kept as objects, hence the dictionary is in general not JSON serializable).
+        """
+        if not self.is_reconstructible:
+            raise ValueError("BasisProjectionOperator with weights given as arrays cannot be serialized.")
+        V_id, W_id = (
+            (self._codomain_symbolic_name, self._domain_symbolic_name)
+            if self._transposed
+            else (self._domain_symbolic_name, self._codomain_symbolic_name)
+        )
+        return {
+            "type": self.__class__.__name__,
+            "params": {
+                "V_id": V_id,
+                "W_id": W_id,
+                "weights": [list(row) for row in self._weights],
+                "transposed": self._transposed,
+                "polar_shift": self._polar_shift,
+                "use_cache": self._use_cache,
+            },
+        }
+
+    @classmethod
+    def from_dict(cls, dct: dict, derham: Derham) -> "BasisProjectionOperator":
+        """Re-create a :class:`BasisProjectionOperator` from :meth:`to_dict` on the given Derham,
+        with the Derham's (global) commuting projector, extraction and boundary operators."""
+        assert dct["type"] == cls.__name__
+        params = dct["params"]
+        V_id, W_id = params["V_id"], params["W_id"]
+        return cls(
+            derham.projectors[W_id],
+            derham.fem_spaces[V_id],
+            [list(row) for row in params["weights"]],
+            V_extraction_op=derham.extraction_ops[V_id],
+            V_boundary_op=derham.boundary_ops[V_id],
+            transposed=params["transposed"],
+            polar_shift=params["polar_shift"],
+            use_cache=params["use_cache"],
+        )
+
     def transpose(self, conjugate=False):
         """
         Returns the transposed operator.

src/struphy/feec/mass.py

@@ -1254,6 +1254,23 @@ def f_call_matrix(e1, e2, e3):
             dry_run=dry_run,
         )
 
+        # weights given at quadrature points or as spline functions are bound to this Derham
+        grid_bound = len(spline_functions) > 0 or (
+            isinstance(weights, list) and any(isinstance(w, xp.ndarray) for row in weights for w in row)
+        )
+        out._creation_info = (
+            None
+            if grid_bound
+            else {
+                "V_id": V_id,
+                "W_id": W_id,
+                "name": name,
+                "weights": weights,
+                "transposed": transposed,
+                "is_transpose": False,
+            }
+        )
+
         if assemble and not dry_run:
             out.assemble()
 
@@ -1465,6 +1482,9 @@ def __init__(
         self._name = name
         self._dry_run = dry_run
 
+        # recipe for re-creating the operator with WeightedMassOperators.create_weighted_mass, see to_dict()
+        self._creation_info: dict | None = None
+
         assert not (dry_run and transposed), "dry_run=True is not supported for transposed operators."
 
         # spline functions that are used as weights in the operator, to be evaluated at quadrature points
@@ -2066,6 +2086,9 @@ def transpose(self, conjugate=False):
         # weights of M in its own (transposed) block order
         M._weights = [[self._weights[n][m] for n in range(len(self._weights))] for m in range(len(self._weights[0]))]
 
+        if self._creation_info is not None:
+            M._creation_info = dict(self._creation_info, is_transpose=not self._creation_info["is_transpose"])
+
         if self._matrix_free:
             if self._symmetry is not None:
                 M.assemble(weights=M._weights)
@@ -2107,6 +2130,9 @@ def assemble(self, weights=None, clear=True):
         assert not self._dry_run, (
             "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)."
         )
+        if weights is not None or not clear:
+            # the data no longer stems from the creation recipe
+            self._creation_info = None
 
         if self._matrix_free:
             if weights is not None:
@@ -2325,6 +2351,67 @@ def assemble(self, weights=None, clear=True):
 
             logger.debug("Done.")
 
+    @property
+    def is_reconstructible(self) -> bool:
+        """Whether the operator can be re-created from :meth:`to_dict` (e.g. on another Derham).
+
+        True for operators created by :meth:`WeightedMassOperators.create_weighted_mass` (and their transposes)
+        whose data has not been modified afterwards (by ``assemble(weights=...)``, in-place arithmetic, ...).
+        """
+        return self._creation_info is not None
+
+    def to_dict(self) -> dict:
+        """Recipe for re-creating the operator with :meth:`WeightedMassOperators.create_weighted_mass`.
+
+        The weights are stored as given at creation. The dictionary is JSON serializable if they are
+        strings (``'Ginv'``, ``'sqrt_g'``, ...) or nested lists of numbers; callables are kept as objects.
+        Re-create the operator (on any Derham) with :meth:`from_dict`.
+        """
+        if self._creation_info is None:
+            raise ValueError(
+                f"WeightedMassOperator {self.name!r} cannot be serialized: it was not created by "
+                "WeightedMassOperators.create_weighted_mass or its data was modified afterwards."
+            )
+        params = dict(self._creation_info)
+        # tuple (1D product of weights) and 2D list (block weights) are different formats; store the
+        # tuple as a list (JSON) and record its type, since its entries may themselves be (3x3) lists
+        params["weights_is_tuple"] = isinstance(params["weights"], tuple)
+        if params["weights_is_tuple"]:
+            params["weights"] = list(params["weights"])
+        return {
+            "type": self.__class__.__name__,
+            "params": params,
+        }
+
+    @classmethod
+    def from_dict(cls, dct: dict, mass_ops: "WeightedMassOperators") -> "WeightedMassOperator":
+        """Re-create a :class:`WeightedMassOperator` from :meth:`to_dict` with the given collection.
+
+        Parameters
+        ----------
+        dct : dict
+            Output of :meth:`to_dict`.
+
+        mass_ops : WeightedMassOperators
+            Collection providing the Derham, domain and matrix_free option of the new operator.
+        """
+        assert dct["type"] == cls.__name__
+        params = dct["params"]
+        name = params["name"]
+        weights = params["weights"]
+        if params["weights_is_tuple"]:
+            weights = tuple(weights)
+
+        out = mass_ops.create_weighted_mass(
+            params["V_id"],
+            params["W_id"],
+            name=name,
+            weights=weights,
+            assemble=True,
+            transposed=params["transposed"],
+        )
+        return out.T if params["is_transpose"] else out
+
     def copy(self, out=None):
         """Create a copy of self, that can potentially be stored in a given WeightedMassOperator.
 
@@ -2357,21 +2444,30 @@ def copy(self, out=None):
             out._weights = [list(row) for row in self._weights]
 
         self._mat.copy(out=out._mat)
+
+        if self._creation_info is None:
+            out._creation_info = None
+        else:
+            out._creation_info = dict(self._creation_info)  # to create a separate dictionary
+
         return out
 
     def __imul__(self, a):
         self._mat *= a
+        self._creation_info = None
         return self
 
     def __iadd__(self, M):
         assert M.domain is self.domain
         assert M.codomain is self.codomain
 
         if isinstance(M, WeightedMassOperator):
+            self._creation_info = None
             self._mat += M._mat
             return self
 
         elif isinstance(M, LinearOperator):
+            self._creation_info = None
             self._mat += M
             return self
 
@@ -2383,10 +2479,12 @@ def __isub__(self, M):
         assert M.codomain is self.codomain
 
         if isinstance(M, WeightedMassOperator):
+            self._creation_info = None
             self._mat -= M._mat
             return self
 
         elif isinstance(M, LinearOperator):
+            self._creation_info = None
             self._mat -= M
             return self
 

src/struphy/feec/psydac_derham.py

@@ -563,6 +563,11 @@ class Derham:
     domain : Domain, optional
         The Struphy domain object for evaluating the mapping F : [0, 1]^3 --> R^3 and the corresponding metric coefficients.
 
+    domain_decomposition : DomainDecomposition, optional
+        Prescribed MPI decomposition of the elements, e.g. ``fine_derham.domain_decomposition.coarsen(...)``
+        for an aligned multigrid level. Must match ``grid.num_elements`` and the periodicity from ``options.bcs``,
+        and be built on ``comm``. If None (default), it is computed from ``comm`` and ``grid.mpi_dims_mask``.
+
     Notes
     -----
     The underlying base sequence is
@@ -579,6 +584,8 @@ def __init__(
         options: DerhamOptions,
         comm: MPI.Intracomm = None,
         domain: Domain = None,
+        *,
+        domain_decomposition: DomainDecomposition | None = None,
     ):
 
         # inputs
@@ -680,6 +687,7 @@ def __init__(
             comm=self.comm,
             mpi_dims_mask=mpi_dims_mask,
             use_feectools=use_feectools,
+            domain_decomposition=domain_decomposition,
         )
 
         # FEM spaces
@@ -1520,6 +1528,7 @@ def init_derham(
         comm=None,
         mpi_dims_mask: tuple[bool, bool, bool] = None,
         use_feectools: bool = True,
+        domain_decomposition: DomainDecomposition | None = None,
     ) -> DiscreteDerham:
         """Return a discrete Derham complex. Allows for the use of tiny-feectools.
 
@@ -1543,12 +1552,29 @@ def init_derham(
 
         use_feectools: bool
             Use slimmed-down fork `feectools` of Psydac.
+
+        domain_decomposition : DomainDecomposition, optional
+            Prescribed decomposition of the elements; if None, it is computed from ``comm`` and ``mpi_dims_mask``.
         """
 
         if use_feectools:
-            self._domain_decomposition = DomainDecomposition(
-                num_elements, spl_kind, comm=comm, mpi_dims_mask=mpi_dims_mask
-            )
+            if domain_decomposition is None:
+                self._domain_decomposition = DomainDecomposition(
+                    num_elements, spl_kind, comm=comm, mpi_dims_mask=mpi_dims_mask
+                )
+            else:
+                assert tuple(domain_decomposition.ncells) == tuple(num_elements), (
+                    f"{domain_decomposition.ncells = } does not match {num_elements = }."
+                )
+                assert tuple(domain_decomposition.periods) == tuple(spl_kind), (
+                    f"{domain_decomposition.periods = } does not match {spl_kind = }."
+                )
+                if domain_decomposition.comm is not None and comm is not None:
+                    # (comm is None in the decomposition when feectools runs with MockMPI)
+                    assert domain_decomposition.comm == comm, (
+                        "domain_decomposition must be built on the Derham communicator."
+                    )
+                self._domain_decomposition = domain_decomposition
 
             _derham = self._discretize_derham(
                 num_elements,

src/struphy/io/options.py

@@ -76,6 +76,7 @@ class LiteralOptions:
     OptsSymmSolver = Literal["pcg", "cg"]
     OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"]
     OptsMassPrecond = Literal["MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None]
+    OptsDiffusionPrecond = Literal["MultiGrid", "MassMatrixPreconditioner", "MassMatrixDiagonalPreconditioner", None]
     OptsSaddlePointSolver = Literal["uzawa"]
     OptsDirectSolver = Literal["SparseSolver", "ScipySparse", "InexactNPInverse", "DirectNPInverse"]
     OptsNonlinearSolver = Literal["Picard", "Newton"]