From 2e6b92db9d9ffd6016cfbcfdc8e659855f304743 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 31 Jul 2026 14:28:55 +0100 Subject: [PATCH 1/4] Copy the nodes an adaptive refinement leaves alone Adaptive refinement only touches part of a mesh, so most fine cells are exact copies of a coarse one. Detect those copied cells via a PETSc SF over the points an adaptive refine_sbr transform preserves, and have prolong/restrict copy their nodes' values directly instead of running them through the transfer kernel, which is both cheaper and exact where evaluation would otherwise be approximate. Also pads the macro-cell coarse-to-fine node map's empty slots (where a coarse cell has fewer fine children than the busiest one in the hierarchy) with a degenerate cell of zero measure, so the map stays rectangular without the kernel double-counting real contributions. --- .github/actions/install/action.yml | 8 + firedrake/cython/mgimpl.pyx | 81 +++++++ firedrake/mg/interface.py | 12 +- firedrake/mg/utils.py | 211 ++++++++++++++++++ .../multigrid/test_adaptive_multigrid.py | 134 +++++------ 5 files changed, 362 insertions(+), 84 deletions(-) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 54c1411d7f..5a08198cd4 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -163,6 +163,14 @@ runs: firedrake-clean pip list + - name: "DROP BEFORE MERGE: install FIAT from firedrakeproject/fiat#267" + shell: bash + run: | + . venv/bin/activate + pip install --verbose --no-build-isolation --no-deps --force-reinstall \ + git+https://github.com/firedrakeproject/fiat.git@pbrubeck/restricted-entity-permutations + pip list | grep -i fiat + - name: Run firedrake-check shell: bash run: | diff --git a/firedrake/cython/mgimpl.pyx b/firedrake/cython/mgimpl.pyx index 4b3a9a3e77..f39c7a8637 100644 --- a/firedrake/cython/mgimpl.pyx +++ b/firedrake/cython/mgimpl.pyx @@ -353,6 +353,87 @@ def adaptive_parent_child_cell_maps(PETSc.DM coarse_dm, return np.asarray(coarse_to_fine), np.asarray(fine_to_coarse) +@cython.boundscheck(False) +@cython.wraparound(False) +def preserved_points(PETSc.DM coarse_dm, + PETSc.Section coarse_cell_numbering, + PETSc.DM fine_dm, + PETSc.Section fine_cell_numbering, + PetscInt nfine, + np.ndarray coarse_to_fine_cells): + """Pair the points an adaptive refinement left alone with their coarse originals. + + A coarse cell that the refinement did not touch is copied into the fine + mesh, so its whole closure is copied with it: the transform preserves the + cone of every point it does not refine, and hence the plex closure of the + copied cell entry by entry. Such a cell has exactly one child, which is + what the right-padding of ``coarse_to_fine_cells`` with -1 identifies. + + :arg coarse_dm: the coarse mesh DMPlex. + :arg coarse_cell_numbering: the coarse mesh's cell numbering section. + :arg fine_dm: the adaptively refined DMPlex. + :arg fine_cell_numbering: the fine mesh's cell numbering section. + :arg nfine: the number of owned fine cells. + :arg coarse_to_fine_cells: the Firedrake-numbered coarse-to-fine cell map. + :returns: an array over the chart of ``fine_dm``, holding for each fine + point the coarse point it was copied from, or -1 if the refinement + changed it. + """ + cdef: + PetscInt ncoarse, max_children, c, i, off, child + PetscInt cStart, cEnd, pStart, pEnd, coarse_size, fine_size + PetscInt *coarse_closure = NULL + PetscInt *fine_closure = NULL + PetscInt[::1] coarse_point, fine_point, fine_to_coarse + PetscInt[:, ::1] coarse_to_fine + + coarse_to_fine = coarse_to_fine_cells + ncoarse = coarse_to_fine.shape[0] + max_children = coarse_to_fine.shape[1] + + # Both cell maps are in Firedrake numbering, so invert each mesh's cell + # numbering section to get back to the plex points the closures live on. + coarse_point = np.full(ncoarse, -1, dtype=IntType) + cStart, cEnd = coarse_dm.getHeightStratum(0) + for c in range(cStart, cEnd): + CHKERR(PetscSectionGetOffset(coarse_cell_numbering.sec, c, &off)) + if 0 <= off < ncoarse: + coarse_point[off] = c + fine_point = np.full(nfine, -1, dtype=IntType) + cStart, cEnd = fine_dm.getHeightStratum(0) + for c in range(cStart, cEnd): + CHKERR(PetscSectionGetOffset(fine_cell_numbering.sec, c, &off)) + if 0 <= off < nfine: + fine_point[off] = c + + pStart, pEnd = fine_dm.getChart() + fine_to_coarse = np.full(pEnd - pStart, -1, dtype=IntType) + for c in range(ncoarse): + child = coarse_to_fine[c, 0] + if child < 0 or (max_children > 1 and coarse_to_fine[c, 1] >= 0): + continue + if coarse_point[c] < 0 or fine_point[child] < 0: + continue + CHKERR(DMPlexGetTransitiveClosure(coarse_dm.dm, coarse_point[c], PETSC_TRUE, + &coarse_size, &coarse_closure)) + CHKERR(DMPlexGetTransitiveClosure(fine_dm.dm, fine_point[child], PETSC_TRUE, + &fine_size, &fine_closure)) + # A cell with one child that the transform nonetheless changed would + # have a closure of its own shape; leave it to the transfer kernel. + if coarse_size == fine_size: + for i in range(coarse_size): + # The closures interleave points with their orientations, and + # only points that carry the same orientation in both meshes + # order their nodes the same way. + if coarse_closure[2*i + 1] == fine_closure[2*i + 1]: + fine_to_coarse[fine_closure[2*i] - pStart] = coarse_closure[2*i] + CHKERR(DMPlexRestoreTransitiveClosure(coarse_dm.dm, coarse_point[c], PETSC_TRUE, + &coarse_size, &coarse_closure)) + CHKERR(DMPlexRestoreTransitiveClosure(fine_dm.dm, fine_point[child], PETSC_TRUE, + &fine_size, &fine_closure)) + return np.asarray(fine_to_coarse) + + # Exposition: # # These next functions compute maps from coarse mesh cells to fine diff --git a/firedrake/mg/interface.py b/firedrake/mg/interface.py index 186e0acddd..629ac97c07 100644 --- a/firedrake/mg/interface.py +++ b/firedrake/mg/interface.py @@ -102,7 +102,11 @@ def prolong(coarse, fine): for d in [coarse, coarse_coords]: d.dat.global_to_local_begin(op2.READ) d.dat.global_to_local_end(op2.READ) - op2.par_loop(kernel, fine.node_set, *kernel_args) + # An adaptive refinement leaves most of the mesh alone, and the nodes + # it preserves are copied rather than evaluated. + node_subset = utils.transfer_node_subset(Vc, Vf) + op2.par_loop(kernel, node_subset, *kernel_args) + utils.prolong_preserved_nodes(coarse, fine) if needs_quadrature: # Transfer to the actual target space @@ -184,7 +188,11 @@ def restrict(fine_dual, coarse_dual): for d in [coarse_coords]: d.dat.global_to_local_begin(op2.READ) d.dat.global_to_local_end(op2.READ) - op2.par_loop(kernel, fine_dual.node_set, *kernel_args) + # Restriction transposes prolongation, so it skips the same fine nodes + # and hands the coarse nodes they pair with their values whole. + node_subset = utils.transfer_node_subset(Vc, Vf) + op2.par_loop(kernel, node_subset, *kernel_args) + utils.restrict_preserved_nodes(fine_dual, coarse_dual) fine_dual = coarse_dual return coarse_dual diff --git a/firedrake/mg/utils.py b/firedrake/mg/utils.py index a7dff1cb2f..7b08970384 100644 --- a/firedrake/mg/utils.py +++ b/firedrake/mg/utils.py @@ -1,6 +1,8 @@ import numpy from fractions import Fraction +from mpi4py import MPI from pyop2 import op2 +from firedrake.petsc import PETSc from firedrake.utils import IntType from firedrake.functionspacedata import entity_dofs_key import finat.ufl @@ -307,6 +309,215 @@ def coarse_cell_child_count(Vc, Vf): return cache.setdefault(key, op2.Dat(dset, counts, dtype=IntType)) +def _preserved_point_sf(coarse_mesh, fine_mesh, coarse_to_fine): + """Create the SF that pairs unrefined points with their coarse originals. + + Adaptive refinement leaves some cells untouched. This SF maps each + unrefined point in ``fine_mesh`` back to the coarse point it came from. + + Parameters + ---------- + coarse_mesh : firedrake.mesh.AbstractMeshTopology + The mesh before refinement. + fine_mesh : firedrake.mesh.AbstractMeshTopology + The mesh after refinement. + coarse_to_fine : numpy.ndarray + The coarse-to-fine cell map that relates the two meshes. + + Returns + ------- + PETSc.SF + An SF with roots on the points of ``coarse_mesh`` and leaves on the + unrefined points of ``fine_mesh``. Returns `None` if refinement + changed every cell, as a uniform refinement does. + + """ + coarse_plex = coarse_mesh.topology_dm + fine_plex = fine_mesh.topology_dm + fine_to_coarse_points = impl.preserved_points( + coarse_plex, coarse_mesh._cell_numbering, + fine_plex, fine_mesh._cell_numbering, + coarse_to_fine, + ) + leaves, = numpy.nonzero(fine_to_coarse_points >= 0) + # A uniform refinement preserves no points. Every rank must agree on + # whether to build the SF at all, not just the ranks with no leaves. + if not fine_plex.comm.tompi4py().allreduce(len(leaves) > 0, op=MPI.LOR): + return None + leaves = leaves.astype(IntType) + # Refinement acts on each rank's own plex. A fine point and the coarse + # point it was copied from always live on the same rank. + remote = numpy.empty((len(leaves), 2), dtype=IntType) + remote[:, 0] = coarse_plex.comm.rank + remote[:, 1] = fine_to_coarse_points[leaves] + pStart, pEnd = coarse_plex.getChart() + point_sf = PETSc.SF().create(comm=coarse_plex.comm) + point_sf.setGraph(pEnd - pStart, leaves, remote) + return point_sf + + +def preserved_node_sf(Vc, Vf): + """Find the nodes that adaptive refinement leaves unchanged. + + An unrefined cell has the same nodes in both spaces. The transfer + operators can then copy values between them instead of evaluating them. + This is cheaper, and exact. + + Parameters + ---------- + Vc : firedrake.functionspaceimpl.WithGeometry + The coarse function space. + Vf : firedrake.functionspaceimpl.WithGeometry + The fine function space, on the next level of the same hierarchy. + + Returns + ------- + PETSc.SF + An SF with roots on the nodes of ``Vc`` and leaves on the matching + nodes of ``Vf``. Returns `None` if no nodes match. + + """ + if Vc.ufl_element() != Vf.ufl_element() or Vc.boundary_set != Vf.boundary_set: + # A space and its counterpart on the refined mesh use the same node + # layout on an unrefined cell only when the element and the boundary + # set both match. + return None + if Vc.extruded or Vf.extruded: + # The DMPlex of an extruded mesh stores only the 2D base mesh. Each + # point there represents a whole vertical column of nodes, and a + # Section cannot address one node within that column. Give up here + # and let the transfer kernel evaluate every node instead. + return None + hierarchy, levelc = get_level(Vc.mesh()) + _, levelf = get_level(Vf.mesh()) + if hierarchy is None or levelc + Fraction(1, hierarchy.refinements_per_level) != levelf: + return None + cache = Vf.mesh().topology._shared_data_cache["hierarchy_preserved_node_sf"] + key = _cache_key(Vc, Vf) + try: + return cache[key] + except KeyError: + coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc] + point_sf = _preserved_point_sf(Vc.mesh().topology, Vf.mesh().topology, + coarse_to_fine) + if point_sf is None: + return cache.setdefault(key, None) + root_section = Vc.dm.getSection() + leaf_section = Vf.dm.getSection() + # `distributeSection` builds its own section over the range of points + # that the SF touches. Only the broadcast root offsets are needed + # here. Pad them back out to the full chart that `createSectionSF` + # expects. + remote_offsets, distributed_section = point_sf.distributeSection(root_section) + pStart, pEnd = leaf_section.getChart() + lpStart, lpEnd = distributed_section.getChart() + offsets = numpy.zeros(pEnd - pStart, dtype=IntType) + offsets[lpStart - pStart:lpEnd - pStart] = remote_offsets + section_sf = point_sf.createSectionSF(root_section, offsets, leaf_section) + # The transfer kernels compute only the owned fine nodes and leave + # the halo to a later exchange. Keep only the owned leaves here too: + # a ghost fine node reduced onto its coarse node would count twice. + nroots, ilocal, iremote = section_sf.getGraph() + owned = ilocal < Vf.node_set.size + trimmed = PETSc.SF().create(comm=section_sf.comm) + trimmed.setGraph(nroots, ilocal[owned], iremote[owned]) + return cache.setdefault(key, trimmed) + + +def transfer_node_subset(Vc, Vf): + """Find the fine nodes that the transfer kernels must evaluate. + + These are the nodes of ``Vf`` that :func:`preserved_node_sf` does not + already account for. Prolongation and restriction can copy the rest. + + Parameters + ---------- + Vc : firedrake.functionspaceimpl.WithGeometry + The coarse function space. + Vf : firedrake.functionspaceimpl.WithGeometry + The fine function space, on the next level of the same hierarchy. + + Returns + ------- + pyop2.types.set.Set or pyop2.types.set.Subset + A subset of the nodes of ``Vf``, or ``Vf.node_set`` itself if + :func:`preserved_node_sf` found no preserved nodes. + + """ + section_sf = preserved_node_sf(Vc, Vf) + if section_sf is None: + return Vf.node_set + cache = Vf.mesh().topology._shared_data_cache["hierarchy_transfer_node_subset"] + key = _cache_key(Vc, Vf) + try: + return cache[key] + except KeyError: + _, preserved, _ = section_sf.getGraph() + nodes = numpy.setdiff1d(numpy.arange(Vf.node_set.size, dtype=IntType), + preserved) + return cache.setdefault(key, op2.Subset(Vf.node_set, nodes)) + + +def prolong_preserved_nodes(coarse, fine): + """Copy coarse values onto the fine nodes that adaptive refinement preserved. + + Parameters + ---------- + coarse : firedrake.function.Function + The function on the coarse mesh. + fine : firedrake.function.Function + The function on the refined mesh. The transfer kernel has already + computed its other nodes. + + """ + from firedrake.halo import _get_mtype + + section_sf = preserved_node_sf(coarse.function_space(), fine.function_space()) + if section_sf is None: + return + mtype, _ = _get_mtype(fine.dat) + # The source coarse node can be a ghost node. Only owned fine nodes are + # written here, the same as the transfer kernel writes. + source = coarse.dat.data_ro_with_halos + target = fine.dat.data_wo + section_sf.bcastBegin(mtype, source, target, MPI.REPLACE) + section_sf.bcastEnd(mtype, source, target, MPI.REPLACE) + + +def restrict_preserved_nodes(fine_dual, coarse_dual): + """Add the contribution of preserved nodes to the coarse dual. + + Prolongation copies a preserved node's value without change. Restriction + is its transpose, so it adds the fine value to the coarse node unchanged. + + Parameters + ---------- + fine_dual : firedrake.cofunction.Cofunction + The cofunction on the refined mesh. + coarse_dual : firedrake.cofunction.Cofunction + The cofunction on the coarse mesh. It already holds the contribution + that the transfer kernel accumulated from the other fine nodes. + + """ + from firedrake.halo import _get_mtype + + coarse_V = coarse_dual.function_space() + section_sf = preserved_node_sf(coarse_V, fine_dual.function_space()) + if section_sf is None: + return + buffer = firedrake.Function(coarse_V) + mtype, _ = _get_mtype(buffer.dat) + source = fine_dual.dat.data_ro + target = buffer.dat.data_wo_with_halos + section_sf.reduceBegin(mtype, source, target, MPI.SUM) + section_sf.reduceEnd(mtype, source, target, MPI.SUM) + # A preserved coarse node can be a ghost on the rank that owns the + # matching fine node. Reduce the contributions onto the owning rank. + buffer.dat.local_to_global_begin(op2.INC) + buffer.dat.local_to_global_end(op2.INC) + coarse_dual.dat.data[...] += buffer.dat.data_ro + + def physical_node_locations(V): element = V.ufl_element() if V.value_shape: diff --git a/tests/firedrake/multigrid/test_adaptive_multigrid.py b/tests/firedrake/multigrid/test_adaptive_multigrid.py index 13acfdee14..a0cc5ae1f8 100644 --- a/tests/firedrake/multigrid/test_adaptive_multigrid.py +++ b/tests/firedrake/multigrid/test_adaptive_multigrid.py @@ -33,8 +33,9 @@ def _linear_expr(mesh): def coarse_mesh(request): dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)} mesher = request.param - # Big enough that refining part of it leaves untouched cells behind, and - # that a coarse cell's child count varies widely across the mesh. + # Big enough that refining part of it leaves untouched cells behind, + # which is the case the transfers copy rather than evaluate, and that + # a coarse cell's child count varies widely across the mesh. if mesher == "firedrake-square": return UnitSquareMesh(4, 4, distribution_parameters=dparams) elif mesher == "firedrake-cube": @@ -331,31 +332,32 @@ def test_adapt_before_uniform_refinement(coarse_mesh, refine): assert (fine_to_coarse[coarse_to_fine, 0] == parents).all() -@pytest.mark.parallel([1, 2, 4]) -@pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_DG0(mh, operator): - """Prolongation & Injection test for DG0""" - V_coarse = FunctionSpace(mh[0], "DG", 0) - V_fine = FunctionSpace(mh[-1], "DG", 0) - u_coarse = Function(V_coarse) - u_fine = Function(V_fine) - xc, *_ = SpatialCoordinate(V_coarse.mesh()) - stepc = conditional(ge(xc, 0), 1, 0) - xf, *_ = SpatialCoordinate(V_fine.mesh()) - stepf = conditional(ge(xf, 0), 1, 0) +def _representable_expr(mesh, degree): + """An expression that a space of the given degree holds exactly on any mesh.""" + x = SpatialCoordinate(mesh) + if degree == 0: + # The only expression a DG0 space holds on every mesh of a hierarchy + # alike is one that is constant on each of the coarsest cells. + return conditional(ge(x[0], 0), 1, 0) + return sum(xi ** degree for xi in x) - if operator == "prolong": - u_coarse.interpolate(stepc) - assert errornorm(stepc, u_coarse) <= 1e-12 - prolong(u_coarse, u_fine) - assert errornorm(stepf, u_fine) <= 1e-12 - if operator == "inject": - u_fine.interpolate(stepf) - assert errornorm(stepf, u_fine) <= 1e-12 +def _copied_nodes(mh, V): + """Count the nodes of ``V`` that the transfers copy rather than evaluate.""" + from firedrake.mg.utils import transfer_node_subset - inject(u_fine, u_coarse) - assert errornorm(stepc, u_coarse) <= 1e-12 + copied = 0 + for level in range(len(mh) - 1): + V_coarse = V.reconstruct(mesh=mh[level]) + V_fine = V.reconstruct(mesh=mh[level + 1]) + subset = transfer_node_subset(V_coarse, V_fine) + # A Subset's indices only ever span the owned range like node_set.size + # does, but transfer_node_subset falls back to the fine node_set + # itself (whose .indices spans the larger owned+halo total_size) when + # nothing is preserved, so cap at node_set.size before differencing. + visited = min(len(subset.indices), V_fine.node_set.size) + copied += V_fine.node_set.size - visited + return mh[0].comm.allreduce(copied, MPI.SUM) def _coarse_cell_integrals(mh, level, u_coarse, u_fine): @@ -490,72 +492,40 @@ def test_dg_injection_ignores_padded_children(mh, family, degree): @pytest.mark.parallel([1, 2, 4]) -@pytest.mark.parametrize("operator", ["prolong", "inject"]) -def test_CG1(mh, operator): - """Prolongation & Injection test for CG1""" - V_coarse = FunctionSpace(mh[0], "CG", 1) - V_fine = FunctionSpace(mh[-1], "CG", 1) - u_coarse = Function(V_coarse) - u_fine = Function(V_fine) - xc, *_ = SpatialCoordinate(V_coarse.mesh()) - xf, *_ = SpatialCoordinate(V_fine.mesh()) - - if operator == "prolong": - u_coarse.interpolate(xc) - assert errornorm(xc, u_coarse) <= 1e-12 - - prolong(u_coarse, u_fine) - assert errornorm(xf, u_fine) <= 1e-12 - if operator == "inject": - u_fine.interpolate(xf) - assert errornorm(xf, u_fine) <= 1e-12 - - inject(u_fine, u_coarse) - assert errornorm(xc, u_coarse) <= 1e-12 - - -@pytest.mark.parallel([1, 2, 4]) -def test_restrict_CG1(mh): - """Test restriction with CG1""" - V_coarse = FunctionSpace(mh[0], "CG", 1) - V_fine = FunctionSpace(mh[-1], "CG", 1) - u_coarse = Function(V_coarse) - u_fine = Function(V_fine) - xc, *_ = SpatialCoordinate(V_coarse.mesh()) +@pytest.mark.parametrize("family, degree", [("DG", 0), ("CG", 1), ("CG", 2), ("CG", 3)]) +def test_transfers(mh, family, degree): + """Prolongation, injection and restriction on an adaptive hierarchy. - u_coarse.interpolate(xc) - prolong(u_coarse, u_fine) - - rf = assemble(conj(TestFunction(V_fine)) * dx) - rc = Cofunction(V_coarse.dual()) - restrict(rf, rc) - - assert np.allclose( - assemble(action(rc, u_coarse)), - assemble(action(rf, u_fine)), - rtol=1e-12 - ) + Degree 3 puts more than one node on an edge, so it is the case that + notices if the nodes of a copied entity come out in the wrong order. + """ + V_coarse = FunctionSpace(mh[0], family, degree) + V_fine = FunctionSpace(mh[-1], family, degree) + expr_coarse = _representable_expr(mh[0], degree) + expr_fine = _representable_expr(mh[-1], degree) + # The cells a refinement leaves alone are transferred by copying their + # nodes, so a hierarchy that refines everything says nothing about them. + assert _copied_nodes(mh, V_coarse) > 0 -@pytest.mark.parallel([1, 2, 4]) -def test_restrict_DG0(mh): - """Test restriction with DG0""" - V_coarse = FunctionSpace(mh[0], "DG", 0) - V_fine = FunctionSpace(mh[-1], "DG", 0) - u_coarse = Function(V_coarse) + u_coarse = Function(V_coarse).interpolate(expr_coarse) u_fine = Function(V_fine) - xc, *_ = SpatialCoordinate(V_coarse.mesh()) - - u_coarse.interpolate(xc) prolong(u_coarse, u_fine) + assert errornorm(expr_fine, u_fine) <= 1e-12 - rf = assemble(conj(TestFunction(V_fine)) * dx) - rc = Cofunction(V_coarse.dual()) - restrict(rf, rc) + u_injected = Function(V_coarse) + inject(Function(V_fine).interpolate(expr_fine), u_injected) + assert errornorm(expr_coarse, u_injected) <= 1e-12 + # Restriction is the transpose of prolongation, which pins the two down + # together: the nodes prolongation copies are the nodes restriction must + # not also accumulate through the kernel. + r_fine = assemble(conj(TestFunction(V_fine)) * dx) + r_coarse = Cofunction(V_coarse.dual()) + restrict(r_fine, r_coarse) assert np.allclose( - assemble(action(rc, u_coarse)), - assemble(action(rf, u_fine)), + assemble(action(r_coarse, u_coarse)), + assemble(action(r_fine, u_fine)), rtol=1e-12 ) From 738661afdea973c946a1468f46c2ca23d0ddbdc4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 5 Aug 2026 16:54:16 +0100 Subject: [PATCH 2/4] Address remaining PR review threads - Drop the temporary FIAT install step now that fiat#267 is merged. - preserved_points computes nfine/ncoarse via num_owned_cells instead of taking nfine as an argument, asserting the passed-in array shape against it rather than trusting it as the source of truth. --- .github/actions/install/action.yml | 8 -------- firedrake/cython/mgimpl.pyx | 8 ++++---- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index 5a08198cd4..54c1411d7f 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -163,14 +163,6 @@ runs: firedrake-clean pip list - - name: "DROP BEFORE MERGE: install FIAT from firedrakeproject/fiat#267" - shell: bash - run: | - . venv/bin/activate - pip install --verbose --no-build-isolation --no-deps --force-reinstall \ - git+https://github.com/firedrakeproject/fiat.git@pbrubeck/restricted-entity-permutations - pip list | grep -i fiat - - name: Run firedrake-check shell: bash run: | diff --git a/firedrake/cython/mgimpl.pyx b/firedrake/cython/mgimpl.pyx index f39c7a8637..5847c12f31 100644 --- a/firedrake/cython/mgimpl.pyx +++ b/firedrake/cython/mgimpl.pyx @@ -359,7 +359,6 @@ def preserved_points(PETSc.DM coarse_dm, PETSc.Section coarse_cell_numbering, PETSc.DM fine_dm, PETSc.Section fine_cell_numbering, - PetscInt nfine, np.ndarray coarse_to_fine_cells): """Pair the points an adaptive refinement left alone with their coarse originals. @@ -373,14 +372,13 @@ def preserved_points(PETSc.DM coarse_dm, :arg coarse_cell_numbering: the coarse mesh's cell numbering section. :arg fine_dm: the adaptively refined DMPlex. :arg fine_cell_numbering: the fine mesh's cell numbering section. - :arg nfine: the number of owned fine cells. :arg coarse_to_fine_cells: the Firedrake-numbered coarse-to-fine cell map. :returns: an array over the chart of ``fine_dm``, holding for each fine point the coarse point it was copied from, or -1 if the refinement changed it. """ cdef: - PetscInt ncoarse, max_children, c, i, off, child + PetscInt ncoarse, nfine, max_children, c, i, off, child PetscInt cStart, cEnd, pStart, pEnd, coarse_size, fine_size PetscInt *coarse_closure = NULL PetscInt *fine_closure = NULL @@ -388,8 +386,10 @@ def preserved_points(PETSc.DM coarse_dm, PetscInt[:, ::1] coarse_to_fine coarse_to_fine = coarse_to_fine_cells - ncoarse = coarse_to_fine.shape[0] + ncoarse = num_owned_cells(coarse_dm) + assert ncoarse == coarse_to_fine.shape[0] max_children = coarse_to_fine.shape[1] + nfine = num_owned_cells(fine_dm) # Both cell maps are in Firedrake numbering, so invert each mesh's cell # numbering section to get back to the plex points the closures live on. From 994901be091fb639344fc205cdf388d75c5320d0 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 8 Aug 2026 09:56:22 +0100 Subject: [PATCH 3/4] Rewrite the clause-stacked docstrings and comments Follow ASD-STE100: short sentences, one idea each, active voice, and the subject named up front rather than buried in a relative clause. The AGENTS.md rules this follows are in #5338. Co-Authored-By: Claude Opus 5 --- firedrake/cython/mgimpl.pyx | 54 +++++++----- firedrake/mg/interface.py | 8 +- firedrake/mg/utils.py | 3 +- .../multigrid/test_adaptive_multigrid.py | 85 +++++++++++-------- 4 files changed, 86 insertions(+), 64 deletions(-) diff --git a/firedrake/cython/mgimpl.pyx b/firedrake/cython/mgimpl.pyx index 5847c12f31..69b7aec26d 100644 --- a/firedrake/cython/mgimpl.pyx +++ b/firedrake/cython/mgimpl.pyx @@ -360,22 +360,34 @@ def preserved_points(PETSc.DM coarse_dm, PETSc.DM fine_dm, PETSc.Section fine_cell_numbering, np.ndarray coarse_to_fine_cells): - """Pair the points an adaptive refinement left alone with their coarse originals. - - A coarse cell that the refinement did not touch is copied into the fine - mesh, so its whole closure is copied with it: the transform preserves the - cone of every point it does not refine, and hence the plex closure of the - copied cell entry by entry. Such a cell has exactly one child, which is - what the right-padding of ``coarse_to_fine_cells`` with -1 identifies. - - :arg coarse_dm: the coarse mesh DMPlex. - :arg coarse_cell_numbering: the coarse mesh's cell numbering section. - :arg fine_dm: the adaptively refined DMPlex. - :arg fine_cell_numbering: the fine mesh's cell numbering section. - :arg coarse_to_fine_cells: the Firedrake-numbered coarse-to-fine cell map. - :returns: an array over the chart of ``fine_dm``, holding for each fine - point the coarse point it was copied from, or -1 if the refinement - changed it. + """Pair unrefined fine points with their coarse originals. + + Adaptive refinement copies an untouched coarse cell into the fine mesh + without change. It therefore preserves the cone of every point in that + cell, and so preserves the whole plex closure, point for point. Such a + cell has exactly one child. The right-padding of ``coarse_to_fine_cells`` + with -1 identifies which cells these are. + + Parameters + ---------- + coarse_dm : PETSc.DM + The coarse mesh DMPlex. + coarse_cell_numbering : PETSc.Section + The cell numbering section of the coarse mesh. + fine_dm : PETSc.DM + The adaptively refined DMPlex. + fine_cell_numbering : PETSc.Section + The cell numbering section of the fine mesh. + coarse_to_fine_cells : numpy.ndarray + The Firedrake-numbered coarse-to-fine cell map. + + Returns + ------- + numpy.ndarray + An array over the chart of ``fine_dm``. For each fine point, it + holds the coarse point it was copied from, or -1 if refinement + changed that point. + """ cdef: PetscInt ncoarse, nfine, max_children, c, i, off, child @@ -418,13 +430,13 @@ def preserved_points(PETSc.DM coarse_dm, &coarse_size, &coarse_closure)) CHKERR(DMPlexGetTransitiveClosure(fine_dm.dm, fine_point[child], PETSC_TRUE, &fine_size, &fine_closure)) - # A cell with one child that the transform nonetheless changed would - # have a closure of its own shape; leave it to the transfer kernel. + # A one-child cell that refinement did change would have a closure + # of a different size. Skip it and let the transfer kernel handle it. if coarse_size == fine_size: for i in range(coarse_size): - # The closures interleave points with their orientations, and - # only points that carry the same orientation in both meshes - # order their nodes the same way. + # Each closure interleaves a point with its orientation. Copy + # a point only when its orientation matches in both meshes: + # only then do the two cells order their nodes the same way. if coarse_closure[2*i + 1] == fine_closure[2*i + 1]: fine_to_coarse[fine_closure[2*i] - pStart] = coarse_closure[2*i] CHKERR(DMPlexRestoreTransitiveClosure(coarse_dm.dm, coarse_point[c], PETSC_TRUE, diff --git a/firedrake/mg/interface.py b/firedrake/mg/interface.py index 629ac97c07..6c42d3c434 100644 --- a/firedrake/mg/interface.py +++ b/firedrake/mg/interface.py @@ -102,8 +102,8 @@ def prolong(coarse, fine): for d in [coarse, coarse_coords]: d.dat.global_to_local_begin(op2.READ) d.dat.global_to_local_end(op2.READ) - # An adaptive refinement leaves most of the mesh alone, and the nodes - # it preserves are copied rather than evaluated. + # Adaptive refinement leaves most of the mesh unchanged. Copy the + # value at those nodes instead of evaluating it there. node_subset = utils.transfer_node_subset(Vc, Vf) op2.par_loop(kernel, node_subset, *kernel_args) utils.prolong_preserved_nodes(coarse, fine) @@ -188,8 +188,8 @@ def restrict(fine_dual, coarse_dual): for d in [coarse_coords]: d.dat.global_to_local_begin(op2.READ) d.dat.global_to_local_end(op2.READ) - # Restriction transposes prolongation, so it skips the same fine nodes - # and hands the coarse nodes they pair with their values whole. + # Restriction is the transpose of prolongation. It skips the same + # fine nodes and adds their value to the matching coarse node. node_subset = utils.transfer_node_subset(Vc, Vf) op2.par_loop(kernel, node_subset, *kernel_args) utils.restrict_preserved_nodes(fine_dual, coarse_dual) diff --git a/firedrake/mg/utils.py b/firedrake/mg/utils.py index 7b08970384..eda06238b1 100644 --- a/firedrake/mg/utils.py +++ b/firedrake/mg/utils.py @@ -8,6 +8,7 @@ import finat.ufl import firedrake from firedrake.cython import mgimpl as impl +from firedrake.halo import _get_mtype def fine_node_to_coarse_node_map(Vf, Vc): @@ -470,7 +471,6 @@ def prolong_preserved_nodes(coarse, fine): computed its other nodes. """ - from firedrake.halo import _get_mtype section_sf = preserved_node_sf(coarse.function_space(), fine.function_space()) if section_sf is None: @@ -499,7 +499,6 @@ def restrict_preserved_nodes(fine_dual, coarse_dual): that the transfer kernel accumulated from the other fine nodes. """ - from firedrake.halo import _get_mtype coarse_V = coarse_dual.function_space() section_sf = preserved_node_sf(coarse_V, fine_dual.function_space()) diff --git a/tests/firedrake/multigrid/test_adaptive_multigrid.py b/tests/firedrake/multigrid/test_adaptive_multigrid.py index a0cc5ae1f8..f3d819d798 100644 --- a/tests/firedrake/multigrid/test_adaptive_multigrid.py +++ b/tests/firedrake/multigrid/test_adaptive_multigrid.py @@ -2,6 +2,8 @@ import numpy as np from mpi4py import MPI from firedrake import * +from firedrake.mg.utils import coarse_cell_to_fine_node_map, transfer_node_subset +from firedrake.utils import complex_mode def corner_adaptive_hierarchy(base, nlevels): @@ -33,9 +35,9 @@ def _linear_expr(mesh): def coarse_mesh(request): dparams = {"overlap_type": (DistributedMeshOverlapType.VERTEX, 1)} mesher = request.param - # Big enough that refining part of it leaves untouched cells behind, - # which is the case the transfers copy rather than evaluate, and that - # a coarse cell's child count varies widely across the mesh. + # Big enough that refining part of it leaves untouched cells behind. + # The transfers copy those cells' nodes instead of evaluating them. + # It also gives a coarse cell's child count a wide range. if mesher == "firedrake-square": return UnitSquareMesh(4, 4, distribution_parameters=dparams) elif mesher == "firedrake-cube": @@ -116,8 +118,10 @@ def test_refine_marked_elements_is_local(): @pytest.mark.parallel([1, 2]) def test_refine_marked_elements_repeats(coarse_mesh): - """A marker value of n refines the marked cells n times, and the cell maps - reach all the way from the original mesh to the n-times-refined one.""" + """A marker value of n refines the marked cells n times. + + The cell maps reach all the way from the original mesh to the + n-times-refined one.""" mesh = coarse_mesh ncells = {} max_children = {} @@ -148,8 +152,9 @@ def test_refine_marked_elements_repeats(coarse_mesh): def test_add_mesh_rejects_unrelated_mesh(): """Cell maps are only meaningful relative to the mesh they were built - against, so a mesh refined from anything but the finest level is refused - rather than silently recorded with somebody else's maps.""" + against. Refuse a mesh refined from anything but the finest level, + instead of silently recording it with maps that belong to another + mesh.""" mh = MeshHierarchy(UnitSquareMesh(2, 2)) other = UnitSquareMesh(4, 4) @@ -215,10 +220,11 @@ def test_CG1_native_transfers_use_adaptive_cell_maps(coarse_mesh): def _assert_adapt_after_uniform_refinement(mh): - """Adaptively refine the finest level of the uniformly-refined hierarchy - ``mh`` by marking a single cell, and check that the cell maps of the level - this adds are sane. Shared by the ``test_adapt_after_uniform_*refinement`` - tests, which only differ in how ``mh`` itself was built. + """Adaptively refine the finest level of the hierarchy ``mh``. + + Mark a single cell and check that the cell maps of the new level are + sane. The ``test_adapt_after_uniform_*refinement`` tests share this + helper; they differ only in how they build ``mh``. """ mesh = mh[-1] level = len(mh) @@ -264,9 +270,10 @@ def test_adapt_after_uniform_netgen_refinement(): @pytest.mark.parallel([1, 2]) @pytest.mark.parametrize("degree", [1, 2]) def test_adapt_preserves_mesh_metadata(degree): - """Adaptive refinement carries the Netgen geometry and flags, and the mesh - construction parameters, over to the refined mesh, so that the refined - mesh can itself be refined again.""" + """Adaptive refinement carries mesh metadata to the refined mesh. + + It copies the Netgen geometry, flags, and construction parameters, so + the refined mesh can itself be refined again.""" from netgen.geom2d import CSG2d, Circle geo = CSG2d() geo.Add(Circle(center=(0, 0), radius=1.0, bc="circle")) @@ -304,8 +311,9 @@ def test_adapt_after_uniform_refinement(coarse_mesh, refine): @pytest.mark.parametrize("refine", [1, 2]) def test_adapt_before_uniform_refinement(coarse_mesh, refine): """An adaptively refined mesh can be uniformly refined into a hierarchy. - Its plex numbers cells by refinement case, so its owned cells are - interleaved with its halo cells, which the cell maps must not assume away. + + Its plex numbers cells by refinement case. This interleaves owned cells + with halo cells, and the cell maps must not assume otherwise. """ netgen_flags = {} if hasattr(coarse_mesh, "netgen_mesh") else None @@ -336,25 +344,23 @@ def _representable_expr(mesh, degree): """An expression that a space of the given degree holds exactly on any mesh.""" x = SpatialCoordinate(mesh) if degree == 0: - # The only expression a DG0 space holds on every mesh of a hierarchy - # alike is one that is constant on each of the coarsest cells. + # A DG0 space holds an expression exactly on every mesh of the + # hierarchy only if it is constant on each coarsest cell. return conditional(ge(x[0], 0), 1, 0) return sum(xi ** degree for xi in x) def _copied_nodes(mh, V): """Count the nodes of ``V`` that the transfers copy rather than evaluate.""" - from firedrake.mg.utils import transfer_node_subset - copied = 0 for level in range(len(mh) - 1): V_coarse = V.reconstruct(mesh=mh[level]) V_fine = V.reconstruct(mesh=mh[level + 1]) subset = transfer_node_subset(V_coarse, V_fine) - # A Subset's indices only ever span the owned range like node_set.size - # does, but transfer_node_subset falls back to the fine node_set - # itself (whose .indices spans the larger owned+halo total_size) when - # nothing is preserved, so cap at node_set.size before differencing. + # A Subset's .indices spans the owned range, like node_set.size does. + # But when nothing is preserved, transfer_node_subset falls back to + # the fine node_set itself, whose .indices spans the larger + # owned+halo total_size. Cap at node_set.size before differencing. visited = min(len(subset.indices), V_fine.node_set.size) copied += V_fine.node_set.size - visited return mh[0].comm.allreduce(copied, MPI.SUM) @@ -438,8 +444,6 @@ def _poison_padding(mh, level, Vc, Vf): Returns the number of slots it overwrote. """ - from firedrake.mg.utils import coarse_cell_to_fine_node_map - children = mh.coarse_to_fine_cells[level][:mh[level].cell_set.size] valid = children >= 0 # Rows carry different numbers of children, so the padded slots do not @@ -496,16 +500,17 @@ def test_dg_injection_ignores_padded_children(mh, family, degree): def test_transfers(mh, family, degree): """Prolongation, injection and restriction on an adaptive hierarchy. - Degree 3 puts more than one node on an edge, so it is the case that - notices if the nodes of a copied entity come out in the wrong order. + Degree 3 puts more than one node on an edge. This catches a copied + entity whose nodes come out in the wrong order. """ V_coarse = FunctionSpace(mh[0], family, degree) V_fine = FunctionSpace(mh[-1], family, degree) expr_coarse = _representable_expr(mh[0], degree) expr_fine = _representable_expr(mh[-1], degree) - # The cells a refinement leaves alone are transferred by copying their - # nodes, so a hierarchy that refines everything says nothing about them. + # The cells that refinement leaves alone are transferred by copying + # their nodes. A hierarchy that refines everything says nothing about + # them. assert _copied_nodes(mh, V_coarse) > 0 u_coarse = Function(V_coarse).interpolate(expr_coarse) @@ -513,13 +518,9 @@ def test_transfers(mh, family, degree): prolong(u_coarse, u_fine) assert errornorm(expr_fine, u_fine) <= 1e-12 - u_injected = Function(V_coarse) - inject(Function(V_fine).interpolate(expr_fine), u_injected) - assert errornorm(expr_coarse, u_injected) <= 1e-12 - - # Restriction is the transpose of prolongation, which pins the two down - # together: the nodes prolongation copies are the nodes restriction must - # not also accumulate through the kernel. + # Restriction is the transpose of prolongation. This ties the two + # together: restriction must not also accumulate, through the kernel, + # the same nodes that prolongation copies. r_fine = assemble(conj(TestFunction(V_fine)) * dx) r_coarse = Cofunction(V_coarse.dual()) restrict(r_fine, r_coarse) @@ -529,6 +530,16 @@ def test_transfers(mh, family, degree): rtol=1e-12 ) + # Injection + u_fine = Function(V_fine).interpolate(expr_fine) + u_injected = Function(V_coarse) + if family in {"DG", "DQ"} and complex_mode: + with pytest.raises(NotImplementedError): + inject(u_fine, u_injected) + else: + inject(u_fine, u_injected) + assert errornorm(expr_coarse, u_injected) <= 1e-12 + @pytest.mark.parallel([1, 2]) def test_mg_jacobi(mh): From 16605b71643abb54d2a1f967a93279aca099b594 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sat, 29 Aug 2026 21:40:20 +0100 Subject: [PATCH 4/4] Say once what a null preserved_node_sf means Name the four reasons no node can be preserved in one Notes section, and let each guard stand on its own. The collective agreement moves to where the verdict is reached, so _preserved_point_sf always returns an SF. Document why restriction accumulates instead of copying. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01GRZG3iuNzgXQWGiuzrjKQo --- firedrake/mg/utils.py | 70 ++++++++++++------- .../multigrid/test_adaptive_multigrid.py | 13 ++-- 2 files changed, 48 insertions(+), 35 deletions(-) diff --git a/firedrake/mg/utils.py b/firedrake/mg/utils.py index b2204dddd2..042bf5edaf 100644 --- a/firedrake/mg/utils.py +++ b/firedrake/mg/utils.py @@ -232,7 +232,7 @@ def _preserved_point_sf(coarse_mesh, fine_mesh, coarse_to_fine): ------- PETSc.SF An SF with roots on the points of ``coarse_mesh`` and leaves on the - unrefined points of ``fine_mesh``. Returns `None` if refinement + unrefined points of ``fine_mesh``. It has no leaves where refinement changed every cell, as a uniform refinement does. """ @@ -243,12 +243,7 @@ def _preserved_point_sf(coarse_mesh, fine_mesh, coarse_to_fine): fine_plex, fine_mesh._cell_numbering, coarse_to_fine, ) - leaves, = numpy.nonzero(fine_to_coarse_points >= 0) - # A uniform refinement preserves no points. Every rank must agree on - # whether to build the SF at all, not just the ranks with no leaves. - if not fine_plex.comm.tompi4py().allreduce(len(leaves) > 0, op=MPI.LOR): - return None - leaves = leaves.astype(IntType) + leaves = numpy.nonzero(fine_to_coarse_points >= 0)[0].astype(IntType) # Refinement acts on each rank's own plex. A fine point and the coarse # point it was copied from always live on the same rank. remote = numpy.empty((len(leaves), 2), dtype=IntType) @@ -260,7 +255,10 @@ def _preserved_point_sf(coarse_mesh, fine_mesh, coarse_to_fine): return point_sf -def preserved_node_sf(Vc, Vf): +def preserved_node_sf( + Vc: firedrake.functionspaceimpl.WithGeometry, + Vf: firedrake.functionspaceimpl.WithGeometry, +) -> PETSc.SF | None: """Find the nodes that adaptive refinement leaves unchanged. An unrefined cell has the same nodes in both spaces. The transfer @@ -276,21 +274,28 @@ def preserved_node_sf(Vc, Vf): Returns ------- - PETSc.SF - An SF with roots on the nodes of ``Vc`` and leaves on the matching - nodes of ``Vf``. Returns `None` if no nodes match. + PETSc.SF or None + An SF with roots on the nodes of ``Vc`` and leaves on the owned nodes + of ``Vf`` that match. `None` says that no node matches anywhere, so + that a caller can skip the copy and evaluate every node instead. + + Notes + ----- + Four things stop any node from matching, and every rank reaches the same + verdict on each of them: + + * the two spaces lay their nodes out differently, because their elements + or their boundary sets differ; + * one of the meshes is extruded. Its DMPlex holds the base mesh alone, so + a point there stands for a whole vertical column of nodes and a + `PETSc.Section` cannot address one node within that column; + * the two meshes are not consecutive levels of one hierarchy; + * the refinement is uniform, and so rebuilt every cell. """ if Vc.ufl_element() != Vf.ufl_element() or Vc.boundary_set != Vf.boundary_set: - # A space and its counterpart on the refined mesh use the same node - # layout on an unrefined cell only when the element and the boundary - # set both match. return None if Vc.extruded or Vf.extruded: - # The DMPlex of an extruded mesh stores only the 2D base mesh. Each - # point there represents a whole vertical column of nodes, and a - # Section cannot address one node within that column. Give up here - # and let the transfer kernel evaluate every node instead. return None hierarchy, levelc = get_level(Vc.mesh()) _, levelf = get_level(Vf.mesh()) @@ -304,8 +309,6 @@ def preserved_node_sf(Vc, Vf): coarse_to_fine = hierarchy.coarse_to_fine_cells[levelc] point_sf = _preserved_point_sf(Vc.mesh().topology, Vf.mesh().topology, coarse_to_fine) - if point_sf is None: - return cache.setdefault(key, None) root_section = Vc.dm.getSection() leaf_section = Vf.dm.getSection() # `distributeSection` builds its own section over the range of points @@ -323,12 +326,20 @@ def preserved_node_sf(Vc, Vf): # a ghost fine node reduced onto its coarse node would count twice. nroots, ilocal, iremote = section_sf.getGraph() owned = ilocal < Vf.node_set.size + # Every rank must agree on whether to copy or to evaluate, or they + # generate different code. A rank with no owned leaf of its own still + # takes part where another rank has one. + if not Vf.mesh().comm.allreduce(bool(owned.any()), op=MPI.LOR): + return cache.setdefault(key, None) trimmed = PETSc.SF().create(comm=section_sf.comm) trimmed.setGraph(nroots, ilocal[owned], iremote[owned]) return cache.setdefault(key, trimmed) -def transfer_node_subset(Vc, Vf): +def transfer_node_subset( + Vc: firedrake.functionspaceimpl.WithGeometry, + Vf: firedrake.functionspaceimpl.WithGeometry, +) -> op2.Set: """Find the fine nodes that the transfer kernels must evaluate. These are the nodes of ``Vf`` that :func:`preserved_node_sf` does not @@ -344,8 +355,9 @@ def transfer_node_subset(Vc, Vf): Returns ------- pyop2.types.set.Set or pyop2.types.set.Subset - A subset of the nodes of ``Vf``, or ``Vf.node_set`` itself if - :func:`preserved_node_sf` found no preserved nodes. + A subset of the nodes of ``Vf``. Where nothing is preserved this is + ``Vf.node_set`` itself, which spares the kernel a level of + indirection that would index every node anyway. """ section_sf = preserved_node_sf(Vc, Vf) @@ -374,7 +386,6 @@ def prolong_preserved_nodes(coarse, fine): computed its other nodes. """ - section_sf = preserved_node_sf(coarse.function_space(), fine.function_space()) if section_sf is None: return @@ -401,8 +412,15 @@ def restrict_preserved_nodes(fine_dual, coarse_dual): The cofunction on the coarse mesh. It already holds the contribution that the transfer kernel accumulated from the other fine nodes. - """ + Notes + ----- + This adds rather than copies, at every stage. A coarse basis function + does not vanish on the cells around a preserved node that refinement did + split, so the same coarse node also collects a contribution from the + kernel. Several preserved fine nodes can likewise reduce onto one coarse + node, and the coarse node can be a ghost on the rank that owns them. + """ coarse_V = coarse_dual.function_space() section_sf = preserved_node_sf(coarse_V, fine_dual.function_space()) if section_sf is None: @@ -413,8 +431,6 @@ def restrict_preserved_nodes(fine_dual, coarse_dual): target = buffer.dat.data_wo_with_halos section_sf.reduceBegin(mtype, source, target, MPI.SUM) section_sf.reduceEnd(mtype, source, target, MPI.SUM) - # A preserved coarse node can be a ghost on the rank that owns the - # matching fine node. Reduce the contributions onto the owning rank. buffer.dat.local_to_global_begin(op2.INC) buffer.dat.local_to_global_end(op2.INC) coarse_dual.dat.data[...] += buffer.dat.data_ro diff --git a/tests/firedrake/multigrid/test_adaptive_multigrid.py b/tests/firedrake/multigrid/test_adaptive_multigrid.py index 69baba63c7..bfdc63abe1 100644 --- a/tests/firedrake/multigrid/test_adaptive_multigrid.py +++ b/tests/firedrake/multigrid/test_adaptive_multigrid.py @@ -2,7 +2,7 @@ import numpy as np from mpi4py import MPI from firedrake import * -from firedrake.mg.utils import transfer_node_subset +from firedrake.mg.utils import preserved_node_sf from firedrake.utils import complex_mode @@ -356,13 +356,10 @@ def _copied_nodes(mh, V): for level in range(len(mh) - 1): V_coarse = V.reconstruct(mesh=mh[level]) V_fine = V.reconstruct(mesh=mh[level + 1]) - subset = transfer_node_subset(V_coarse, V_fine) - # A Subset's .indices spans the owned range, like node_set.size does. - # But when nothing is preserved, transfer_node_subset falls back to - # the fine node_set itself, whose .indices spans the larger - # owned+halo total_size. Cap at node_set.size before differencing. - visited = min(len(subset.indices), V_fine.node_set.size) - copied += V_fine.node_set.size - visited + section_sf = preserved_node_sf(V_coarse, V_fine) + if section_sf is not None: + _, preserved, _ = section_sf.getGraph() + copied += len(preserved) return mh[0].comm.allreduce(copied, MPI.SUM)