Support MatIS and BDDC (#4405)

* Support mat_type="is"

* Support mat_type="nest", sub_mat_type="is"

* Propagate sub_mat_type and only set sub_mat_type on the diagonal blocks

* refine BDDC customization, use callables for gradient and divergence

* Simplify coordinates handling of bddcpc driver

* tabulate_exterior_derivative as MatIS

* FDMPC: Support other variants

* bddc: attach constants of the H(grad) space to discrete gradient

* add mat_type kawrg to local_to_global_map

* Test BDDC for H(curl) and H(div)

---------

Co-authored-by: Stefano Zampini <stefano.zampini@gmail.com>

Author: pbrubeck

firedrake/assemble.py

@@ -364,7 +364,9 @@ def allocate(self):
             else:
                 test, trial = self._form.arguments()
                 sparsity = ExplicitMatrixAssembler._make_sparsity(test, trial, self._mat_type, self._sub_mat_type, self.maps_and_regions)
-                return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType, options_prefix=self._options_prefix)
+                return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType,
+                                     sub_mat_type=self._sub_mat_type,
+                                     options_prefix=self._options_prefix)
         else:
             raise NotImplementedError("Only implemented for rank = 2 and diagonal = False")
 
@@ -1324,12 +1326,12 @@ def _get_mat_type(mat_type, sub_mat_type, arguments):
                for arg in arguments
                for V in arg.function_space()):
             mat_type = "nest"
-    if mat_type not in {"matfree", "aij", "baij", "nest", "dense"}:
+    if mat_type not in {"matfree", "aij", "baij", "nest", "dense", "is"}:
         raise ValueError(f"Unrecognised matrix type, '{mat_type}'")
     if sub_mat_type is None:
         sub_mat_type = parameters.parameters["default_sub_matrix_type"]
-    if sub_mat_type not in {"aij", "baij"}:
-        raise ValueError(f"Invalid submatrix type, '{sub_mat_type}' (not 'aij' or 'baij')")
+    if sub_mat_type not in {"aij", "baij", "is"}:
+        raise ValueError(f"Invalid submatrix type, '{sub_mat_type}' (not 'aij', 'baij', or 'is')")
     return mat_type, sub_mat_type
 
 
@@ -1373,6 +1375,7 @@ def allocate(self):
                                                           self._sub_mat_type,
                                                           self._make_maps_and_regions())
         return matrix.Matrix(self._form, self._bcs, self._mat_type, sparsity, ScalarType,
+                             sub_mat_type=self._sub_mat_type,
                              options_prefix=self._options_prefix,
                              fc_params=self._form_compiler_params)
 
@@ -1490,8 +1493,10 @@ def _apply_bc(self, tensor, bc, u=None):
             # Set diagonal entries on bc nodes to 1 if the current
             # block is on the matrix diagonal and its index matches the
             # index of the function space the bc is defined on.
+            if op2tensor.handle.getType() == "is":
+                # Flag the entire matrix as assembled before indexing the diagonal block
+                op2tensor.handle.assemble()
             op2tensor[index, index].set_local_diagonal_entries(bc.nodes, idx=component, diag_val=self.weight)
-
             # Handle off-diagonal block involving real function space.
             # "lgmaps" is correctly constructed in _matrix_arg, but
             # is ignored by PyOP2 in this case.
@@ -2067,16 +2072,18 @@ def collect_lgmaps(self):
                 row_bcs, col_bcs = self._filter_bcs(i, j)
                 # the tensor is already indexed
                 rlgmap, clgmap = self._tensor.local_to_global_maps
-                rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap)
-                clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap)
+                mat_type = self._tensor.handle.getType()
+                rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap, mat_type=mat_type)
+                clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap, mat_type=mat_type)
                 return ((rlgmap, clgmap),)
             else:
                 lgmaps = []
                 for i, j in self.get_indicess():
                     row_bcs, col_bcs = self._filter_bcs(i, j)
                     rlgmap, clgmap = self._tensor[i, j].local_to_global_maps
-                    rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap)
-                    clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap)
+                    mat_type = self._tensor[i, j].handle.getType()
+                    rlgmap = self.test_function_space[i].local_to_global_map(row_bcs, rlgmap, mat_type=mat_type)
+                    clgmap = self.trial_function_space[j].local_to_global_map(col_bcs, clgmap, mat_type=mat_type)
                     lgmaps.append((rlgmap, clgmap))
                 return tuple(lgmaps)
         else:

firedrake/functionspaceimpl.py

@@ -856,10 +856,31 @@ def boundary_nodes(self, sub_domain):
         return self._shared_data.boundary_nodes(self, sub_domain)
 
     @PETSc.Log.EventDecorator()
-    def local_to_global_map(self, bcs, lgmap=None):
+    def local_to_global_map(self, bcs, lgmap=None, mat_type=None):
         r"""Return a map from process local dof numbering to global dof numbering.
 
-        If BCs is provided, mask out those dofs which match the BC nodes."""
+        Parameters
+        ----------
+        bcs: [firedrake.bcs.BCBase]
+            If provided, mask out those dofs which match the BC nodes.
+        lgmap: PETSc.LGMap
+            The base local-to-global map, which might be partially masked.
+        mat_type: str
+            The matrix assembly type. This is required as different matrix types
+            handle the LGMap differently for MixedFunctionSpace.
+
+        Note
+        ----
+            For a :func:`.VectorFunctionSpace` or :func:`.TensorFunctionSpace` the returned
+            LGMap will be the scalar one, unless the bcs are imposed on a particular component.
+            For a :class:`MixedFunctionSpace` the returned LGMap is unblocked,
+            unless mat_type == "is".
+
+        Returns
+        -------
+        PETSc.LGMap
+            A local-to-global map with masked BC dofs.
+        """
         # Caching these things is too complicated, since it depends
         # not just on the bcs, but also the parent space, and anything
         # this space has been recursively split out from [e.g. inside
@@ -884,10 +905,16 @@ def local_to_global_map(self, bcs, lgmap=None):
                 bsize = lgmap.getBlockSize()
                 assert bsize == self.block_size
         else:
-            # MatBlock case, LGMap is already unrolled.
-            indices = lgmap.block_indices.copy()
+            # MatBlock case, the LGMap is implementation dependent
             bsize = lgmap.getBlockSize()
-            unblocked = True
+            assert bsize == self.block_size
+            if mat_type == "is":
+                indices = lgmap.indices.copy()
+                unblocked = False
+            else:
+                # LGMap is already unrolled
+                indices = lgmap.block_indices.copy()
+                unblocked = True
         nodes = []
         for bc in bcs:
             if bc.function_space().component is not None:
@@ -997,7 +1024,7 @@ def __hash__(self):
         return hash((self.mesh(), self.dof_dset, self.ufl_element(),
                      self.boundary_set))
 
-    def local_to_global_map(self, bcs, lgmap=None):
+    def local_to_global_map(self, bcs, lgmap=None, mat_type=None):
         return lgmap or self.dof_dset.lgmap
 
     def collapse(self):
@@ -1204,7 +1231,7 @@ def exterior_facet_node_map(self):
         function space nodes."""
         return op2.MixedMap(s.exterior_facet_node_map() for s in self)
 
-    def local_to_global_map(self, bcs):
+    def local_to_global_map(self, bcs, lgmap=None, mat_type=None):
         r"""Return a map from process local dof numbering to global dof numbering.
 
         If BCs is provided, mask out those dofs which match the BC nodes."""
@@ -1452,7 +1479,7 @@ def top_nodes(self):
         ":class:`RealFunctionSpace` objects have no bottom nodes."
         return None
 
-    def local_to_global_map(self, bcs, lgmap=None):
+    def local_to_global_map(self, bcs, lgmap=None, mat_type=None):
         assert len(bcs) == 0
         return None
 

firedrake/preconditioners/bddc.py

@@ -4,8 +4,13 @@
 from firedrake.petsc import PETSc
 from firedrake.dmhooks import get_function_space, get_appctx
 from firedrake.ufl_expr import TestFunction, TrialFunction
-from ufl import curl, div, HCurl, HDiv, inner, dx
+from firedrake.function import Function
+from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace
+from firedrake.preconditioners.fdm import tabulate_exterior_derivative
+from firedrake.preconditioners.hiptmair import curl_to_grad
+from ufl import H1, H2, inner, dx, JacobianDeterminant
 from pyop2.utils import as_tuple
+import gem
 import numpy
 
 __all__ = ("BDDCPC",)
@@ -22,17 +27,13 @@ class BDDCPC(PCBase):
     - ``'bddc_pc_bddc_dirichlet'`` to set sub-KSPs on subdomain interiors,
     - ``'bddc_pc_bddc_coarse'`` to set the coarse solver KSP.
 
-    This PC also inspects optional arguments supplied in the application context:
-    - ``'discrete_gradient'`` for problems in H(curl), this sets the arguments
-    (a Mat tabulating the gradient of the auxiliary H1 space) and
+    This PC also inspects optional callbacks supplied in the application context:
+    - ``'get_discrete_gradient'`` for 3D problems in H(curl), this is a callable that
+    provide the arguments (a Mat tabulating the gradient of the auxiliary H1 space) and
     keyword arguments supplied to ``PETSc.PC.setBDDCDiscreteGradient``.
-    - ``'divergence_mat'`` for 3D problems in H(div), this sets the Mat with the
-    assembled bilinear form testing the divergence against an L2 space.
-
-    Notes
-    -----
-    Currently the Mat type IS is only supported by FDMPC.
-
+    - ``'get_divergence_mat'`` for problems in H(div) (resp. 2D H(curl)), this is
+    provide the arguments (a Mat with the assembled bilinear form testing the divergence
+    (curl) against an L2 space) and keyword arguments supplied to ``PETSc.PC.setDivergenceMat``.
     """
 
     _prefix = "bddc_"
@@ -44,6 +45,8 @@ def initialize(self, pc):
         self.prefix = (pc.getOptionsPrefix() or "") + self._prefix
 
         V = get_function_space(dm)
+        variant = V.ufl_element().variant()
+        sobolev_space = V.ufl_element().sobolev_space
 
         # Create new PC object as BDDC type
         bddcpc = PETSc.PC().create(comm=pc.comm)
@@ -53,13 +56,17 @@ def initialize(self, pc):
         bddcpc.setType(PETSc.PC.Type.BDDC)
 
         opts = PETSc.Options(bddcpc.getOptionsPrefix())
-        if V.ufl_element().variant() == "fdm" and "pc_bddc_use_local_mat_graph" not in opts:
-            # Disable computation of disconected components of subdomain interfaces
+        # Do not use CSR of local matrix to define dofs connectivity unless requested
+        # Using the CSR only makes sense for H1/H2 problems
+        is_h1h2 = sobolev_space in [H1, H2]
+        if "pc_bddc_use_local_mat_graph" not in opts and (not is_h1h2 or not is_lagrange(V.finat_element)):
             opts["pc_bddc_use_local_mat_graph"] = False
 
+        # Handle boundary dofs
         ctx = get_appctx(dm)
-        bcs = tuple(ctx._problem.bcs)
-        if V.extruded:
+        bcs = tuple(ctx._problem.dirichlet_bcs())
+        mesh = V.mesh().unique()
+        if mesh.extruded and not mesh.extruded_periodic:
             boundary_nodes = numpy.unique(numpy.concatenate(list(map(V.boundary_nodes, ("on_boundary", "top", "bottom")))))
         else:
             boundary_nodes = V.boundary_nodes("on_boundary")
@@ -69,46 +76,44 @@ def initialize(self, pc):
             dir_nodes = numpy.unique(numpy.concatenate([bcdofs(bc, ghost=False) for bc in bcs]))
         neu_nodes = numpy.setdiff1d(boundary_nodes, dir_nodes)
 
-        V.dof_dset.lgmap.apply(dir_nodes, result=dir_nodes)
+        dir_nodes = V.dof_dset.lgmap.apply(dir_nodes)
         dir_bndr = PETSc.IS().createGeneral(dir_nodes, comm=pc.comm)
         bddcpc.setBDDCDirichletBoundaries(dir_bndr)
 
-        V.dof_dset.lgmap.apply(neu_nodes, result=neu_nodes)
+        neu_nodes = V.dof_dset.lgmap.apply(neu_nodes)
         neu_bndr = PETSc.IS().createGeneral(neu_nodes, comm=pc.comm)
         bddcpc.setBDDCNeumannBoundaries(neu_bndr)
 
         appctx = self.get_appctx(pc)
-        sobolev_space = V.ufl_element().sobolev_space
+        degree = max(as_tuple(V.ufl_element().degree()))
+
+        # Set coordinates only if corner selection is requested
+        # There's no API to query from PC
+        if "pc_bddc_corner_selection" in opts:
+            W = VectorFunctionSpace(V.mesh(), "Lagrange", degree, variant=variant)
+            coords = Function(W).interpolate(V.mesh().coordinates)
+            bddcpc.setCoordinates(coords.dat.data_ro.repeat(V.block_size, axis=0))
 
         tdim = V.mesh().topological_dimension
-        degree = max(as_tuple(V.ufl_element().degree()))
         if tdim >= 2 and V.finat_element.formdegree == tdim-1:
-            B = appctx.get("divergence_mat", None)
-            if B is None:
-                from firedrake.assemble import assemble
-                d = {HCurl: curl, HDiv: div}[sobolev_space]
-                Q = V.reconstruct(family="DG", degree=degree-1)
-                b = inner(d(TrialFunction(V)), TestFunction(Q)) * dx(degree=2*(degree-1))
-                B = assemble(b, mat_type="matfree")
-            bddcpc.setBDDCDivergenceMat(B.petscmat)
-        elif sobolev_space == HCurl:
-            gradient = appctx.get("discrete_gradient", None)
-            if gradient is None:
-                from firedrake.preconditioners.fdm import tabulate_exterior_derivative
-                from firedrake.preconditioners.hiptmair import curl_to_grad
-                Q = V.reconstruct(element=curl_to_grad(V.ufl_element()))
-                gradient = tabulate_exterior_derivative(Q, V)
-                corners = get_vertex_dofs(Q)
-                gradient.compose('_elements_corners', corners)
+            allow_repeated = P.getISAllowRepeated()
+            get_divergence = appctx.get("get_divergence_mat", get_divergence_mat)
+            divergence = get_divergence(V, mat_type="is", allow_repeated=allow_repeated)
+            try:
+                div_args, div_kwargs = divergence
+            except ValueError:
+                div_args = (divergence,)
+                div_kwargs = dict()
+            bddcpc.setBDDCDivergenceMat(*div_args, **div_kwargs)
+
+        elif tdim >= 3 and V.finat_element.formdegree == 1:
+            get_gradient = appctx.get("get_discrete_gradient", get_discrete_gradient)
+            gradient = get_gradient(V)
+            try:
+                grad_args, grad_kwargs = gradient
+            except ValueError:
                 grad_args = (gradient,)
-                grad_kwargs = {'order': degree}
-            else:
-                try:
-                    grad_args, grad_kwargs = gradient
-                except ValueError:
-                    grad_args = (gradient,)
-                    grad_kwargs = dict()
-
+                grad_kwargs = dict()
             bddcpc.setBDDCDiscreteGradient(*grad_args, **grad_kwargs)
 
         bddcpc.setFromOptions()
@@ -127,9 +132,67 @@ def applyTranspose(self, pc, x, y):
         self.pc.applyTranspose(x, y)
 
 
-def get_vertex_dofs(V):
-    W = V.reconstruct(element=restrict(V.ufl_element(), "vertex"))
+def get_restricted_dofs(V, domain):
+    W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), domain))
     indices = get_restriction_indices(V, W)
-    V.dof_dset.lgmap.apply(indices, result=indices)
-    vertex_dofs = PETSc.IS().createGeneral(indices, comm=V.comm)
-    return vertex_dofs
+    indices = V.dof_dset.lgmap.apply(indices)
+    return PETSc.IS().createGeneral(indices, comm=V.comm)
+
+
+def get_divergence_mat(V, mat_type="is", allow_repeated=False):
+    from firedrake import assemble
+    degree = max(as_tuple(V.ufl_element().degree()))
+    Q = TensorFunctionSpace(V.mesh(), "DG", 0, variant=f"integral({degree-1})", shape=V.value_shape[:-1])
+    B = tabulate_exterior_derivative(V, Q, mat_type=mat_type, allow_repeated=allow_repeated)
+
+    Jdet = JacobianDeterminant(V.mesh())
+    s = assemble(inner(TrialFunction(Q)*(1/Jdet), TestFunction(Q))*dx(degree=0), diagonal=True)
+    with s.dat.vec as svec:
+        B.diagonalScale(svec, None)
+    return (B,), {}
+
+
+def get_discrete_gradient(V):
+    from firedrake import Constant
+    from firedrake.nullspace import VectorSpaceBasis
+
+    Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element()))
+    gradient = tabulate_exterior_derivative(Q, V)
+    basis = Function(Q)
+    try:
+        basis.interpolate(Constant(1))
+    except NotImplementedError:
+        basis.project(Constant(1))
+    nsp = VectorSpaceBasis([basis])
+    nsp.orthonormalize()
+    gradient.setNullSpace(nsp.nullspace())
+    if not is_lagrange(Q.finat_element):
+        vdofs = get_restricted_dofs(Q, "vertex")
+        gradient.compose('_elements_corners', vdofs)
+
+    degree = max(as_tuple(Q.ufl_element().degree()))
+    grad_args = (gradient,)
+    grad_kwargs = {'order': degree}
+    return grad_args, grad_kwargs
+
+
+def is_lagrange(finat_element):
+    """Returns whether finat_element.dual_basis consists only of point evaluation dofs."""
+    try:
+        Q, ps = finat_element.dual_basis
+    except NotImplementedError:
+        return False
+    # Inspect the weight matrix
+    # Lagrange elements have gem.Delta as the only terminal nodes
+    children = [Q]
+    while children:
+        nodes = []
+        for c in children:
+            if isinstance(c, gem.Delta):
+                pass
+            elif isinstance(c, gem.gem.Terminal):
+                return False
+            else:
+                nodes.extend(c.children)
+        children = nodes
+    return True