diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 7c22a556bb..3d668f486f 100644 --- a/firedrake/preconditioners/bddc.py +++ b/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 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 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,58 @@ 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) + P, assembleP = create_matis(P, "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, "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 +118,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,6 +150,13 @@ 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 @@ -123,7 +164,8 @@ 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 +174,99 @@ def applyTranspose(self, pc, x, y): self.pc.applyTranspose(x, y) +class BrokenDirichletBC(DirichletBC): + def __init__(self, bc): + self.bc = bc + V = bc.function_space().broken_space() + g = bc._original_arg + 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_function(u.function_space(), val=u.dat) + return numpy.flatnonzero(u.dat.data) + + +def create_matis(Amat, 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_argument(arg, cellwise): + return arg.reconstruct(function_space=local_space(arg.function_space(), cellwise)) + + def local_integral(it): + extra_domain_integral_type_map = dict(it.extra_domain_integral_type_map()) + extra_domain_integral_type_map[it.ufl_domain()] = it.integral_type() + return it.reconstruct(domain=local_mesh(it.ufl_domain()), + extra_domain_integral_type_map=extra_domain_integral_type_map) + + def local_bc(bc, cellwise): + V = bc.function_space() + Vsub = local_space(V, False) + sub_domain = list(bc.sub_domain) + if "on_boundary" in sub_domain: + sub_domain.remove("on_boundary") + sub_domain.extend(V.mesh().unique().exterior_facets.unique_markers) + + valid_markers = Vsub.mesh().unique().exterior_facets.unique_markers + sub_domain = list(set(sub_domain) & set(valid_markers)) + bc = bc.reconstruct(V=Vsub, g=0, sub_domain=sub_domain) + if cellwise: + bc = BrokenDirichletBC(bc) + return bc + + def local_to_global_map(V, cellwise): + u = Function(V) + u.dat.data_wo[:] = numpy.arange(*V.dof_dset.layout_vec.getOwnershipRange()) + + Vsub = local_space(V, False) + usub = Function(Vsub).assign(u) + if cellwise: + usub = broken_function(usub.function_space(), val=usub.dat) + 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 + + local_form = replace(form, {arg: local_argument(arg, cellwise) for arg in form.arguments()}) + local_form = Form(list(map(local_integral, 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() + + 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.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(), V.ufl_element()[domain]) indices = get_restriction_indices(V, W) @@ -176,6 +311,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: diff --git a/tests/firedrake/regression/test_bddc.py b/tests/firedrake/regression/test_bddc.py index 06f28425be..937e76eaf5 100644 --- a/tests/firedrake/regression/test_bddc.py +++ b/tests/firedrake/regression/test_bddc.py @@ -9,15 +9,16 @@ def rg(): return RandomGenerator(PCG64(seed=123456789)) -def bddc_params(): +def bddc_params(mat_type="is", cellwise=False): chol = { "pc_type": "cholesky", "pc_factor_mat_solver_type": DEFAULT_DIRECT_SOLVER, } sp = { - "mat_type": "is", + "mat_type": mat_type, "pc_type": "python", "pc_python_type": "firedrake.BDDCPC", + "bddc_cellwise": cellwise, "bddc_pc_bddc_neumann": chol, "bddc_pc_bddc_dirichlet": chol, "bddc_pc_bddc_coarse": chol, @@ -26,13 +27,13 @@ def bddc_params(): def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, atol=0): - sp_bddc = bddc_params() - if not cellwise: + mat_type = "matfree" if cellwise and variant != "fdm" else "is" + sp_bddc = bddc_params(mat_type=mat_type, cellwise=cellwise) + if variant != "fdm": assert not condense sp = sp_bddc elif condense: - assert variant == "fdm" sp = { "pc_type": "python", "pc_python_type": "firedrake.FacetSplitPC", @@ -52,7 +53,6 @@ def solver_parameters(cellwise=False, condense=False, variant=None, rtol=1E-10, "facet_fdm_fieldsplit_1": sp_bddc, } else: - assert variant == "fdm" sp = { "pc_type": "python", "pc_python_type": "firedrake.FDMPC", @@ -187,12 +187,12 @@ def test_bddc_aij_quad(rg, mh, family, degree, vector): @pytest.mark.parallel -@pytest.mark.parametrize("family,degree", [("CG", 3), ("N1curl", 3), ("N1div", 3)]) -def test_bddc_aij_simplex(rg, family, degree): +@pytest.mark.parametrize("family,degree,cellwise", [("CG", 3, False), ("CG", 3, True), ("N1curl", 3, False), ("N1div", 3, False)]) +def test_bddc_aij_simplex(rg, family, degree, cellwise): """Test h-dependence of condition number by measuring iteration counts""" variant = None bcs = True base = UnitCubeMesh(2, 2, 2) meshes = MeshHierarchy(base, 2) - sqrt_kappa = [solve_riesz_map(rg, m, family, degree, variant, bcs) for m in meshes] + sqrt_kappa = [solve_riesz_map(rg, m, family, degree, variant, bcs, cellwise=cellwise) for m in meshes] assert (np.diff(sqrt_kappa) <= 0.5).all(), str(sqrt_kappa)