BDDC: cellwise subdomains, matfree MatIS, and user-defined primal vertices (#4757)

Author: pbrubeck

firedrake/preconditioners/bddc.py

@@ -1,3 +1,5 @@
+from itertools import repeat
+
 from firedrake.preconditioners.base import PCBase
 from firedrake.preconditioners.patch import bcdofs
 from firedrake.preconditioners.facet_split import get_restriction_indices
@@ -6,11 +8,16 @@
 from firedrake.ufl_expr import TestFunction, TrialFunction
 from firedrake.function import Function
 from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace
-from firedrake.preconditioners.fdm import tabulate_exterior_derivative
+from firedrake.preconditioners.fdm import broken_function, tabulate_exterior_derivative
 from firedrake.preconditioners.hiptmair import curl_to_grad
-from ufl import H1, H2, inner, dx, JacobianDeterminant
+from firedrake.parloops import par_loop, INC, READ
+from firedrake.utils import cached_property
+from firedrake.bcs import DirichletBC
+from firedrake.mesh import Submesh
+from ufl import Form, H1, H2, JacobianDeterminant, dx, inner, replace
+from finat.ufl import BrokenElement
+from pyop2.mpi import COMM_SELF
 from pyop2.utils import as_tuple
-import gem
 import numpy
 
 __all__ = ("BDDCPC",)
@@ -23,6 +30,8 @@ class BDDCPC(PCBase):
 
     Internally, this PC creates a PETSc PCBDDC object that can be controlled by
     the options:
+    - ``'bddc_cellwise'`` to set up a MatIS on cellwise subdomains if P.type == python,
+    - ``'bddc_matfree'`` to set up a matrix-free MatIS if A.type == python,
     - ``'bddc_pc_bddc_neumann'`` to set sub-KSPs on subdomains excluding corners,
     - ``'bddc_pc_bddc_dirichlet'`` to set sub-KSPs on subdomain interiors,
     - ``'bddc_pc_bddc_coarse'`` to set the coarse solver KSP.
@@ -34,36 +43,59 @@ class BDDCPC(PCBase):
     - ``'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``.
+    - ``'primal_markers'`` a Function marking degrees of freedom of the solution space to be included in the
+    coarse space. Any nonzero value is counted as a marked degree of freedom.
+    If a DG(0) Function is provided, then all degrees of freedom on the cell are marked.
+    Alternatively, ``'primal_markers'`` can be a list of the global degrees of freedom to
+    be supplied directly to ``PETSc.PC.setBDDCPrimalVerticesIS``.
     """
 
     _prefix = "bddc_"
 
     def initialize(self, pc):
-        # Get context from pc
-        _, P = pc.getOperators()
-        dm = pc.getDM()
-        self.prefix = (pc.getOptionsPrefix() or "") + self._prefix
+        prefix = (pc.getOptionsPrefix() or "") + self._prefix
 
+        dm = pc.getDM()
         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)
         bddcpc.incrementTabLevel(1, parent=pc)
-        bddcpc.setOptionsPrefix(self.prefix)
-        bddcpc.setOperators(*pc.getOperators())
+        bddcpc.setOptionsPrefix(prefix)
         bddcpc.setType(PETSc.PC.Type.BDDC)
 
         opts = PETSc.Options(bddcpc.getOptionsPrefix())
+        matfree = opts.getBool("matfree", False)
+
+        # Set operators
+        assemblers = []
+        A, P = pc.getOperators()
+        if P.type == "python":
+            # Reconstruct P as MatIS
+            cellwise = opts.getBool("cellwise", False)
+            P, assembleP = create_matis(P, "aij", cellwise=cellwise)
+            assemblers.append(assembleP)
+
+        if P.type != "is":
+            raise ValueError(f"Expecting P to be either 'matfree' or 'is', not {P.type}.")
+
+        if A.type == "python" and matfree:
+            # Reconstruct A as MatIS
+            A, assembleA = create_matis(A, "matfree", cellwise=P.getISAllowRepeated())
+            assemblers.append(assembleA)
+        bddcpc.setOperators(A, P)
+        self.assemblers = assemblers
+
         # 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)):
+        is_h1h2 = V.ufl_element().sobolev_space in {H1, H2}
+        if "pc_bddc_use_local_mat_graph" not in opts and (not is_h1h2 or not V.finat_element.has_pointwise_dual_basis):
             opts["pc_bddc_use_local_mat_graph"] = False
 
-        # Handle boundary dofs
+        # Get context from DM
         ctx = get_appctx(dm)
+
+        # Handle boundary dofs
         bcs = tuple(ctx._problem.dirichlet_bcs())
         mesh = V.mesh().unique()
         if mesh.extruded and not mesh.extruded_periodic:
@@ -85,16 +117,17 @@ def initialize(self, pc):
         bddcpc.setBDDCNeumannBoundaries(neu_bndr)
 
         appctx = self.get_appctx(pc)
-        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)
+            degree = max(as_tuple(V.ufl_element().degree()))
+            variant = V.ufl_element().variant()
+            W = VectorFunctionSpace(mesh, "Lagrange", degree, variant=variant)
+            coords = Function(W).interpolate(mesh.coordinates)
             bddcpc.setCoordinates(coords.dat.data_ro.repeat(V.block_size, axis=0))
 
-        tdim = V.mesh().topological_dimension
+        tdim = mesh.topological_dimension
         if tdim >= 2 and V.finat_element.formdegree == tdim-1:
             allow_repeated = P.getISAllowRepeated()
             get_divergence = appctx.get("get_divergence_mat", get_divergence_mat)
@@ -116,14 +149,22 @@ def initialize(self, pc):
                 grad_kwargs = dict()
             bddcpc.setBDDCDiscreteGradient(*grad_args, **grad_kwargs)
 
+        # Set the user-defined primal (coarse) degrees of freedom
+        primal_markers = appctx.get("primal_markers")
+        if primal_markers is not None:
+            primal_indices = get_primal_indices(V, primal_markers)
+            primal_is = PETSc.IS().createGeneral(primal_indices.astype(PETSc.IntType), comm=pc.comm)
+            bddcpc.setBDDCPrimalVerticesIS(primal_is)
+
         bddcpc.setFromOptions()
         self.pc = bddcpc
 
     def view(self, pc, viewer=None):
         self.pc.view(viewer=viewer)
 
     def update(self, pc):
-        pass
+        for c in self.assemblers:
+            c()
 
     def apply(self, pc, x, y):
         self.pc.apply(x, y)
@@ -132,6 +173,104 @@ def applyTranspose(self, pc, x, y):
         self.pc.applyTranspose(x, y)
 
 
+class BrokenDirichletBC(DirichletBC):
+    def __init__(self, bc):
+        self.bc = bc
+        V = bc.function_space().broken_space()
+        g = bc._original_arg
+        super().__init__(V, g, bc.sub_domain)
+
+    @cached_property
+    def nodes(self):
+        u = Function(self.bc.function_space())
+        self.bc.set(u, 1)
+        u = broken_function(u.function_space(), val=u.dat)
+        return numpy.flatnonzero(u.dat.data)
+
+
+def create_matis(Amat, local_mat_type, cellwise=False):
+    from firedrake.assemble import get_assembler
+
+    def local_mesh(mesh):
+        key = "local_submesh"
+        cache = mesh._shared_data_cache["local_submesh_cache"]
+        try:
+            return cache[key]
+        except KeyError:
+            if mesh.comm.size > 1:
+                submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, reorder=False, comm=COMM_SELF)
+            else:
+                submesh = None
+            return cache.setdefault(key, submesh)
+
+    def local_space(V, cellwise):
+        mesh = local_mesh(V.mesh().unique())
+        element = BrokenElement(V.ufl_element()) if cellwise else None
+        return V.reconstruct(mesh=mesh, element=element)
+
+    def local_argument(arg, cellwise):
+        return arg.reconstruct(function_space=local_space(arg.function_space(), cellwise))
+
+    def local_integral(it):
+        extra_domain_integral_type_map = dict(it.extra_domain_integral_type_map())
+        extra_domain_integral_type_map[it.ufl_domain()] = it.integral_type()
+        return it.reconstruct(domain=local_mesh(it.ufl_domain()),
+                              extra_domain_integral_type_map=extra_domain_integral_type_map)
+
+    def local_bc(bc, cellwise):
+        V = bc.function_space()
+        Vsub = local_space(V, False)
+        sub_domain = list(bc.sub_domain)
+        if "on_boundary" in sub_domain:
+            sub_domain.remove("on_boundary")
+            sub_domain.extend(V.mesh().unique().exterior_facets.unique_markers)
+
+        valid_markers = Vsub.mesh().unique().exterior_facets.unique_markers
+        sub_domain = list(set(sub_domain) & set(valid_markers))
+        bc = bc.reconstruct(V=Vsub, g=0, sub_domain=sub_domain)
+        if cellwise:
+            bc = BrokenDirichletBC(bc)
+        return bc
+
+    def local_to_global_map(V, cellwise):
+        u = Function(V)
+        u.dat.data_wo[:] = numpy.arange(*V.dof_dset.layout_vec.getOwnershipRange())
+
+        Vsub = local_space(V, False)
+        usub = Function(Vsub).assign(u)
+        if cellwise:
+            usub = broken_function(usub.function_space(), val=usub.dat)
+        indices = usub.dat.data_ro.astype(PETSc.IntType)
+        return PETSc.LGMap().create(indices, comm=V.comm)
+
+    assert Amat.type == "python"
+    ctx = Amat.getPythonContext()
+    form = ctx.a
+    bcs = ctx.bcs
+
+    local_form = replace(form, {arg: local_argument(arg, cellwise) for arg in form.arguments()})
+    local_form = Form(list(map(local_integral, local_form.integrals())))
+    local_bcs = tuple(map(local_bc, bcs, repeat(cellwise)))
+
+    assembler = get_assembler(local_form, bcs=local_bcs, mat_type=local_mat_type)
+    tensor = assembler.assemble()
+
+    rmap = local_to_global_map(form.arguments()[0].function_space(), cellwise)
+    cmap = local_to_global_map(form.arguments()[1].function_space(), cellwise)
+
+    Amatis = PETSc.Mat().createIS(Amat.getSizes(), comm=Amat.getComm())
+    Amatis.setISAllowRepeated(cellwise)
+    Amatis.setLGMap(rmap, cmap)
+    Amatis.setISLocalMat(tensor.petscmat)
+    Amatis.setUp()
+    Amatis.assemble()
+
+    def update():
+        assembler.assemble(tensor=tensor)
+        Amatis.assemble()
+    return Amatis, update
+
+
 def get_restricted_dofs(V, domain):
     W = FunctionSpace(V.mesh(), V.ufl_element()[domain])
     indices = get_restriction_indices(V, W)
@@ -166,7 +305,7 @@ def get_discrete_gradient(V):
     nsp = VectorSpaceBasis([basis])
     nsp.orthonormalize()
     gradient.setNullSpace(nsp.nullspace())
-    if not is_lagrange(Q.finat_element):
+    if not Q.finat_element.has_pointwise_dual_basis:
         vdofs = get_restricted_dofs(Q, "vertex")
         gradient.compose('_elements_corners', vdofs)
 
@@ -176,23 +315,25 @@ def get_discrete_gradient(V):
     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
+def get_primal_indices(V, primal_markers):
+    if isinstance(primal_markers, Function):
+        marker_space = primal_markers.function_space()
+        if marker_space == V:
+            markers = primal_markers
+        elif marker_space.finat_element.space_dimension() == 1:
+            shapes = (V.finat_element.space_dimension(), V.block_size)
+            domain = "{[i,j]: 0 <= i < %d and 0 <= j < %d}" % shapes
+            instructions = """
+            for i, j
+                w[i,j] = w[i,j] + t[0]
+            end
+            """
+            markers = Function(V)
+            par_loop((domain, instructions), dx, {"w": (markers, INC), "t": (primal_markers, READ)})
+        else:
+            raise ValueError(f"Expecting markers in either {V.ufl_element()} or DG(0).")
+        primal_indices = numpy.flatnonzero(markers.dat.data >= 1E-12)
+        primal_indices += V.dof_dset.layout_vec.getOwnershipRange()[0]
+    else:
+        primal_indices = numpy.asarray(primal_markers, dtype=PETSc.IntType)
+    return primal_indices