Support MatIS and BDDC

.github/workflows/core.yml

@@ -210,6 +210,7 @@ jobs:
           pip install -U pip
           pip install --group ./firedrake-repo/pyproject.toml:ci
 
+          pip install -I "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/fix/permutations"
           firedrake-clean
           pip list
 

firedrake/assemble.py

@@ -353,7 +353,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")
 
@@ -1295,12 +1297,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
 
 
@@ -1344,6 +1346,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)
 
@@ -1366,6 +1369,18 @@ def _make_sparsity(test, trial, mat_type, sub_mat_type, maps_and_regions):
         except SparsityFormatError:
             raise ValueError("Monolithic matrix assembly not supported for systems "
                              "with R-space blocks")
+
+        # TODO reconstruct dof_dset with the unghosted lgmap
+        if mat_type == "is":
+            rmap = unghosted_lgmap(sparsity._dsets[0].lgmap, test.function_space())
+            cmap = unghosted_lgmap(sparsity._dsets[1].lgmap, trial.function_space())
+            sparsity._lgmaps = (rmap, cmap)
+        elif mat_type == "nest" and sub_mat_type == "is":
+            for i, j in numpy.ndindex(sparsity.shape):
+                block = sparsity[i, j]
+                rmap = unghosted_lgmap(block._dsets[0].lgmap, test.function_space()[i])
+                cmap = unghosted_lgmap(block._dsets[1].lgmap, trial.function_space()[j])
+                block._lgmaps = (rmap, cmap)
         return sparsity
 
     def _make_maps_and_regions(self):
@@ -1461,7 +1476,6 @@ def _apply_bc(self, tensor, bc, u=None):
             # block is on the matrix diagonal and its index matches the
             # index of the function space the bc is defined on.
             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.
@@ -2177,3 +2191,49 @@ def index_function_spaces(form, indices):
             return tuple(a.ufl_function_space()[i] for i, a in zip(indices, form.arguments()))
         else:
             raise AssertionError
+
+
+def masked_lgmap(lgmap, mask, block=True):
+    if block:
+        indices = lgmap.block_indices.copy()
+        bsize = lgmap.getBlockSize()
+    else:
+        indices = lgmap.indices.copy()
+        bsize = 1
+
+    if len(mask) > 0:
+        indices[mask] = -1
+    return PETSc.LGMap().create(indices=indices, bsize=bsize, comm=lgmap.comm)
+
+
+def unghosted_lgmap(lgmap, V, block=True):
+    if block:
+        ndofs = lgmap.getBlockIndices().size
+    else:
+        ndofs = lgmap.getIndices().size
+    mask = numpy.arange(ndofs, dtype=PETSc.IntType)
+
+    mesh = V._mesh
+    mesh_dm = mesh.topology_dm
+    start, end = mesh_dm.getHeightStratum(0)
+    for i, W in enumerate(V):
+        iset = V.dof_dset.local_ises[i]
+        W_local_indices = iset.indices
+        bsize = 1 if block else iset.getBlockSize()
+        section = W.dm.getDefaultSection()
+        for seed in range(start, end):
+            # Do not loop over ghost cells
+            if mesh_dm.getLabelValue("pyop2_ghost", seed) != -1:
+                continue
+            closure, _ = mesh_dm.getTransitiveClosure(seed, useCone=True)
+            for p in closure:
+                dof = section.getDof(p)
+                if dof <= 0:
+                    continue
+                off = section.getOffset(p)
+                # Local indices within W
+                W_indices = slice(bsize * off, bsize * (off + dof))
+                mask[W_local_indices[W_indices]] = -1
+
+    mask = mask[mask > -1]
+    return masked_lgmap(lgmap, mask, block=block)

firedrake/preconditioners/bddc.py

@@ -4,8 +4,11 @@
 from firedrake.petsc import PETSc
 from firedrake.dmhooks import get_function_space, get_appctx
 from firedrake.ufl_expr import TestFunction, TrialFunction
+from firedrake.function import Function
 from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace
-from ufl import curl, div, HCurl, HDiv, inner, dx
+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 numpy
 
@@ -23,17 +26,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_"
@@ -45,6 +44,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)
@@ -54,12 +55,15 @@ 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 variant in {"fdm", "demkowicz", "demkowiczmass"}):
             opts["pc_bddc_use_local_mat_graph"] = False
 
+        # Handle boundary dofs
         ctx = get_appctx(dm)
-        bcs = tuple(ctx._problem.bcs)
+        bcs = tuple(ctx._problem.dirichlet_bcs())
         if V.extruded:
             boundary_nodes = numpy.unique(numpy.concatenate(list(map(V.boundary_nodes, ("on_boundary", "top", "bottom")))))
         else:
@@ -70,52 +74,45 @@ 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]
-                if V.shape == ():
-                    make_function_space = FunctionSpace
-                elif len(V.shape) == 1:
-                    make_function_space = VectorFunctionSpace
-                else:
-                    make_function_space = TensorFunctionSpace
-                Q = make_function_space(V.mesh(), "DG", 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 = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element()))
-                gradient = tabulate_exterior_derivative(Q, V)
-                corners = get_vertex_dofs(Q)
-                gradient.compose('_elements_corners', corners)
+            A, P = pc.getOperators()
+            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()
@@ -134,9 +131,36 @@ def applyTranspose(self, pc, x, y):
         self.pc.applyTranspose(x, y)
 
 
-def get_vertex_dofs(V):
-    W = FunctionSpace(V.mesh(), 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="aij", 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):
+    Q = FunctionSpace(V.mesh(), curl_to_grad(V.ufl_element()))
+    gradient = tabulate_exterior_derivative(Q, V)
+
+    variant = Q.ufl_element().variant()
+    if variant in {"fdm", "demkowicz", "demkowiczmass"}:
+        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

firedrake/preconditioners/facet_split.py

@@ -1,7 +1,7 @@
 from functools import partial
 from mpi4py import MPI
 from pyop2 import op2, PermutedMap
-from finat.ufl import RestrictedElement, MixedElement, TensorElement, VectorElement
+from finat.ufl import BrokenElement, RestrictedElement, MixedElement, TensorElement, VectorElement
 from firedrake.petsc import PETSc
 from firedrake.preconditioners.base import PCBase
 import firedrake.dmhooks as dmhooks
@@ -206,7 +206,9 @@ def restrict(ele, restriction_domain):
     if isinstance(ele, VectorElement):
         return type(ele)(restrict(ele._sub_element, restriction_domain), dim=ele.num_sub_elements)
     elif isinstance(ele, TensorElement):
-        return type(ele)(restrict(ele._sub_element, restriction_domain), shape=ele._shape, symmetry=ele._symmety)
+        return type(ele)(restrict(ele._sub_element, restriction_domain), shape=ele._shape, symmetry=ele._symmetry)
+    elif restriction_domain == "broken":
+        return BrokenElement(ele)
     else:
         return RestrictedElement(ele, restriction_domain)
 
@@ -239,7 +241,7 @@ def restricted_dofs(celem, felem):
     cdofs = celem.entity_dofs()
     fdofs = felem.entity_dofs()
     for dim in sorted(cdofs):
-        for entity in cdofs[dim]:
+        for entity in sorted(cdofs[dim]):
             ndofs = len(cdofs[dim][entity])
             indices[cdofs[dim][entity]] = fdofs[dim][entity][:ndofs]
     return indices
@@ -248,6 +250,9 @@ def restricted_dofs(celem, felem):
 def get_restriction_indices(V, W):
     """Return the list of dofs in the space V such that W = V[indices].
     """
+    if V.cell_node_map() is W.cell_node_map():
+        return numpy.arange(V.dof_dset.layout_vec.getSizes()[0], dtype=PETSc.IntType)
+
     vdat = V.make_dat(val=numpy.arange(V.dof_count, dtype=PETSc.IntType))
     wdats = [Wsub.make_dat(val=numpy.full((Wsub.dof_count,), -1, dtype=PETSc.IntType)) for Wsub in W]
     wdat = wdats[0] if len(W) == 1 else op2.MixedDat(wdats)