From dceb95771a0b88f48788c2f6f6f58fcb166c2519 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Sat, 24 Oct 2026 17:54:09 +1100 Subject: [PATCH 1/2] Fix submesh injection when parent ownership differs across ranks Reduce source values and multiplicities through the parent field DM before owner synchronisation. This preserves spatial injection without global coordinate gathering, including add mode and untouched exterior values. A parent owner may retain no extracted-region cells even when another rank has the corresponding source node. Add independent P1/P2 vector replace/add regressions. Reproducer: annulus 2h buffers at h=1/32 on 16 ranks lose four inner-wall source values despite correct saved submesh topography. Static compilation and diff checks pass; matching versioned int32 installation, serial/MPI regression gates and affected scientific reruns remain pending on Gadi. --- .../discretisation/discretisation_mesh.py | 49 ++++++++++++++++--- ...est_1087_submesh_prolongation_ownership.py | 38 ++++++++++++++ 2 files changed, 80 insertions(+), 7 deletions(-) create mode 100644 tests/parallel/test_1087_submesh_prolongation_ownership.py diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index ccfca8c36..095d8acd3 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -2540,6 +2540,16 @@ def prolongate(self, sub_var, parent_var, mode="replace"): ``"replace"`` overwrites parent values at submesh DOFs. ``"add"`` adds submesh values into parent. + Notes + ----- + A parent DOF owner need not retain any extracted-region cells. Local + coordinate matches can therefore exist only on another rank's parent + ghosts. Reduce matching contributions and their multiplicities through + the parent field DM before writing owner values. Averaging duplicate + shared copies preserves injection, including ``add`` semantics; unmatched + parent DOFs keep their original values. No global coordinate gather or + interpolation is used. + Raises ------ ValueError @@ -2554,15 +2564,40 @@ def prolongate(self, sub_var, parent_var, mode="replace"): sub_rows, parent_rows = self._build_dof_map(parent_var, sub_var) - new_data = numpy.array(parent_var.data) - - if mode == "replace": - new_data[parent_rows] = sub_var.data[sub_rows] - elif mode == "add": - new_data[parent_rows] += sub_var.data[sub_rows] - else: + if mode not in ("replace", "add"): raise ValueError(f"mode must be 'replace' or 'add', got '{mode}'") + new_data = numpy.array(parent_var.data) + indexset, field_dm = self.parent.dm.createSubDM(parent_var.field_id) + values = parent_var.vec.duplicate() + counts = parent_var.vec.duplicate() + global_values = parent_var._gvec.duplicate() + global_counts = parent_var._gvec.duplicate() + try: + values.set(0.0) + counts.set(0.0) + values.array.reshape(new_data.shape)[parent_rows] = sub_var.data[sub_rows] + counts.array.reshape(new_data.shape)[parent_rows] = 1.0 + global_values.set(0.0) + global_counts.set(0.0) + field_dm.localToGlobal(values, global_values, addv=PETSc.InsertMode.ADD_VALUES) + field_dm.localToGlobal(counts, global_counts, addv=PETSc.InsertMode.ADD_VALUES) + field_dm.globalToLocal(global_values, values, addv=False) + field_dm.globalToLocal(global_counts, counts, addv=False) + multiplicity = counts.array.reshape(new_data.shape) + matched = multiplicity > 0 + injected = values.array.reshape(new_data.shape)[matched] / multiplicity[matched] + if mode == "replace": + new_data[matched] = injected + else: + new_data[matched] += injected + finally: + values.destroy() + counts.destroy() + global_values.destroy() + global_counts.destroy() + field_dm.destroy() + indexset.destroy() parent_var.pack_raw_data_to_petsc(new_data, sync=True) parent_var._data_is_dirty = True diff --git a/tests/parallel/test_1087_submesh_prolongation_ownership.py b/tests/parallel/test_1087_submesh_prolongation_ownership.py new file mode 100644 index 000000000..ed3b23cf9 --- /dev/null +++ b/tests/parallel/test_1087_submesh_prolongation_ownership.py @@ -0,0 +1,38 @@ +"""Injection must reach parent owners that have no local extracted-region node. + +Check P1/P2 vector fields, nonzero untouched parent values and both transfer +operations. Global coordinate gathering belongs only to this small independent +test oracle; the production transfer communicates through the parent PETSc DM. +""" + +import numpy as np +from mpi4py import MPI +import pytest +from scipy.spatial import cKDTree +import underworld3 as uw + + +@pytest.mark.parametrize("degree", (1, 2)) +@pytest.mark.parametrize("mode", ("replace", "add")) +def test_submesh_injection_preserves_all_parent_copies(degree, mode): + parent = uw.meshing.BoxInternalBoundary( + elementRes=(12, 12), zelementRes=(10, 2), zintCoord=.85, + minCoords=(0., 0.), maxCoords=(1., 1.), simplex=True, qdegree=2, + ) + child = parent.extract_region("Outer") + source = uw.discretisation.MeshVariable("Child", child, 2, degree=degree) + target = uw.discretisation.MeshVariable("Parent", parent, 2, degree=degree) + x = np.asarray(source.coords) + source.data[...] = np.column_stack((1 + x[:, 0] + 2*x[:, 1], 3 - x[:, 0])) + target.data[...] = 7. + child.prolongate(source, target, mode=mode) + + global_coords = np.concatenate(uw.mpi.comm.allgather(x)) + distance, _ = cKDTree(global_coords).query(np.asarray(target.coords)) + selected = distance < 1e-10 + expected = np.full(target.data.shape, 7.) + points = np.asarray(target.coords)[selected] + values = np.column_stack((1 + points[:, 0] + 2*points[:, 1], 3 - points[:, 0])) + expected[selected] = values + (7. if mode == "add" else 0.) + local = float(np.max(np.abs(target.data - expected), initial=0.)) + assert uw.mpi.comm.allreduce(local, op=MPI.MAX) < 1e-12 From 17b27409a404d17d8938ddae0943a4315f1d9397 Mon Sep 17 00:00:00 2001 From: Tyagi Date: Sun, 25 Oct 2026 02:28:52 +1100 Subject: [PATCH 2/2] Preserve parent partitions when extracting submeshes Avoid implicit child redistribution when a filtered region or surface occupies only one MPI rank. Preserve the parent-derived point/DOF mappings and retain the existing parent-owner reduction. Make surface emptiness validation collective so ranks outside a valid surface do not fail before PETSc extraction. Add deterministic P1/P2 replace/add coverage with the region concentrated on one parent rank, plus empty-rank re-extraction/cell-size and surface-label tests. The four new transfer cases fail on the original PR with errors 5.0/4.0 before this change. Validation on an isolated rebuilt Mac PETSc 3.25 int32/Open MPI runtime: 20 serial passed, one existing expected failure; 23 passed on each of two ranks; 10 passed on each of four ranks. Python parsing, Black test formatting, deprecated-pattern gate, CI test reachability and diff checks pass. Shared runtime and benchmark results unchanged. Next: verify the fresh Linux/MPICH CI run. Underworld development team with AI support from Claude Code --- .../design/submesh-solver-architecture.md | 9 ++ .../discretisation/discretisation_mesh.py | 33 ++++--- ...est_1087_submesh_prolongation_ownership.py | 87 ++++++++++++++++--- 3 files changed, 100 insertions(+), 29 deletions(-) diff --git a/docs/developer/design/submesh-solver-architecture.md b/docs/developer/design/submesh-solver-architecture.md index 866347b84..d745aaf7c 100644 --- a/docs/developer/design/submesh-solver-architecture.md +++ b/docs/developer/design/submesh-solver-architecture.md @@ -49,6 +49,15 @@ lineage link (`parent`, registration in `parent._registered_submeshes`) and supports explicit field transfer back and forth. There are three flavours, and they share one usage pattern: +Region and surface extraction preserve the parent-derived partition, even when +only one rank contains the extracted mesh. PETSc's `isDistributed()` reports +whether multiple ranks own points; a false result does not require redistributing +an extracted child. Implicit redistribution would invalidate rank-local parent +maps and the saved subpoint mapping. Shared field values are still reduced through +the parent field DM, since a parent DOF owner need not own an extracted cell. +Independent child repartitioning requires migrating the parent--child mapping; +it is not provided by coordinate-coincidence injection. + > **get a child → build a solver on it → map fields back and forth** | | Subdomain | Resolution level | Surface | diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 095d8acd3..1e9fff598 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/src/underworld3/discretisation/discretisation_mesh.py @@ -432,6 +432,7 @@ def __init__( units=None, verbose=False, *args, + _preserve_partition=False, **kwargs, ): self.instance = Mesh.mesh_instances @@ -546,7 +547,9 @@ def __init__( self._build_coarsened_hierarchy(coarsening) else: - if not self.dm.isDistributed(): + # An extracted region can occupy one rank without needing a new + # partition. Redistributing it breaks its rank-local parent maps. + if not _preserve_partition and not self.dm.isDistributed(): self.sf1 = self.dm.distribute() # On manifold meshes (dim != cdim — e.g. SphericalManifold @@ -2042,6 +2045,7 @@ def extract_region(self, label_name, label_value=None): boundaries=sub_boundaries, coordinate_system_type=self.CoordinateSystemType, verbose=False, + _preserve_partition=True, ) # Store lineage @@ -2129,33 +2133,25 @@ def extract_surface(self, label_name, label_value=None, verbose=False): f"Available: {[b.name for b in self.boundaries]}" ) - # --- loud-fail if the parent has no faces marked with this label --- + # Reject globally empty surfaces, not ranks outside the surface. # NB: never call getStratumIS(v) for a value not in the live value set # — that hard-aborts PETSc on some labels (cf. the "Centre" pseudo- # label). Probe getValueIS() first. label = self.dm.getLabel(label_name) - if label is None: - raise ValueError( - f"Parent DM has no label '{label_name}'. " - f"Cannot extract a surface submesh from it." - ) - vals_is = label.getValueIS() + vals_is = label.getValueIS() if label is not None else None live_values = ( set(int(v) for v in vals_is.getIndices()) if vals_is is not None else set() ) - if int(label_value) not in live_values: - raise ValueError( - f"Label '{label_name}' has no stratum with value {label_value} " - f"on the parent (live values: {sorted(live_values)}). " - f"There is no surface to extract." - ) - sis = label.getStratumIS(label_value) - if sis is None or sis.getSize() == 0: + sis = label.getStratumIS(label_value) if int(label_value) in live_values else None + local_faces = sis.getSize() if sis is not None else 0 + if uw.mpi.comm.allreduce(local_faces) == 0: raise ValueError( - f"Label '{label_name}' (value {label_value}) has an empty face " - f"stratum on the parent. There is no surface to extract." + f"Label '{label_name}' (value {label_value}) has a globally empty " + f"face stratum. There is no surface to extract." ) + # The PETSc extraction is collective and needs a label on empty ranks too. + self.dm.createLabel(label_name) # Stage 1: cd-1 DM (with the phantom upward-DAG stratum). sub_with_phantoms = petsc_dm_create_submesh_from_label( @@ -2198,6 +2194,7 @@ def extract_surface(self, label_name, label_value=None, verbose=False): boundaries=sub_boundaries, coordinate_system_type=self.CoordinateSystemType, verbose=verbose, + _preserve_partition=True, ) # Submesh lineage — same shape as extract_region diff --git a/tests/parallel/test_1087_submesh_prolongation_ownership.py b/tests/parallel/test_1087_submesh_prolongation_ownership.py index ed3b23cf9..357510bb8 100644 --- a/tests/parallel/test_1087_submesh_prolongation_ownership.py +++ b/tests/parallel/test_1087_submesh_prolongation_ownership.py @@ -12,27 +12,92 @@ import underworld3 as uw -@pytest.mark.parametrize("degree", (1, 2)) -@pytest.mark.parametrize("mode", ("replace", "add")) -def test_submesh_injection_preserves_all_parent_copies(degree, mode): +pytestmark = [pytest.mark.level_2, pytest.mark.tier_b] + + +def _thin_slab_parent(partition: str) -> uw.discretisation.Mesh: parent = uw.meshing.BoxInternalBoundary( - elementRes=(12, 12), zelementRes=(10, 2), zintCoord=.85, - minCoords=(0., 0.), maxCoords=(1., 1.), simplex=True, qdegree=2, + elementRes=(12, 12), + zelementRes=(10, 2), + zintCoord=0.85, + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + simplex=True, + qdegree=2, ) + if partition == "default" or uw.mpi.size == 1: + return parent + + # Concentrating the outer slab on the last rank makes the empty-child + # case independent of Gmsh triangulation and the default partitioner. + dm = parent.dm.clone() + dm.setCoordinatesLocal(parent.dm.getCoordinatesLocal()) + start, end = dm.getHeightStratum(0) + lower, upper = [], [] + for cell in range(start, end): + _, centre, _ = dm.computeCellGeometryFVM(cell) + (lower if centre[1] < 0.5 else upper).append(cell - start) + sizes = [0] * uw.mpi.size + sizes[0], sizes[-1] = len(lower), len(upper) + partitioner = dm.getPartitioner() + partitioner.setType("shell") + partitioner.setShellPartition(uw.mpi.size, sizes, lower + upper) + migration = dm.distribute() + concentrated = uw.discretisation.Mesh(dm, boundaries=parent.boundaries, qdegree=2) + concentrated.regions = parent.regions + if migration is not None: + migration.destroy() + return concentrated + + +@pytest.mark.parametrize("partition", ("default", "single_rank_region")) +@pytest.mark.parametrize("degree", (1, 2)) +@pytest.mark.parametrize("mode", ("replace", "add")) +def test_submesh_injection_preserves_all_parent_copies(degree, mode, partition): + parent = _thin_slab_parent(partition) child = parent.extract_region("Outer") source = uw.discretisation.MeshVariable("Child", child, 2, degree=degree) target = uw.discretisation.MeshVariable("Parent", parent, 2, degree=degree) x = np.asarray(source.coords) - source.data[...] = np.column_stack((1 + x[:, 0] + 2*x[:, 1], 3 - x[:, 0])) - target.data[...] = 7. + source.array[:, 0, :] = np.column_stack((1 + x[:, 0] + 2 * x[:, 1], 3 - x[:, 0])) + target.array[:, 0, :] = 7.0 child.prolongate(source, target, mode=mode) global_coords = np.concatenate(uw.mpi.comm.allgather(x)) distance, _ = cKDTree(global_coords).query(np.asarray(target.coords)) selected = distance < 1e-10 - expected = np.full(target.data.shape, 7.) + expected = np.full(target.array[:, 0, :].shape, 7.0) points = np.asarray(target.coords)[selected] - values = np.column_stack((1 + points[:, 0] + 2*points[:, 1], 3 - points[:, 0])) - expected[selected] = values + (7. if mode == "add" else 0.) - local = float(np.max(np.abs(target.data - expected), initial=0.)) + values = np.column_stack((1 + points[:, 0] + 2 * points[:, 1], 3 - points[:, 0])) + expected[selected] = values + (7.0 if mode == "add" else 0.0) + local = float(np.max(np.abs(target.array[:, 0, :] - expected), initial=0.0)) assert uw.mpi.comm.allreduce(local, op=MPI.MAX) < 1e-12 + if partition == "single_rank_region" and uw.mpi.size > 1: + start, end = child.dm.getHeightStratum(0) + counts = uw.mpi.comm.allgather(end - start) + assert counts[-1] > 0 and all(count == 0 for count in counts[:-1]) + + +def test_region_reextraction_keeps_empty_ranks_and_cell_sizes(): + parent = _thin_slab_parent("single_rank_region") + child = parent.extract_region("Outer") + start, end = child.dm.getHeightStratum(0) + counts = uw.mpi.comm.allgather(end - start) + child._re_extract_from_parent() + start, end = child.dm.getHeightStratum(0) + assert uw.mpi.comm.allgather(end - start) == counts + child.cell_size() + sizes = np.asarray(child._cell_size_variable.array[:, 0, 0]) + assert sizes.size == end - start + assert np.all(sizes > 0.0) + + +def test_surface_extraction_keeps_empty_ranks(): + parent = _thin_slab_parent("single_rank_region") + surface = parent.extract_surface("Top") + start, end = surface.dm.getHeightStratum(0) + counts = uw.mpi.comm.allgather(end - start) + assert counts[-1] > 0 and all(count == 0 for count in counts[:-1]) + assert surface.dim == 1 and surface.cdim == 2 + with pytest.raises(ValueError): + parent.extract_surface("Top", label_value=99999)