From 9e9cb958f37375000dfdac99ee9063df578320f7 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 1 Dec 2025 16:02:46 +0000 Subject: [PATCH 01/20] Submesh on subcommunicator --- firedrake/cython/dmcommon.pyx | 13 +++++++++---- firedrake/cython/petschdr.pxi | 1 - firedrake/mesh.py | 17 +++++++++++++++-- firedrake/mg/mesh.py | 2 +- 4 files changed, 25 insertions(+), 8 deletions(-) diff --git a/firedrake/cython/dmcommon.pyx b/firedrake/cython/dmcommon.pyx index 999d08f82e..2623f769e7 100644 --- a/firedrake/cython/dmcommon.pyx +++ b/firedrake/cython/dmcommon.pyx @@ -3819,7 +3819,8 @@ def submesh_create(PETSc.DM dm, PetscInt subdim, label_name, PetscInt label_value, - PetscBool ignore_label_halo): + PetscBool ignore_label_halo, + PETSc.Comm comm): """Create submesh. Parameters @@ -3834,12 +3835,12 @@ def submesh_create(PETSc.DM dm, Value in the label ignore_label_halo : bool If labeled points in the halo are ignored. + comm: PETSc.Comm + An optional sub-communicator to define the submesh. """ cdef: - PETSc.DM subdm = PETSc.DMPlex() PETSc.DMLabel label, temp_label - PETSc.SF ownership_transfer_sf = PETSc.SF() char *temp_label_name = "firedrake_submesh_temp_label" PetscInt pStart, pEnd, p, i, stratum_size PETSc.PetscIS stratum_is = NULL @@ -3863,7 +3864,11 @@ def submesh_create(PETSc.DM dm, CHKERR(ISRestoreIndices(stratum_is, &stratum_indices)) CHKERR(ISDestroy(&stratum_is)) # Make submesh using temp_label. - CHKERR(DMPlexFilter(dm.dm, temp_label.dmlabel, label_value, ignore_label_halo, PETSC_TRUE, &ownership_transfer_sf.sf, &subdm.dm)) + subdm, ownership_transfer_sf = dm.filter(label=temp_label, + value=label_value, + ignoreHalo=ignore_label_halo, + sanitizeSubMesh=PETSC_TRUE, + comm=comm) # Destroy temp_label. dm.removeLabel(temp_label_name) subdm.removeLabel(temp_label_name) diff --git a/firedrake/cython/petschdr.pxi b/firedrake/cython/petschdr.pxi index fa9d10c300..d7b1478ca0 100644 --- a/firedrake/cython/petschdr.pxi +++ b/firedrake/cython/petschdr.pxi @@ -76,7 +76,6 @@ cdef extern from "petscdmplex.h" nogil: PetscErrorCode DMPlexLabelComplete(PETSc.PetscDM, PETSc.PetscDMLabel) PetscErrorCode DMPlexDistributeOverlap(PETSc.PetscDM,PetscInt,PETSc.PetscSF*,PETSc.PetscDM*) - PetscErrorCode DMPlexFilter(PETSc.PetscDM,PETSc.PetscDMLabel,PetscInt,PetscBool,PetscBool,PETSc.PetscSF*,PETSc.PetscDM*) PetscErrorCode DMPlexGetSubpointIS(PETSc.PetscDM,PETSc.PetscIS*) PetscErrorCode DMPlexGetSubpointMap(PETSc.PetscDM,PETSc.PetscDMLabel*) PetscErrorCode DMPlexSetSubpointMap(PETSc.PetscDM,PETSc.PetscDMLabel) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index df48b11acf..d06c9f0c05 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4731,7 +4731,7 @@ def SubDomainData(geometric_expr): return op2.Subset(m.cell_set, indices) -def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None): +def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo=False, comm=None): """Construct a submesh from a given mesh. Parameters @@ -4746,6 +4746,10 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None): Name of the label to search ``subdomain_id`` in. name : str Name of the submesh. + ignore_halo : bool + Excludes the halo from the submesh. + comm : PETSc.Comm + An optional sub-communicator to define the submesh. Returns ------- @@ -4815,7 +4819,15 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None): elif subdim == dim - 1: label_name = dmcommon.FACE_SETS_LABEL name = name or _generate_default_submesh_name(mesh.name) - subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, False) + if comm is None: + subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm) + else: + # FIXME + subplex, _ = plex.filter(label=plex.getLabel(label_name), + value=subdomain_id, + sanitizeSubMesh=True, + ignoreHalo=ignore_halo, + comm=comm) subplex.setName(_generate_default_mesh_topology_name(name)) if subplex.getDimension() != subdim: raise RuntimeError(f"Found subplex dim ({subplex.getDimension()}) != expected ({subdim})") @@ -4823,6 +4835,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None): subplex, submesh_parent=mesh, name=name, + comm=subplex.comm, distribution_parameters={ "partition": False, "overlap_type": (DistributedMeshOverlapType.NONE, 0), diff --git a/firedrake/mg/mesh.py b/firedrake/mg/mesh.py index a646189776..a88bf42504 100644 --- a/firedrake/mg/mesh.py +++ b/firedrake/mg/mesh.py @@ -136,7 +136,7 @@ def MeshHierarchy(mesh, refinement_levels, # This is algorithmically guaranteed. dm_cell_type, = mesh.dm_cell_types tdim = mesh.topology_dm.getDimension() - cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "celltype", dm_cell_type, True) + cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "celltype", dm_cell_type, True, None) cdm.removeLabel("pyop2_core") cdm.removeLabel("pyop2_owned") cdm.removeLabel("pyop2_ghost") From a1cca747b5c54049798f95871efae38bc459a5b8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 1 Dec 2025 17:02:07 +0000 Subject: [PATCH 02/20] test --- firedrake/mesh.py | 3 ++- tests/firedrake/submesh/test_submesh_comm.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 tests/firedrake/submesh/test_submesh_comm.py diff --git a/firedrake/mesh.py b/firedrake/mesh.py index d06c9f0c05..7dcc3329b8 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4819,7 +4819,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= elif subdim == dim - 1: label_name = dmcommon.FACE_SETS_LABEL name = name or _generate_default_submesh_name(mesh.name) - if comm is None: + if comm is None or comm == mesh.comm: subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm) else: # FIXME @@ -4836,6 +4836,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= submesh_parent=mesh, name=name, comm=subplex.comm, + reorder=False, distribution_parameters={ "partition": False, "overlap_type": (DistributedMeshOverlapType.NONE, 0), diff --git a/tests/firedrake/submesh/test_submesh_comm.py b/tests/firedrake/submesh/test_submesh_comm.py new file mode 100644 index 0000000000..62dded4987 --- /dev/null +++ b/tests/firedrake/submesh/test_submesh_comm.py @@ -0,0 +1,18 @@ +import pytest +import numpy as np +from firedrake import * +from firedrake.petsc import PETSc + + +@pytest.mark.parallel([1, 3]) +def test_submesh_comm_self(): + comm = PETSc.COMM_SELF + nx = 4 + mesh = UnitSquareMesh(nx, nx, quadrilateral=True) + tdim = mesh.topological_dimension + submesh = Submesh(mesh, tdim, None, ignore_halo=True, comm=comm) + + assert submesh.submesh_parent is mesh + assert submesh.comm.size == 1 + assert mesh.cell_set.size == submesh.cell_set.size + assert np.allclose(mesh.coordinates.dat.data_ro, submesh.coordinates.dat.data_ro) From 586c545390ff0a7afd7922cd4c13deeb1f68fef6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 1 Dec 2025 17:04:08 +0000 Subject: [PATCH 03/20] PETSc version --- firedrake/__init__.py | 2 +- setup.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/__init__.py b/firedrake/__init__.py index 70dec1d6ca..a9a2a300c9 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -3,7 +3,7 @@ # the specific version, here we are more permissive. This is to catch the # case where users don't update their PETSc for a really long time or # accidentally install a too-new release that isn't yet supported. -PETSC_SUPPORTED_VERSIONS = ">=3.23.0" +PETSC_SUPPORTED_VERSIONS = ">=3.25.0" def init_petsc(): diff --git a/setup.py b/setup.py index d601613b53..7ad41e34b8 100644 --- a/setup.py +++ b/setup.py @@ -20,7 +20,7 @@ # Ensure that the PETSc getting linked against is compatible -petsctools.init(version_spec=">=3.23.0") +petsctools.init(version_spec=">=3.25.0") import petsc4py From 9e69d998d67f3712fe72e6671f9d316eeffd32b2 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 1 Dec 2025 17:06:54 +0000 Subject: [PATCH 04/20] BDDC: create_matis --- firedrake/preconditioners/bddc.py | 169 +++++++++++++++++++++++++++--- 1 file changed, 152 insertions(+), 17 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 24e47320f2..ff8e6a9d92 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 restrict, get_restriction_indices @@ -6,9 +9,14 @@ 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, unghosted_lgmap 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.utils import as_tuple import gem import numpy @@ -23,6 +31,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 +44,59 @@ 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) + lgmap = unghosted_lgmap(V, V.dof_dset.lgmap, allow_repeated=cellwise) + P, assembleP = create_matis(P, lgmap, lgmap, "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, *P.getLGMap(), "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,75 @@ def applyTranspose(self, pc, x, y): self.pc.applyTranspose(x, y) +class BrokenDirichletBC(DirichletBC): + def __init__(self, V, g, bc): + self.bc = bc + 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 = broken_function(u.function_space(), val=u.dat) + return numpy.flatnonzero(u_broken.dat.data_ro) + + +def create_matis(Amat, rmap, cmap, 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, 0, ignore_halo=True, comm=PETSc.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_bc(bc, cellwise): + V = local_space(bc.function_space(), cellwise) + if cellwise: + return BrokenDirichletBC(V, 0, bc) + else: + return bc.reconstruct(V=V, g=0) + + def local_coefficient(c): + V = local_space(c.function_space(), False) + return Function(V, val=c.dat.data) + + assert Amat.type == "python" + ctx = Amat.getPythonContext() + form = ctx.a + bcs = ctx.bcs + + repl = {arg: arg.reconstruct(function_space=local_space(arg.function_space(), cellwise)) + for arg in form.arguments()} + repl.update({c: local_coefficient(c) for c in form.coefficients()}) + local_form = replace(form, repl) + + new_mesh = local_form.arguments()[0].function_space().mesh() + local_form = Form([it.reconstruct(domain=new_mesh) for it in 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() + + 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(), restrict(V.ufl_element(), domain)) indices = get_restriction_indices(V, W) @@ -176,6 +287,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: From ed2f1abd1c6384bb808e0e91fa11c6b293dbd5e1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 1 Dec 2025 17:35:20 +0000 Subject: [PATCH 05/20] reorder kwarg --- firedrake/__init__.py | 2 +- firedrake/mesh.py | 8 +++++--- setup.py | 2 +- tests/firedrake/submesh/test_submesh_comm.py | 5 ++--- 4 files changed, 9 insertions(+), 8 deletions(-) diff --git a/firedrake/__init__.py b/firedrake/__init__.py index a9a2a300c9..2a1390b53c 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -3,7 +3,7 @@ # the specific version, here we are more permissive. This is to catch the # case where users don't update their PETSc for a really long time or # accidentally install a too-new release that isn't yet supported. -PETSC_SUPPORTED_VERSIONS = ">=3.25.0" +PETSC_SUPPORTED_VERSIONS = ">=3.24.1" def init_petsc(): diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 7dcc3329b8..e9634b9bbb 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4731,7 +4731,7 @@ def SubDomainData(geometric_expr): return op2.Subset(m.cell_set, indices) -def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo=False, comm=None): +def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo=False, reorder=True, comm=None): """Construct a submesh from a given mesh. Parameters @@ -4747,7 +4747,9 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= name : str Name of the submesh. ignore_halo : bool - Excludes the halo from the submesh. + Whether to exclude the halo from the submesh. + reorder : bool + Whether to reorder the mesh entities. comm : PETSc.Comm An optional sub-communicator to define the submesh. @@ -4836,7 +4838,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= submesh_parent=mesh, name=name, comm=subplex.comm, - reorder=False, + reorder=reorder, distribution_parameters={ "partition": False, "overlap_type": (DistributedMeshOverlapType.NONE, 0), diff --git a/setup.py b/setup.py index 7ad41e34b8..9651870b79 100644 --- a/setup.py +++ b/setup.py @@ -20,7 +20,7 @@ # Ensure that the PETSc getting linked against is compatible -petsctools.init(version_spec=">=3.25.0") +petsctools.init(version_spec=">=3.24.1") import petsc4py diff --git a/tests/firedrake/submesh/test_submesh_comm.py b/tests/firedrake/submesh/test_submesh_comm.py index 62dded4987..1429401edc 100644 --- a/tests/firedrake/submesh/test_submesh_comm.py +++ b/tests/firedrake/submesh/test_submesh_comm.py @@ -8,9 +8,8 @@ def test_submesh_comm_self(): comm = PETSc.COMM_SELF nx = 4 - mesh = UnitSquareMesh(nx, nx, quadrilateral=True) - tdim = mesh.topological_dimension - submesh = Submesh(mesh, tdim, None, ignore_halo=True, comm=comm) + mesh = UnitSquareMesh(nx, nx, quadrilateral=True, reorder=False) + submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, comm=comm, reorder=False) assert submesh.submesh_parent is mesh assert submesh.comm.size == 1 From 26e505cd6edc27defa899c44044eb0bc1bfbaf6e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 2 Dec 2025 17:41:58 +0000 Subject: [PATCH 06/20] Optional comm kwarg --- firedrake/cython/dmcommon.pyx | 4 ++-- firedrake/mg/mesh.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/firedrake/cython/dmcommon.pyx b/firedrake/cython/dmcommon.pyx index 2623f769e7..b6ea9b6306 100644 --- a/firedrake/cython/dmcommon.pyx +++ b/firedrake/cython/dmcommon.pyx @@ -3820,7 +3820,7 @@ def submesh_create(PETSc.DM dm, label_name, PetscInt label_value, PetscBool ignore_label_halo, - PETSc.Comm comm): + comm=None): """Create submesh. Parameters @@ -3835,7 +3835,7 @@ def submesh_create(PETSc.DM dm, Value in the label ignore_label_halo : bool If labeled points in the halo are ignored. - comm: PETSc.Comm + comm : PETSc.Comm | None An optional sub-communicator to define the submesh. """ diff --git a/firedrake/mg/mesh.py b/firedrake/mg/mesh.py index a88bf42504..a646189776 100644 --- a/firedrake/mg/mesh.py +++ b/firedrake/mg/mesh.py @@ -136,7 +136,7 @@ def MeshHierarchy(mesh, refinement_levels, # This is algorithmically guaranteed. dm_cell_type, = mesh.dm_cell_types tdim = mesh.topology_dm.getDimension() - cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "celltype", dm_cell_type, True, None) + cdm = dmcommon.submesh_create(mesh.topology_dm, tdim, "celltype", dm_cell_type, True) cdm.removeLabel("pyop2_core") cdm.removeLabel("pyop2_owned") cdm.removeLabel("pyop2_ghost") From 6fb20e0fa6e2e0c39749a31ccb1132f295183373 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 7 Dec 2025 17:27:13 +0000 Subject: [PATCH 07/20] FML: fix subtraction and support BaseForm --- firedrake/fml/form_manipulation_language.py | 35 +++++++++------ tests/firedrake/regression/test_fml.py | 48 ++++++++++++++++++++- 2 files changed, 68 insertions(+), 15 deletions(-) diff --git a/firedrake/fml/form_manipulation_language.py b/firedrake/fml/form_manipulation_language.py index 8eebc94e49..dafba9d99a 100644 --- a/firedrake/fml/form_manipulation_language.py +++ b/firedrake/fml/form_manipulation_language.py @@ -93,7 +93,7 @@ class Term(object): __slots__ = ["form", "labels"] - def __init__(self, form: ufl.Form, label_dict: Mapping = None): + def __init__(self, form: ufl.BaseForm, label_dict: Mapping = None): """ Parameters @@ -216,6 +216,9 @@ def __mul__( __rmul__ = __mul__ + def __neg__(self): + return -1 * self + def __truediv__( self, other: Union[float, Constant, ufl.algebra.Product] @@ -274,7 +277,7 @@ def __init__(self, *terms: Sequence[Term]): def __add__( self, - other: Union[ufl.Form, Term, "LabelledForm"] + other: Union[ufl.BaseForm, Term, "LabelledForm"] ) -> "LabelledForm": """Add a form, term or labelled form to this labelled form. @@ -289,12 +292,12 @@ def __add__( A labelled form containing the terms. """ - if isinstance(other, ufl.Form): - return LabelledForm(*self, Term(other)) - elif type(other) is Term: + if type(other) is Term: return LabelledForm(*self, other) elif type(other) is LabelledForm: return LabelledForm(*self, *other) + elif isinstance(other, ufl.BaseForm): + return LabelledForm(*self, Term(other)) elif other is None: return self else: @@ -304,7 +307,7 @@ def __add__( def __sub__( self, - other: Union[ufl.Form, Term, "LabelledForm"] + other: Union[ufl.BaseForm, Term, "LabelledForm"] ) -> "LabelledForm": """Subtract a form, term or labelled form from this labelled form. @@ -320,15 +323,18 @@ def __sub__( """ if type(other) is Term: - return LabelledForm(*self, Constant(-1.)*other) + return LabelledForm(*self, -other) elif type(other) is LabelledForm: - return LabelledForm(*self, *[Constant(-1.)*t for t in other]) + return LabelledForm(*self, *[-t for t in other]) elif other is None: return self else: # Make new Term for other and subtract it return LabelledForm(*self, Term(Constant(-1.)*other)) + def __rsub__(self, other): + return other + (-self) + def __mul__( self, other: Union[float, Constant, ufl.algebra.Product] @@ -371,6 +377,9 @@ def __truediv__( __rmul__ = __mul__ + def __neg__(self): + return -1 * self + def __iter__(self) -> Sequence: """Iterable of the terms in the labelled form.""" return iter(self.terms) @@ -427,7 +436,7 @@ def label_map( return new_labelled_form @property - def form(self) -> ufl.Form: + def form(self) -> ufl.BaseForm: """Provide the whole form from the labelled form. Raises @@ -437,7 +446,7 @@ def form(self) -> ufl.Form: Returns ------- - ufl.Form + ufl.BaseForm The whole form corresponding to all the terms. """ @@ -479,7 +488,7 @@ def __init__( def __call__( self, - target: Union[ufl.Form, Term, LabelledForm], + target: Union[ufl.BaseForm, Term, LabelledForm], value: Any = None ) -> Union[Term, LabelledForm]: """Apply the label to a form or term. @@ -494,7 +503,7 @@ def __call__( Raises ------ ValueError - If the `target` is not a ufl.Form, Term or + If the `target` is not a ufl.BaseForm, Term or LabelledForm. Returns @@ -514,7 +523,7 @@ def __call__( self.value = self.default_value if isinstance(target, LabelledForm): return LabelledForm(*(self(t, value) for t in target.terms)) - elif isinstance(target, ufl.Form): + elif isinstance(target, ufl.BaseForm): return LabelledForm(Term(target, {self.label: self.value})) elif isinstance(target, Term): new_labels = target.labels.copy() diff --git a/tests/firedrake/regression/test_fml.py b/tests/firedrake/regression/test_fml.py index f639c8e870..f6c02d55a1 100644 --- a/tests/firedrake/regression/test_fml.py +++ b/tests/firedrake/regression/test_fml.py @@ -9,9 +9,12 @@ PeriodicUnitSquareMesh, FunctionSpace, Constant, MixedFunctionSpace, TestFunctions, Function, split, inner, dx, SpatialCoordinate, as_vector, pi, sin, div, - NonlinearVariationalProblem, NonlinearVariationalSolver + NonlinearVariationalProblem, NonlinearVariationalSolver, + UnitIntervalMesh, Cofunction, TestFunction, interpolate ) -from firedrake.fml import subject, replace_subject, keep, drop, Label +from firedrake.fml import subject, replace_subject, keep, drop, Label, LabelledForm + +import pytest def test_fml(): @@ -121,3 +124,44 @@ def test_fml(): X.assign(X_np1) implicit_solver.solve() X.assign(X_np1) + + +@pytest.fixture +def V(): + mesh = UnitIntervalMesh(2) + return FunctionSpace(mesh, "CG", 1) + + +@pytest.mark.parametrize("baseform", ("form", "cofunction", "interpolate")) +def test_fml_baseform(V, baseform): + if baseform == "form": + form = inner(1, TestFunction(V))*dx + + elif baseform == "cofunction": + form = Cofunction(V.dual()) + form.assign(42) + + elif baseform == "interpolate": + W = V.reconstruct(degree=2) + w = Cofunction(W.dual()) + w.assign(42) + form = interpolate(TestFunction(V), w) + + # Test that we can wrap this BaseForm with a Label + has_labels = lambda term: len(term.labels) > 0 + label = Label("label") + F = label(form) + assert isinstance(F, LabelledForm) + assert F.label_map(has_labels, map_if_true=keep, map_if_false=drop).form == form + + # Test that we can linearly combine this BaseForm with a LabelledForm + Fcombs = [form + label(form), + label(form) + form, + label(form) - form] + for F2 in Fcombs: + assert isinstance(F2, LabelledForm) + assert F2.label_map(has_labels, map_if_true=keep, map_if_false=drop).form == form + + F3 = form - label(form) + assert isinstance(F3, LabelledForm) + assert F3.label_map(has_labels, map_if_true=keep, map_if_false=drop).form == -form From 4a91deea28c840adb360a05dd06db721fbb620c8 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 14:33:03 +0000 Subject: [PATCH 08/20] Cleanup --- firedrake/__init__.py | 4 +++- firedrake/mesh.py | 34 +++++++++++++++++----------------- 2 files changed, 20 insertions(+), 18 deletions(-) diff --git a/firedrake/__init__.py b/firedrake/__init__.py index 2a1390b53c..98c36f15f8 100644 --- a/firedrake/__init__.py +++ b/firedrake/__init__.py @@ -3,8 +3,10 @@ # the specific version, here we are more permissive. This is to catch the # case where users don't update their PETSc for a really long time or # accidentally install a too-new release that isn't yet supported. -PETSC_SUPPORTED_VERSIONS = ">=3.24.1" +PETSC_SUPPORTED_VERSIONS = ">=3.24.2" +# TODO RELEASE +# PETSC_SUPPORTED_VERSIONS = ">=3.25" def init_petsc(): import os diff --git a/firedrake/mesh.py b/firedrake/mesh.py index e9634b9bbb..663e6cc7b9 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4740,8 +4740,9 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= Parent mesh (`MeshGeometry`). subdim : int Topological dimension of the submesh. - subdomain_id : int + subdomain_id : int | None Subdomain ID representing the submesh. + `None` defines the submesh owned by the sub-communicator. label_name : str Name of the label to search ``subdomain_id`` in. name : str @@ -4813,23 +4814,22 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= raise NotImplementedError("Can not create a submesh of a ``VertexOnlyMesh``") plex = mesh.topology_dm dim = plex.getDimension() - if subdim not in [dim, dim - 1]: - raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") - if label_name is None: - if subdim == dim: - label_name = dmcommon.CELL_SETS_LABEL - elif subdim == dim - 1: - label_name = dmcommon.FACE_SETS_LABEL - name = name or _generate_default_submesh_name(mesh.name) - if comm is None or comm == mesh.comm: - subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm) + if subdomain_id is None: + # Filter the plex with PETSc's default label (cells owned by comm) + if subdim != dim: + raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") + subplex, _ = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=comm) else: - # FIXME - subplex, _ = plex.filter(label=plex.getLabel(label_name), - value=subdomain_id, - sanitizeSubMesh=True, - ignoreHalo=ignore_halo, - comm=comm) + if subdim not in [dim, dim - 1]: + raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") + if label_name is None: + if subdim == dim: + label_name = dmcommon.CELL_SETS_LABEL + elif subdim == dim - 1: + label_name = dmcommon.FACE_SETS_LABEL + subplex = dmcommon.submesh_create(plex, subdim, label_name, subdomain_id, ignore_halo, comm=comm) + + name = name or _generate_default_submesh_name(mesh.name) subplex.setName(_generate_default_mesh_topology_name(name)) if subplex.getDimension() != subdim: raise RuntimeError(f"Found subplex dim ({subplex.getDimension()}) != expected ({subdim})") From fb522090a2d49e2f9b1449eeac8349aee9451423 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 16:16:10 +0000 Subject: [PATCH 09/20] fixup --- firedrake/mesh.py | 8 +++++--- tests/firedrake/submesh/test_submesh_comm.py | 9 ++++----- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 109b4d109e..ba0882f618 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -3012,7 +3012,7 @@ def make_mesh_from_coordinates(coordinates, name, tolerance=0.5): def make_mesh_from_mesh_topology(topology, name, tolerance=0.5): - """Make mesh from tpology. + """Make mesh from topology. Parameters ---------- @@ -4820,11 +4820,13 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= if subdomain_id is None: # Filter the plex with PETSc's default label (cells owned by comm) if label_name is not None: - raise ValueError(f"Submesh on sub-communicator needs label_name == None.") + raise ValueError("subdomain_id == None requires label_name == None.") if subdim != dim: raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") subplex, _ = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=comm) else: + if comm is not None and comm != mesh.comm: + raise NotImplementedError("Submesh on subcommunicator not implemented on cell subsets.") if subdim not in [dim, dim - 1]: raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") if label_name is None: @@ -4842,7 +4844,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= subplex, submesh_parent=mesh, name=name, - comm=subplex.comm, + comm=subplex.comm.tompi4py(), reorder=reorder, distribution_parameters={ "partition": False, diff --git a/tests/firedrake/submesh/test_submesh_comm.py b/tests/firedrake/submesh/test_submesh_comm.py index 1429401edc..670bae7aa1 100644 --- a/tests/firedrake/submesh/test_submesh_comm.py +++ b/tests/firedrake/submesh/test_submesh_comm.py @@ -1,17 +1,16 @@ import pytest import numpy as np from firedrake import * -from firedrake.petsc import PETSc @pytest.mark.parallel([1, 3]) def test_submesh_comm_self(): - comm = PETSc.COMM_SELF + comm = COMM_SELF + subdomain_id = None nx = 4 mesh = UnitSquareMesh(nx, nx, quadrilateral=True, reorder=False) - submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, comm=comm, reorder=False) - + submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, comm=comm, reorder=False) assert submesh.submesh_parent is mesh assert submesh.comm.size == 1 - assert mesh.cell_set.size == submesh.cell_set.size + assert submesh.cell_set.size == mesh.cell_set.size assert np.allclose(mesh.coordinates.dat.data_ro, submesh.coordinates.dat.data_ro) From fdae167491af3bac19a14705f4d211c73f757bad Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 16:59:10 +0000 Subject: [PATCH 10/20] fixup --- firedrake/preconditioners/bddc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index ff8e6a9d92..ee381baf64 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -197,7 +197,7 @@ def local_mesh(mesh): return cache[key] except KeyError: if mesh.comm.size > 1: - submesh = Submesh(mesh, mesh.topological_dimension, 0, ignore_halo=True, comm=PETSc.COMM_SELF) + submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, comm=PETSc.COMM_SELF) else: submesh = None return cache.setdefault(key, submesh) From c125fd3708aeff98e21d8e1c1acab1eb053d3d1d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 17:02:03 +0000 Subject: [PATCH 11/20] fix --- firedrake/mesh.py | 1 - 1 file changed, 1 deletion(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 5e0350d512..59fde76119 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4824,7 +4824,6 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= if subdim != dim: raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") subplex, _ = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=comm) - comm = subplex.comm.tompi4py() else: if comm is not None and comm != mesh.comm: raise NotImplementedError("Submesh on subcommunicator not implemented on cell subsets.") From 2f2e77d30d9a170b7f89ed6bcfd05c212aedcc95 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 17:31:11 +0000 Subject: [PATCH 12/20] add assemble test --- firedrake/mesh.py | 1 + tests/firedrake/submesh/test_submesh_comm.py | 24 ++++++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 59fde76119..5e0350d512 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4824,6 +4824,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= if subdim != dim: raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") subplex, _ = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=comm) + comm = subplex.comm.tompi4py() else: if comm is not None and comm != mesh.comm: raise NotImplementedError("Submesh on subcommunicator not implemented on cell subsets.") diff --git a/tests/firedrake/submesh/test_submesh_comm.py b/tests/firedrake/submesh/test_submesh_comm.py index 21bc336b16..5a3ff8fc94 100644 --- a/tests/firedrake/submesh/test_submesh_comm.py +++ b/tests/firedrake/submesh/test_submesh_comm.py @@ -14,3 +14,27 @@ def test_create_submesh_comm_self(): assert submesh.comm.size == 1 assert submesh.cell_set.size == mesh.cell_set.size assert np.allclose(mesh.coordinates.dat.data_ro, submesh.coordinates.dat.data_ro) + + +@pytest.mark.parallel([1, 3]) +def test_assemble_submesh_comm_self(): + comm = COMM_SELF + subdomain_id = None + nx = 6 + ny = 5 + x = -np.cos(np.linspace(0, np.pi, nx)) + y = -np.cos(np.linspace(0, np.pi, ny)) + mesh = TensorRectangleMesh(x, y) + submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, comm=comm) + + Vsub = FunctionSpace(submesh, "DG", 0) + Asub = assemble(inner(TrialFunction(Vsub), TestFunction(Vsub))*dx) + + V = FunctionSpace(submesh, "DG", 0) + A = assemble(inner(TrialFunction(V), TestFunction(V))*dx) + + i0, j0, v0 = A.petscmat.getValuesCSR() + i1, j1, v1 = Asub.petscmat.getValuesCSR() + assert np.array_equal(i0, i1) + assert np.array_equal(j0, j1) + assert np.allclose(v0, v1) From 99610109fa0f59bee0252f153c8a27708e0267be Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 17:42:38 +0000 Subject: [PATCH 13/20] add assemble test --- firedrake/mesh.py | 1 + tests/firedrake/submesh/test_submesh_comm.py | 27 ++++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 59fde76119..5e0350d512 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4824,6 +4824,7 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= if subdim != dim: raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") subplex, _ = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=comm) + comm = subplex.comm.tompi4py() else: if comm is not None and comm != mesh.comm: raise NotImplementedError("Submesh on subcommunicator not implemented on cell subsets.") diff --git a/tests/firedrake/submesh/test_submesh_comm.py b/tests/firedrake/submesh/test_submesh_comm.py index 21bc336b16..e6483f0bf0 100644 --- a/tests/firedrake/submesh/test_submesh_comm.py +++ b/tests/firedrake/submesh/test_submesh_comm.py @@ -14,3 +14,30 @@ def test_create_submesh_comm_self(): assert submesh.comm.size == 1 assert submesh.cell_set.size == mesh.cell_set.size assert np.allclose(mesh.coordinates.dat.data_ro, submesh.coordinates.dat.data_ro) + + +@pytest.mark.parallel([1, 3]) +def test_assemble_submesh_comm_self(): + comm = COMM_SELF + subdomain_id = None + nx = 6 + ny = 5 + px = -np.cos(np.linspace(0, np.pi, nx)) + py = -np.cos(np.linspace(0, np.pi, ny)) + mesh = TensorRectangleMesh(px, py, reorder=False) + submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, comm=comm, reorder=False) + + x = SpatialCoordinate(mesh) + V = FunctionSpace(mesh, "DG", 0) + c = Function(V).interpolate(1 + dot(x, x)) + A = assemble(inner(TrialFunction(V) * c, TestFunction(V))*dx) + + Vsub = FunctionSpace(submesh, "DG", 0) + Asub = assemble(inner(TrialFunction(Vsub) * c, TestFunction(Vsub))*dx) + + i0, j0, v0 = A.petscmat.getValuesCSR() + j0 -= V.dof_dset.layout_vec.getOwnershipRange()[0] + i1, j1, v1 = Asub.petscmat.getValuesCSR() + assert np.array_equal(i0, i1) + assert np.array_equal(j0, j1) + assert np.allclose(v0, v1) From 825d2ef1ff9e4dd002b4a42a3f001b060e8a0f67 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 17:42:38 +0000 Subject: [PATCH 14/20] add assemble test --- tests/firedrake/submesh/test_submesh_comm.py | 27 ++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/tests/firedrake/submesh/test_submesh_comm.py b/tests/firedrake/submesh/test_submesh_comm.py index 21bc336b16..e6483f0bf0 100644 --- a/tests/firedrake/submesh/test_submesh_comm.py +++ b/tests/firedrake/submesh/test_submesh_comm.py @@ -14,3 +14,30 @@ def test_create_submesh_comm_self(): assert submesh.comm.size == 1 assert submesh.cell_set.size == mesh.cell_set.size assert np.allclose(mesh.coordinates.dat.data_ro, submesh.coordinates.dat.data_ro) + + +@pytest.mark.parallel([1, 3]) +def test_assemble_submesh_comm_self(): + comm = COMM_SELF + subdomain_id = None + nx = 6 + ny = 5 + px = -np.cos(np.linspace(0, np.pi, nx)) + py = -np.cos(np.linspace(0, np.pi, ny)) + mesh = TensorRectangleMesh(px, py, reorder=False) + submesh = Submesh(mesh, mesh.topological_dimension, subdomain_id, ignore_halo=True, comm=comm, reorder=False) + + x = SpatialCoordinate(mesh) + V = FunctionSpace(mesh, "DG", 0) + c = Function(V).interpolate(1 + dot(x, x)) + A = assemble(inner(TrialFunction(V) * c, TestFunction(V))*dx) + + Vsub = FunctionSpace(submesh, "DG", 0) + Asub = assemble(inner(TrialFunction(Vsub) * c, TestFunction(Vsub))*dx) + + i0, j0, v0 = A.petscmat.getValuesCSR() + j0 -= V.dof_dset.layout_vec.getOwnershipRange()[0] + i1, j1, v1 = Asub.petscmat.getValuesCSR() + assert np.array_equal(i0, i1) + assert np.array_equal(j0, j1) + assert np.allclose(v0, v1) From 6134381af35232f62cf6d3c58ee2478568b0b18e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 10 Dec 2025 18:07:21 +0000 Subject: [PATCH 15/20] use correct COMM_SELF --- firedrake/mesh.py | 1 - firedrake/preconditioners/bddc.py | 3 ++- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 5e0350d512..59fde76119 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -4824,7 +4824,6 @@ def Submesh(mesh, subdim, subdomain_id, label_name=None, name=None, ignore_halo= if subdim != dim: raise NotImplementedError(f"Found submesh dim ({subdim}) and parent dim ({dim})") subplex, _ = plex.filter(sanitizeSubMesh=True, ignoreHalo=ignore_halo, comm=comm) - comm = subplex.comm.tompi4py() else: if comm is not None and comm != mesh.comm: raise NotImplementedError("Submesh on subcommunicator not implemented on cell subsets.") diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index ee381baf64..a86e28f237 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -17,6 +17,7 @@ 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 @@ -197,7 +198,7 @@ def local_mesh(mesh): return cache[key] except KeyError: if mesh.comm.size > 1: - submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, comm=PETSc.COMM_SELF) + submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, comm=COMM_SELF) else: submesh = None return cache.setdefault(key, submesh) From 554d8756b1205f367ca9a0e274f7a7e3f90851bc Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 15 Dec 2025 15:26:25 +0000 Subject: [PATCH 16/20] Apply suggestions from code review --- firedrake/preconditioners/bddc.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index a86e28f237..5d8c295b0a 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -215,10 +215,6 @@ def local_bc(bc, cellwise): else: return bc.reconstruct(V=V, g=0) - def local_coefficient(c): - V = local_space(c.function_space(), False) - return Function(V, val=c.dat.data) - assert Amat.type == "python" ctx = Amat.getPythonContext() form = ctx.a @@ -226,7 +222,6 @@ def local_coefficient(c): repl = {arg: arg.reconstruct(function_space=local_space(arg.function_space(), cellwise)) for arg in form.arguments()} - repl.update({c: local_coefficient(c) for c in form.coefficients()}) local_form = replace(form, repl) new_mesh = local_form.arguments()[0].function_space().mesh() From 090ff8632d5477ceb7422d55c674144924e7bc51 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Mon, 15 Dec 2025 15:55:12 +0000 Subject: [PATCH 17/20] do not reorder --- firedrake/preconditioners/bddc.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 5d8c295b0a..d69895a969 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -198,7 +198,7 @@ def local_mesh(mesh): return cache[key] except KeyError: if mesh.comm.size > 1: - submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, comm=COMM_SELF) + submesh = Submesh(mesh, mesh.topological_dimension, None, ignore_halo=True, reorder=False, comm=COMM_SELF) else: submesh = None return cache.setdefault(key, submesh) @@ -220,6 +220,8 @@ def local_bc(bc, cellwise): form = ctx.a bcs = ctx.bcs + mesh = form.arguments()[0].function_space().mesh() + assert not mesh._did_reordering or len(form.coefficients()) == 0 repl = {arg: arg.reconstruct(function_space=local_space(arg.function_space(), cellwise)) for arg in form.arguments()} local_form = replace(form, repl) From 7a278125c441b9fd76214e9d19dd2448466d2842 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 6 Jan 2026 19:35:04 -0600 Subject: [PATCH 18/20] create_matis: support reordered meshes --- firedrake/preconditioners/bddc.py | 64 +++++++++++++++++++------------ 1 file changed, 40 insertions(+), 24 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 5377823f5d..23731e5e86 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -9,7 +9,7 @@ from firedrake.ufl_expr import TestFunction, TrialFunction from firedrake.function import Function from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace -from firedrake.preconditioners.fdm import broken_function, tabulate_exterior_derivative, unghosted_lgmap +from firedrake.preconditioners.fdm import broken_function, tabulate_exterior_derivative from firedrake.preconditioners.hiptmair import curl_to_grad from firedrake.parloops import par_loop, INC, READ from firedrake.utils import cached_property @@ -74,8 +74,7 @@ def initialize(self, pc): if P.type == "python": # Reconstruct P as MatIS cellwise = opts.getBool("cellwise", False) - lgmap = unghosted_lgmap(V, V.dof_dset.lgmap, allow_repeated=cellwise) - P, assembleP = create_matis(P, lgmap, lgmap, "aij", cellwise=cellwise) + P, assembleP = create_matis(P, "aij", cellwise=cellwise) assemblers.append(assembleP) if P.type != "is": @@ -83,7 +82,7 @@ def initialize(self, pc): if A.type == "python" and matfree: # Reconstruct A as MatIS - A, assembleA = create_matis(A, *P.getLGMap(), "matfree", cellwise=P.getISAllowRepeated()) + A, assembleA = create_matis(A, "matfree", cellwise=P.getISAllowRepeated()) assemblers.append(assembleA) bddcpc.setOperators(A, P) self.assemblers = assemblers @@ -176,19 +175,16 @@ def applyTranspose(self, pc, x, y): class BrokenDirichletBC(DirichletBC): - def __init__(self, V, g, bc): - self.bc = bc - super().__init__(V, g, bc.sub_domain) + def __init__(self, V, g, nodes): + self._nodes = nodes + super().__init__(V, g, nodes) @cached_property def nodes(self): - u = Function(self.bc.function_space()) - self.bc.set(u, 1) - u_broken = broken_function(u.function_space(), val=u.dat) - return numpy.flatnonzero(u_broken.dat.data_ro) + return self._nodes -def create_matis(Amat, rmap, cmap, local_mat_type, cellwise=False): +def create_matis(Amat, local_mat_type, cellwise=False): from firedrake.assemble import get_assembler def local_mesh(mesh): @@ -208,31 +204,51 @@ def local_space(V, cellwise): 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 = local_space(bc.function_space(), cellwise) + u = Function(bc.function_space()) + bc.set(u, 1) if cellwise: - return BrokenDirichletBC(V, 0, bc) - else: - return bc.reconstruct(V=V, g=0) + u = broken_function(u.function_space(), val=u.dat) + Vsub = local_space(bc.function_space(), cellwise) + usub = Function(Vsub).assign(u) + nodes = numpy.flatnonzero(usub.dat.data) + return BrokenDirichletBC(Vsub, 0, nodes) + + 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 - mesh = form.arguments()[0].function_space().mesh() - assert not mesh._did_reordering or len(form.coefficients()) == 0 - repl = {arg: arg.reconstruct(function_space=local_space(arg.function_space(), cellwise)) - for arg in form.arguments()} - local_form = replace(form, repl) - - new_mesh = local_form.arguments()[0].function_space().mesh() - local_form = Form([it.reconstruct(domain=new_mesh) for it in local_form.integrals()]) + 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) From 4a5564fa0dc65cd27725d709e74fd4b072ce4704 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 6 Jan 2026 22:03:44 -0600 Subject: [PATCH 19/20] cleanup --- firedrake/preconditioners/bddc.py | 34 ++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index 23731e5e86..fd598c591b 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -175,13 +175,18 @@ def applyTranspose(self, pc, x, y): class BrokenDirichletBC(DirichletBC): - def __init__(self, V, g, nodes): - self._nodes = nodes - super().__init__(V, g, nodes) + 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): - return self._nodes + 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): @@ -214,14 +219,19 @@ def local_integral(it): extra_domain_integral_type_map=extra_domain_integral_type_map) def local_bc(bc, cellwise): - u = Function(bc.function_space()) - bc.set(u, 1) + V = bc.function_space() + Vsub = local_space(V, False) + sub_domain = list(bc.sub_domain) + if "on_boundary" in sub_domain: + valid_markers = set(V.mesh().unique().exterior_facets.unique_markers) + valid_markers &= set(Vsub.mesh().unique().exterior_facets.unique_markers) + sub_domain.extend(valid_markers) + sub_domain.remove("on_boundary") + + bc = bc.reconstruct(V=Vsub, g=0, sub_domain=sub_domain) if cellwise: - u = broken_function(u.function_space(), val=u.dat) - Vsub = local_space(bc.function_space(), cellwise) - usub = Function(Vsub).assign(u) - nodes = numpy.flatnonzero(usub.dat.data) - return BrokenDirichletBC(Vsub, 0, nodes) + bc = BrokenDirichletBC(bc) + return bc def local_to_global_map(V, cellwise): u = Function(V) @@ -241,8 +251,8 @@ def local_to_global_map(V, cellwise): 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() From 91d078660a79fd73356d2153e19471858a198289 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 7 Jan 2026 20:21:04 -0600 Subject: [PATCH 20/20] add a test --- firedrake/preconditioners/bddc.py | 6 +++--- tests/firedrake/regression/test_bddc.py | 18 +++++++++--------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/firedrake/preconditioners/bddc.py b/firedrake/preconditioners/bddc.py index fd598c591b..3d668f486f 100644 --- a/firedrake/preconditioners/bddc.py +++ b/firedrake/preconditioners/bddc.py @@ -223,11 +223,11 @@ def local_bc(bc, cellwise): Vsub = local_space(V, False) sub_domain = list(bc.sub_domain) if "on_boundary" in sub_domain: - valid_markers = set(V.mesh().unique().exterior_facets.unique_markers) - valid_markers &= set(Vsub.mesh().unique().exterior_facets.unique_markers) - sub_domain.extend(valid_markers) 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) 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)