Fix MPI submesh injection across parent ownership boundaries

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 |

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:
+        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}' 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:
-            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
@@ -2540,6 +2537,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 +2561,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

tests/parallel/test_1087_submesh_prolongation_ownership.py

@@ -0,0 +1,103 @@
+"""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
+
+
+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=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.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.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.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)