Fixes for BDDC

firedrake/preconditioners/bddc.py

@@ -15,10 +15,11 @@
 from firedrake.parloops import par_loop, INC, READ
 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 ufl import Form, H1, H2, JacobianDeterminant, div, dx, inner, replace
+from finat.ufl import BrokenElement, MixedElement
 from pyop2.mpi import COMM_SELF
 from pyop2.utils import as_tuple
+from pyop2 import op2
 import numpy
 
 __all__ = ("BDDCPC",)
@@ -33,9 +34,12 @@ class BDDCPC(PCBase):
     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_use_discrete_gradient'`` to provide the discrete gradient.
+    - ``'bddc_use_divergence_mat'`` to provide the divergence matrix.
     - ``'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.
+    - ``'bddc_pc_bddc_corner_selection'`` to geometrically select the subdomain corners.
 
     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
@@ -87,11 +91,15 @@ def initialize(self, pc):
         bddcpc.setOperators(A, P)
         self.assemblers = assemblers
 
+        # we may inject some options, we remove them after calling setFromOptions
+        rem_opts = []
+
         # 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 = 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
+            rem_opts.append("pc_bddc_use_local_mat_graph")
 
         # Get context from DM
         ctx = get_appctx(dm)
@@ -119,17 +127,24 @@ def initialize(self, pc):
 
         appctx = self.get_appctx(pc)
 
-        # Set coordinates only if corner selection is requested
+        # Set coordinates if corner selection is requested or needed
         # There's no API to query from PC
-        if "pc_bddc_corner_selection" in opts:
-            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))
-
+        entity_dofs = V.finat_element.entity_dofs()
+        vdofs = entity_dofs[min(entity_dofs)]
+        has_vertex_dofs = any(len(vdofs[v]) > 0 for v in vdofs)
+        corner_selection = opts.getBool("pc_bddc_corner_selection") if "pc_bddc_corner_selection" in opts else has_vertex_dofs
+        if corner_selection:
+            if "pc_bddc_corner_selection" not in opts:
+                opts["pc_bddc_corner_selection"] = True
+                rem_opts.append("pc_bddc_corner_selection")
+            bddcpc.setCoordinates(get_entity_coordinates(V))
+
+        # Provide extra information for H(div) and H(curl) problems
         tdim = mesh.topological_dimension
-        if tdim >= 2 and V.finat_element.formdegree == tdim-1:
+        use_divergence = opts.getBool("use_divergence_mat", tdim >= 2 and V.finat_element.formdegree == tdim-1)
+        use_gradient = opts.getBool("use_discrete_gradient", tdim >= 3 and V.finat_element.formdegree == 1)
+
+        if use_divergence:
             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)
@@ -139,8 +154,7 @@ def initialize(self, pc):
                 div_args = (divergence,)
                 div_kwargs = dict()
             bddcpc.setBDDCDivergenceMat(*div_args, **div_kwargs)
-
-        elif tdim >= 3 and V.finat_element.formdegree == 1:
+        if use_gradient:
             get_gradient = appctx.get("get_discrete_gradient", get_discrete_gradient)
             gradient = get_gradient(V)
             try:
@@ -157,7 +171,13 @@ def initialize(self, pc):
             primal_is = PETSc.IS().createGeneral(primal_indices.astype(PETSc.IntType), comm=pc.comm)
             bddcpc.setBDDCPrimalVerticesIS(primal_is)
 
+        if "pc_bddc_check_level" not in opts and "debug" in opts:
+            opts.setValue("pc_bddc_check_level", opts["debug"])
+            rem_opts.append("pc_bddc_check_level")
         bddcpc.setFromOptions()
+        for opt in rem_opts:
+            del opts[opt]
+
         self.pc = bddcpc
 
     def view(self, pc, viewer=None):
@@ -189,7 +209,7 @@ def nodes(self):
         return numpy.flatnonzero(u.dat.data)
 
 
-def create_matis(Amat, local_mat_type, cellwise=False):
+def create_matis(a, local_mat_type, cellwise=False, bcs=()):
     from firedrake.assemble import get_assembler
 
     def local_mesh(mesh):
@@ -235,7 +255,8 @@ def local_bc(bc, cellwise):
 
     def local_to_global_map(V, cellwise):
         u = Function(V)
-        u.dat.data_wo[:] = numpy.arange(*V.dof_dset.layout_vec.getOwnershipRange())
+        shp = u.dat.data_ro.shape
+        u.dat.data_wo[...] = numpy.arange(*V.dof_dset.layout_vec.getOwnershipRange()).reshape(shp)
 
         Vsub = local_space(V, False)
         usub = Function(Vsub).assign(u)
@@ -244,10 +265,18 @@ def local_to_global_map(V, cellwise):
         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
+    if isinstance(a, Form):
+        form = a
+        args = a.arguments()
+        comm = args[0].function_space().comm
+        sizes = tuple(arg.function_space().dof_dset.layout_vec.getSizes() for arg in args)
+    elif isinstance(a, PETSc.Mat):
+        assert a.type == "python"
+        ctx = a.getPythonContext()
+        form = ctx.a
+        bcs = ctx.bcs
+        comm = a.comm
+        sizes = a.getSizes()
 
     local_form = replace(form, {arg: local_argument(arg, cellwise) for arg in form.arguments()})
     local_form = Form(list(map(local_integral, local_form.integrals())))
@@ -259,7 +288,7 @@ def local_to_global_map(V, cellwise):
     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 = PETSc.Mat().createIS(sizes, comm=comm)
     Amatis.setISAllowRepeated(cellwise)
     Amatis.setLGMap(rmap, cmap)
     Amatis.setISLocalMat(tensor.petscmat)
@@ -283,12 +312,20 @@ 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)
+    if V.finat_element.complex.is_macrocell() or V.finat_element.formdegree != Q.finat_element.formdegree-1:
+        form = inner(div(TrialFunction(V)), TestFunction(Q)) * dx
+        if mat_type == "is" and allow_repeated:
+            B, _ = create_matis(form, "aij", allow_repeated)
+        else:
+            B = assemble(form, mat_type=mat_type).petscmat
+    else:
+        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,), {}
 
 
@@ -338,3 +375,95 @@ def get_primal_indices(V, primal_markers):
     else:
         primal_indices = numpy.asarray(primal_markers, dtype=PETSc.IntType)
     return primal_indices
+
+
+def get_entity_coordinates(V):
+    """
+    Return a Function on VectorFunctionSpace(mesh, V.ufl_element()) containing
+    the physical coordinates of the entity associated with each degree of freedom of V.
+    """
+    mesh = V.mesh()
+    gdim = mesh.geometric_dimension
+
+    element = V.ufl_element()
+    scalar_element = element.sub_elements[0] if isinstance(element, MixedElement) else element
+    V_target = VectorFunctionSpace(mesh, scalar_element, dim=gdim)
+    out_coords = Function(V_target)
+
+    X = mesh.coordinates.function_space()
+    coords_element = X.ufl_element()
+    if (max(as_tuple(coords_element.degree())) == 1
+            and coords_element.family() == {"Lagrange", "Q"}):
+        CG1 = X
+        cg1_coords = mesh.coordinates
+    else:
+        CG1 = VectorFunctionSpace(mesh, "Lagrange", 1, dim=gdim)
+        cg1_coords = Function(CG1).interpolate(mesh.coordinates)
+
+    V_entity_ids = V.finat_element.entity_dofs()
+    active_entities = [
+        (dim, entity)
+        for dim in V_entity_ids
+        for entity in V_entity_ids[dim]
+        if len(V_entity_ids[dim][entity]) > 0
+    ]
+
+    def flatten_entity_ids(entities, entity_ids):
+        offsets = numpy.zeros(len(entities) + 1, dtype=PETSc.IntType)
+        flatmap = []
+
+        for idx, (dim, ent_num) in enumerate(entities):
+            offsets[idx] = len(flatmap)
+            flatmap.extend(entity_ids[dim][ent_num])
+
+        offsets[-1] = len(flatmap)
+        return offsets, numpy.array(flatmap, dtype=PETSc.IntType)
+
+    cg1_closure_ids = CG1.finat_element.entity_closure_dofs()
+    cg1_offsets, cg1_flatmap = flatten_entity_ids(active_entities, cg1_closure_ids)
+    v_offsets, v_flatmap = flatten_entity_ids(active_entities, V_entity_ids)
+    num_entities = len(active_entities)
+
+    kernel_code = f"""
+    void compute_entity_coordinates(PetscScalar *coords, PetscScalar *cg1_coords) {{
+
+        const PetscInt cg1_offsets[{num_entities + 1}] = {{ {", ".join(map(str, cg1_offsets))} }};
+        const PetscInt cg1_flatmap[{len(cg1_flatmap)}] = {{ {", ".join(map(str, cg1_flatmap))} }};
+        const PetscInt v_offsets[{num_entities + 1}] = {{ {", ".join(map(str, v_offsets))} }};
+        const PetscInt v_flatmap[{len(v_flatmap)}] = {{ {", ".join(map(str, v_flatmap))} }};
+
+        // Loop over the flat entity index
+        for (PetscInt e = 0; e < {num_entities}; ++e) {{
+            PetscInt v_start = v_offsets[e];
+            PetscInt v_end = v_offsets[e + 1];
+            PetscInt cg1_start = cg1_offsets[e];
+            PetscInt cg1_end = cg1_offsets[e + 1];
+            PetscInt num_cg1_dofs = cg1_end - cg1_start;
+
+            // Compute entity barycenter
+            PetscScalar bary[{gdim}] = {{0.0}};
+            for (PetscInt j = cg1_start; j < cg1_end; ++j) {{
+                PetscInt src_dof = cg1_flatmap[j];
+                for (PetscInt c = 0; c < {gdim}; ++c) {{
+                    bary[c] += cg1_coords[src_dof * {gdim} + c];
+                }}
+            }}
+            for (PetscInt c = 0; c < {gdim}; ++c) {{
+                bary[c] /= (PetscScalar)num_cg1_dofs;
+            }}
+
+            // Assign the entity barycenter to all DOFs on the entity
+            for (PetscInt i = v_start; i < v_end; ++i) {{
+                PetscInt dest_dof = v_flatmap[i];
+                for (PetscInt c = 0; c < {gdim}; ++c) {{
+                    coords[dest_dof * {gdim} + c] = bary[c];
+                }}
+            }}
+        }}
+    }}
+    """
+    kernel = op2.Kernel(kernel_code, "compute_entity_coordinates")
+    op2.par_loop(kernel, mesh.cell_set,
+                 out_coords.dat(op2.WRITE, out_coords.cell_node_map()),
+                 cg1_coords.dat(op2.READ, cg1_coords.cell_node_map()))
+    return out_coords.dat.data.repeat(V.block_size, axis=0)