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 ccfca8c36..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: + 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 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..357510bb8 --- /dev/null +++ b/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)