Swarm.migrate: shrink the working set instead of reclassifying every round

src/underworld3/swarm.py

@@ -76,6 +76,25 @@ def __setitem__(self, key, value):
         raise ValueError(self._GUIDANCE)
 
 
+def _coordinate_row_keys(coords):
+    """One opaque key per coordinate row, for exact set membership.
+
+    Used by :meth:`Swarm.migrate` to remember which points a rank has already
+    claimed, so the round loop does not reclassify them. The identity has to be
+    the coordinate rather than the index: ``dm.migrate`` does not preserve the
+    local ordering, so an index from the previous round refers to a different
+    particle after the move.
+
+    The comparison is on the exact bit pattern, so this matches only points
+    whose coordinates are identical — which is what migration produces, since
+    it moves values without touching them.
+    """
+
+    a = np.ascontiguousarray(coords, dtype=np.float64)
+
+    return a.view(np.dtype((np.void, a.dtype.itemsize * a.shape[1]))).reshape(-1)
+
+
 class SwarmType(Enum):
     """
     PETSc swarm type specification.
@@ -3728,6 +3747,11 @@ def migrate(
             swarm_coord_array,
         )
 
+        # The working set for the round loop below. Seeded with what this first
+        # pass claimed; it only ever grows, because a point in this rank's
+        # domain stays in it.
+        claimed_keys = _coordinate_row_keys(swarm_coord_array[in_or_not])
+
         num_points_in_domain = np.count_nonzero(in_or_not == True)
         num_points_not_in_domain = np.count_nonzero(in_or_not == False)
         not_my_points = np.where(in_or_not == False)[0]
@@ -3787,7 +3811,43 @@ def migrate(
                 uw.mpi.barrier()
 
                 swarm_coord_array = self.dm.getField("DMSwarmPIC_coor").reshape(-1, self.cdim)
-                in_or_not = self.mesh.points_in_domain(swarm_coord_array)
+
+                # Only classify what we have not already claimed. A point this
+                # rank has found to be in its domain stays in its domain: the
+                # coordinates do not change during migration and neither does
+                # the mesh. Re-testing it on every round is what made the
+                # per-rank classification work grow with the round count (a
+                # fixed 47k-point global set was offered to points_in_domain
+                # 238k times at np=4).
+                #
+                # The bookkeeping is by COORDINATE rather than by index because
+                # `dm.migrate` does not preserve the local ordering — measured,
+                # retained points are not left at the front — so an index from
+                # the previous round means nothing after the move. Two
+                # particles sharing a coordinate share the answer, so a
+                # collision is harmless.
+                keys_now = _coordinate_row_keys(swarm_coord_array)
+                already = np.isin(keys_now, claimed_keys)
+
+                # UNCONDITIONAL: points_in_domain is COLLECTIVE — it reaches
+                # get_max_radius() before any short-circuit, precisely so that a
+                # rank with nothing to classify still joins the reduction (the
+                # #405 treatment, stated in its own source). Calling it only
+                # when this rank has undecided points deadlocks as soon as one
+                # rank runs out of them, which is what an earlier draft of this
+                # did at np=4.
+                in_or_not = already.copy()
+                undecided = np.where(~already)[0]
+                in_or_not[undecided] = self.mesh.points_in_domain(
+                    swarm_coord_array[undecided]
+                )
+
+                newly_claimed = np.where(in_or_not & ~already)[0]
+                if newly_claimed.size:
+                    claimed_keys = np.concatenate(
+                        [claimed_keys, keys_now[newly_claimed]]
+                    )
+
                 self.dm.restoreField("DMSwarmPIC_coor")
 
                 num_points_in_domain = np.count_nonzero(in_or_not == True)

tests/parallel/test_0777_migrate_working_set_mpi.py

@@ -0,0 +1,136 @@
+"""Swarm.migrate's round loop: correctness, and the collective it must not skip.
+
+The loop classifies points once per round. It used to re-offer the whole local
+array to `points_in_domain` every round, so the classification work per rank grew
+with the round count — a fixed 47k-point set was offered 238k times at np=4. It
+now remembers what it has already claimed and only classifies the rest.
+
+The claimed set is keyed by COORDINATE, not by index: `dm.migrate` does not
+preserve the local ordering (measured — retained points are not left at the
+front), so an index from the previous round names a different particle after the
+move.
+
+The second test is here because the first draft of that change deadlocked. It
+made the `points_in_domain` call conditional on this rank having undecided
+points, and `points_in_domain` is collective — it reaches `get_max_radius()`
+before any short-circuit precisely so a rank with nothing to classify still
+joins the reduction (the #405 treatment, stated in its own source). At np=4 a
+rank ran out of undecided points, skipped the reduction, and the job hung with
+every rank inside the call.
+
+Run under MPI::
+
+    mpirun -np 4 python -m pytest --with-mpi \
+        tests/parallel/test_0777_migrate_working_set_mpi.py
+"""
+
+import numpy as np
+import pytest
+
+import underworld3 as uw
+
+pytestmark = [
+    pytest.mark.mpi(min_size=2),
+    pytest.mark.timeout(180),
+    pytest.mark.level_1,
+    pytest.mark.tier_a,
+]
+
+SEED = 20260819
+
+
+def _scattered_swarm(cell_size=0.15, n_points=2000):
+    """A swarm whose points are spread over the WHOLE domain.
+
+    Populate alone leaves every particle already owned by its own rank, so
+    nothing migrates and the round loop never runs — the case that makes this
+    file look like it is testing something when it is not.
+    """
+
+    mesh = uw.meshing.UnstructuredSimplexBox(
+        minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=cell_size,
+        regular=False, qdegree=2)
+    swarm = uw.swarm.Swarm(mesh=mesh)
+    swarm.populate(fill_param=3)
+
+    rng = np.random.default_rng(SEED)
+    everywhere = rng.uniform(0.01, 0.99, size=(n_points, 2))
+    mine = np.array_split(everywhere, uw.mpi.size)[uw.mpi.rank]
+
+    coords = swarm.dm.getField("DMSwarmPIC_coor").reshape(-1, mesh.dim)
+    n = min(coords.shape[0], mine.shape[0])
+    coords[:n] = mine[:n]
+    swarm.dm.restoreField("DMSwarmPIC_coor")
+
+    return mesh, swarm
+
+
+def _local_coords(swarm, dim):
+    out = swarm.dm.getField("DMSwarmPIC_coor").reshape(-1, dim).copy()
+    swarm.dm.restoreField("DMSwarmPIC_coor")
+
+    return out
+
+
+def test_premise_the_fixture_actually_migrates():
+    """Without points crossing ranks the round loop never runs."""
+
+    mesh, swarm = _scattered_swarm()
+    before = _local_coords(swarm, mesh.dim)
+    owned_before = int(np.count_nonzero(mesh.points_in_domain(before)))
+    strangers = uw.mpi.comm.allreduce(before.shape[0] - owned_before)
+
+    assert strangers > 0, (
+        "every point is already owned by the rank holding it, so migration has "
+        "nothing to do and neither of the tests below exercises the loop")
+
+
+def test_migrate_returns_and_every_point_lands_on_a_rank_that_holds_it():
+    """The loop terminates, and its answer is right.
+
+    The deadlock this guards against leaves every rank inside `points_in_domain`,
+    so the failure is a timeout rather than an assertion.
+    """
+
+    mesh, swarm = _scattered_swarm()
+    before = _local_coords(swarm, mesh.dim)
+    n_before = uw.mpi.comm.allreduce(before.shape[0])
+
+    swarm.migrate()
+
+    after = _local_coords(swarm, mesh.dim)
+    n_after = uw.mpi.comm.allreduce(after.shape[0])
+
+    assert n_after == n_before, (
+        f"migration changed the global particle count: {n_before} -> {n_after}")
+
+    if after.shape[0]:
+        owned = mesh.points_in_domain(after)
+        assert bool(owned.all()), (
+            f"rank {uw.mpi.rank} holds {int(np.count_nonzero(~owned))} points "
+            "its own mesh does not contain")
+    else:
+        mesh.points_in_domain(after)      # collective: join the reduction
+
+
+def test_the_partition_is_the_same_at_every_rank_count():
+    """The global set of owned coordinates is the input set, whatever np is.
+
+    Keyed on the global multiset rather than per-rank counts, which are a
+    partition detail. This is what a change to the claimed-set bookkeeping would
+    break: dropping a point, or claiming one twice.
+    """
+
+    mesh, swarm = _scattered_swarm()
+    before = _local_coords(swarm, mesh.dim)
+    swarm.migrate()
+    after = _local_coords(swarm, mesh.dim)
+
+    def gathered_rows(a):
+        rows = uw.mpi.comm.allgather(np.ascontiguousarray(a))
+        stacked = np.concatenate([r for r in rows if r.size], axis=0)
+        return stacked[np.lexsort((stacked[:, 1], stacked[:, 0]))]
+
+    assert np.array_equal(gathered_rows(before), gathered_rows(after)), (
+        "the multiset of particle coordinates changed across migration — a "
+        "point was lost or duplicated")