BDDC: cellwise subdomains, matfree MatIS, and user-defined primal vertices

firedrake/preconditioners/bddc.py

@@ -1,3 +1,6 @@
+from itertools import repeat
+from functools import partial
+
 from firedrake.preconditioners.base import PCBase
 from firedrake.preconditioners.patch import bcdofs
 from firedrake.preconditioners.facet_split import restrict, get_restriction_indices
@@ -6,9 +9,15 @@
 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, unghosted_lgmap
 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
@@ -23,6 +32,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 +45,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 to be included in the
+    coarse space. If a DG(0) Function is provided, then all degrees of freedom on the cell
+    are marked as primal. Alternatively, primal_markers can be a list with the global
+    degrees of freedom to be included in the coarse space.
     """
 
     _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)
+            lgmap = unghosted_lgmap(V, V.dof_dset.lgmap, allow_repeated=cellwise)
+            P, assembleP = create_matis(P, lgmap, lgmap, "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, *P.getLGMap(), "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]
+        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 is_lagrange(V.finat_element)):
             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 +119,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 +151,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 +175,72 @@ def applyTranspose(self, pc, x, y):
         self.pc.applyTranspose(x, y)
 
 
+class BrokenDirichletBC(DirichletBC):
+    def __init__(self, V, g, bc):
+        self.bc = bc
+        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 = broken_function(u.function_space(), val=u.dat)
+        return numpy.flatnonzero(u_broken.dat.data_ro)
+
+
+def create_matis(Amat, rmap, cmap, 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_bc(bc, cellwise):
+        V = local_space(bc.function_space(), cellwise)
+        if cellwise:
+            return BrokenDirichletBC(V, 0, bc)
+        else:
+            return bc.reconstruct(V=V, g=0)
+
+    assert Amat.type == "python"
+    ctx = Amat.getPythonContext()
+    form = ctx.a
+    bcs = ctx.bcs
+
+    mesh = form.arguments()[0].function_space().mesh()
+    assert not mesh._did_reordering or len(form.coefficients()) == 0
+    repl = {arg: arg.reconstruct(function_space=local_space(arg.function_space(), cellwise))
+            for arg in form.arguments()}
+    local_form = replace(form, repl)
+
+    new_mesh = local_form.arguments()[0].function_space().mesh()
+    local_form = Form([it.reconstruct(domain=new_mesh) for it in 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()
+
+    Amatis = PETSc.Mat().createIS(Amat.getSizes(), comm=Amat.getComm())
+    Amatis.setISAllowRepeated(cellwise)
+    Amatis.setLGMap(rmap, cmap)
+    Amatis.setISLocalMat(tensor.petscmat)
+    Amatis.setUp()
+    return Amatis, partial(assembler.assemble, tensor=tensor)
+
+
 def get_restricted_dofs(V, domain):
     W = FunctionSpace(V.mesh(), restrict(V.ufl_element(), domain))
     indices = get_restriction_indices(V, W)
@@ -176,6 +285,30 @@ def get_discrete_gradient(V):
     return grad_args, grad_kwargs
 
 
+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)
+    return primal_indices
+
+
 def is_lagrange(finat_element):
     """Returns whether finat_element.dual_basis consists only of point evaluation dofs."""
     try: