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
@@ -119,17 +119,21 @@ def initialize(self, pc):
 
         appctx = self.get_appctx(pc)
 
-        # Set coordinates only if corner selection is requested
+        # Set coordinates for spaces with DOFs on vertices
+        # or if corner selection is requested
         # 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))
+        corner_selection = "pc_bddc_corner_selection" in opts
+        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)
+        if corner_selection or has_vertex_dofs:
+            bddcpc.setCoordinates(get_entity_coordinates(V))
+
+        use_gradient = "use_discrete_gradient" in opts
+        use_divergence = "use_divergence_mat" in opts
 
         tdim = mesh.topological_dimension
-        if tdim >= 2 and V.finat_element.formdegree == tdim-1:
+        if use_divergence or (tdim >= 2 and V.finat_element.formdegree == tdim-1):
             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)
@@ -140,7 +144,7 @@ def initialize(self, pc):
                 div_kwargs = dict()
             bddcpc.setBDDCDivergenceMat(*div_args, **div_kwargs)
 
-        elif tdim >= 3 and V.finat_element.formdegree == 1:
+        elif use_gradient or (tdim >= 3 and V.finat_element.formdegree == 1):
             get_gradient = appctx.get("get_discrete_gradient", get_discrete_gradient)
             gradient = get_gradient(V)
             try:
@@ -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):
@@ -235,7 +239,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 +249,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 +272,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 +296,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 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 +359,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

@@ -10,7 +10,7 @@ def rg():
     return RandomGenerator(PCG64(seed=123456789))
 
 
-def bddc_params(mat_type="is", cellwise=False):
+def bddc_params(mat_type="is", cellwise=False, adaptive=False, use_divergence=False, use_gradient=False):
     chol = {
         "pc_type": "cholesky",
         "pc_factor_mat_solver_type": DEFAULT_DIRECT_SOLVER,
@@ -20,16 +20,27 @@ def bddc_params(mat_type="is", cellwise=False):
         "pc_type": "python",
         "pc_python_type": "firedrake.BDDCPC",
         "bddc_cellwise": cellwise,
+        "bddc_use_discrete_gradient": use_gradient,
+        "bddc_use_divergence_mat": use_divergence,
         "bddc_pc_bddc_neumann": chol,
         "bddc_pc_bddc_dirichlet": chol,
         "bddc_pc_bddc_coarse": chol,
+        "bddc_pc_bddc_corner_selection": True,
     }
+    if adaptive:
+        sp.update({
+            "bddc_pc_bddc_use_deluxe_scaling": None,
+            "bddc_pc_bddc_adaptive_userdefined": None,
+            "bddc_pc_bddc_deluxe_zerorows": False,
+            "bddc_pc_bddc_adaptive_threshold": 5,
+            "bddc_pc_bddc_adaptive_nmin": 1,
+        })
     return sp
 
 
-def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, atol=0):
+def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, atol=0, **kwargs):
     mat_type = "matfree" if cellwise and variant != "fdm" else "is"
-    sp_bddc = bddc_params(mat_type=mat_type, cellwise=cellwise)
+    sp_bddc = bddc_params(mat_type=mat_type, cellwise=cellwise, **kwargs)
     if variant != "fdm":
         assert not condense
         sp = sp_bddc
@@ -75,7 +86,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,18 +111,34 @@ 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
 
     u_exact = rg.uniform(V, -1, 1)
     L = replace(a, {u: u_exact})
     bcs = [DirichletBC(V, u_exact, sub) for sub in dirichlet_ids]
+
+    # Near nullspace
     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 = []
@@ -131,16 +158,17 @@ def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, cond
     problem = LinearVariationalProblem(a, L, uh, bcs=bcs)
 
     rtol = 1E-8
-    sp = solver_parameters(cellwise=cellwise, condense=condense, variant=variant, rtol=rtol)
+    sp = solver_parameters(cellwise=cellwise, condense=condense, variant=variant, rtol=rtol,
+                           use_divergence=elasticity, adaptive=elasticity)
     sp.setdefault("ksp_view_singularvalues", None)
     solver = LinearVariationalSolver(problem, near_nullspace=nsp,
                                      solver_parameters=sp, appctx=appctx)
     solver.solve()
     uerr = Function(V).assign(uh - u_exact)
     assert (assemble(a(uerr, uerr)) / assemble(a(u_exact, u_exact))) ** 0.5 < rtol
 
-    ew = solver.snes.ksp.computeEigenvalues()
-    assert min(ew) >= 1.0
+    ew = solver.snes.ksp.computeEigenvalues().real
+    assert min(ew) >= 1.0 - 1e-6
     kappa = max(abs(ew)) / min(abs(ew))
     return kappa ** 0.5
 
@@ -254,6 +282,20 @@ def test_bddc_aij_simplex(rg, family, degree, cellwise):
     assert (np.diff(sqrt_kappa) <= 0.5).all(), str(sqrt_kappa)
 
 
+@pytest.mark.parallel(3)
+@pytest.mark.parametrize("family,degree,cellwise", [("CG", 2, False), ("GN", 1, False), ("MTW", 1, False)])
+def test_bddc_elasticity_aij_simplex(rg, family, degree, cellwise):
+    """Test h-dependence of condition number by measuring iteration counts"""
+    base = UnitSquareMesh(2, 2)
+    meshes = MeshHierarchy(base, 3)
+    dim = base.topological_dimension
+    vector = (family == "CG")
+    variant = "alfeld" if family == "CG" and degree < 2*dim else None
+    bcs = True
+    sqrt_kappa = [solve_riesz_map(rg, m, family, degree, variant, bcs, cellwise=cellwise, vector=vector, elasticity=True) for m in meshes]
+    assert (np.diff(sqrt_kappa) <= 1.0).all(), str(sqrt_kappa)
+
+
 @pytest.mark.parallel([1, 3])
 @pytest.mark.parametrize("cellwise", (True, False))
 @pytest.mark.parametrize("local_mat_type", ("aij", "matfree"))