Fixes for BDDC

firedrake/preconditioners/bddc.py

@@ -15,7 +15,7 @@
 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 ufl import Form, H1, H2, JacobianDeterminant, div, dx, inner, replace
 from finat.ufl import BrokenElement
 from pyop2.mpi import COMM_SELF
 from pyop2.utils import as_tuple
@@ -118,6 +118,10 @@ def initialize(self, pc):
         bddcpc.setBDDCNeumannBoundaries(neu_bndr)
 
         appctx = self.get_appctx(pc)
+        entity_dofs = V.finat_element.entity_dofs()
+        vdofs = entity_dofs[min(entity_dofs)]
+        if any(len(vdofs[v]) > 0 for v in vdofs):
+            bddcpc.setCoordinates(get_entity_coordinates(V))
 
         # Set coordinates only if corner selection is requested
         # There's no API to query from PC
@@ -189,7 +193,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):
@@ -244,10 +248,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 +271,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 +295,17 @@ 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():
+        form = inner(div(TrialFunction(V)), TestFunction(Q)) * dx
+        B, _ = create_matis(form, "aij", allow_repeated)
+    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 +355,33 @@ 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 with the coordinates of the entity associated to each
+    degree of freedom of a FunctionSpace.
+    """
+    mesh = V.mesh()
+    plex = mesh.topology_dm
+    num_dofs = V.dof_dset.layout_vec.getSizes()[0]
+    section = V.dm.getLocalSection()
+    vstart, vend = plex.getDepthStratum(0)
+    coords = numpy.empty((num_dofs, mesh.geometric_dimension))
+
+    if mesh.extruded:
+        P1 = VectorFunctionSpace(mesh, "Lagrange", 1)
+        plex_coords = Function(P1).interpolate(mesh.coordinates).dat.data_ro_with_halos
+    else:
+        plex_coords = plex.getCoordinatesLocal().getArray().reshape(-1, mesh.geometric_dimension)
+    for p in range(*plex.getChart()):
+        dof = section.getDof(p)
+        if dof <= 0:
+            continue
+        off = section.getOffset(p)
+        V_slice = slice(off, off + dof)
+
+        closure, _ = plex.getTransitiveClosure(p, useCone=True)
+        pverts = [q - vstart for q in closure if vstart <= q < vend]
+        coords[V_slice] = numpy.average(plex_coords[pverts], axis=0)
+
+    return coords

tests/firedrake/regression/test_bddc.py

@@ -75,7 +75,7 @@ def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10,
     return sp
 
 
-def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, condense=False, vector=False, threshold=None):
+def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, condense=False, vector=False, threshold=None, elasticity=False):
     """Solve the riesz map for a random manufactured solution and return the
        square root of the estimated condition number."""
     dirichlet_ids = []
@@ -100,9 +100,13 @@ def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, cond
         HCurl: curl,
         HDiv: div,
     }[V.ufl_element().sobolev_space]
-
     formdegree = V.finat_element.formdegree
-    if formdegree == 0:
+
+    if elasticity:
+        gamma = Constant(1E4)
+        a = (inner(grad(u) + grad(u).T, grad(v)) * dx
+             + inner(div(u) * gamma, div(v)) * dx)
+    elif formdegree == 0:
         a = inner(d(u), d(v)) * dx
     else:
         a = (inner(u, v) + inner(d(u), d(v))) * dx
@@ -111,7 +115,17 @@ def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, cond
     L = replace(a, {u: u_exact})
     bcs = [DirichletBC(V, u_exact, sub) for sub in dirichlet_ids]
     nsp = None
-    if formdegree == 0:
+    if elasticity:
+        x = SpatialCoordinate(mesh)
+        RM = [Constant(ej) for ej in np.eye(len(x))]
+        if len(x) == 2:
+            RM.append(perp(x))
+        else:
+            constants = list(RM)
+            RM.extend(cross(c, x) for c in constants)
+        nsp = VectorSpaceBasis([Function(V).interpolate(r) for r in RM])
+        nsp.orthonormalize()
+    elif formdegree == 0:
         b = np.zeros(V.value_shape)
         expr = Constant(b)
         basis = []
@@ -254,6 +268,17 @@ def test_bddc_aij_simplex(rg, family, degree, cellwise):
     assert (np.diff(sqrt_kappa) <= 0.5).all(), str(sqrt_kappa)
 
 
+@pytest.mark.parametrize("family,degree,cellwise", [("GN", 1, True), ("MTW", 1, True)])
+def test_bddc_elasticity_aij_simplex(rg, family, degree, cellwise):
+    """Test h-dependence of condition number by measuring iteration counts"""
+    variant = None
+    bcs = True
+    base = UnitSquareMesh(2, 2)
+    meshes = MeshHierarchy(base, 2)
+    sqrt_kappa = [solve_riesz_map(rg, m, family, degree, variant, bcs, cellwise=cellwise, elasticity=True) for m in meshes]
+    assert (np.diff(sqrt_kappa) <= 0.5).all(), str(sqrt_kappa)
+
+
 @pytest.mark.parallel([1, 3])
 @pytest.mark.parametrize("cellwise", (True, False))
 @pytest.mark.parametrize("local_mat_type", ("aij", "matfree"))