diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index f35dd689f3..366bd4413d 100644 --- a/firedrake/preconditioners/bddc.py +++ b/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 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 +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 with 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 diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index a359dff4f3..46c530c703 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/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,9 +111,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 @@ -110,8 +125,20 @@ def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, cond 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,7 +158,8 @@ 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) @@ -139,8 +167,8 @@ def solve_riesz_map(rg, mesh, family, degree, variant, bcs, cellwise=False, cond 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, 2) + 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"))