diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 41146b2f4..d9a0b825a 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -502,6 +502,12 @@ def global_evaluate_nd( expr, evaluation_swarm.dm.migrate(remove_sent_points=True) uw.mpi.barrier() + # Invalidate cached data after bare-bones dm.migrate — + # particle count and values changed but Swarm.migrate() was bypassed. + evaluation_swarm._particle_coordinates._canonical_data = None + for var in evaluation_swarm._vars.values(): + if hasattr(var, "_canonical_data"): + var._canonical_data = None index = original_index.array[:,0,0] diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 72048f8f6..aa92620d7 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -3315,14 +3315,13 @@ def migrate( for index in indices: self.dm.removePointAtIndex(index) - # CRITICAL FIX: Invalidate cached data after removing particles - # The _particle_coordinates variable caches data - must refresh after DM changes - self._particle_coordinates._canonical_data = None - - # Also invalidate caches for all swarm variables - for var in self._vars.values(): - if hasattr(var, "_canonical_data"): - var._canonical_data = None + # Invalidate all cached data after migration. + # Any particle movement (send, receive, or balanced swap) makes + # cached arrays stale — both size and values may have changed. + self._particle_coordinates._canonical_data = None + for var in self._vars.values(): + if hasattr(var, "_canonical_data"): + var._canonical_data = None return diff --git a/tests/parallel/test_0760_swarm_cache_migration.py b/tests/parallel/test_0760_swarm_cache_migration.py new file mode 100644 index 000000000..d27475afa --- /dev/null +++ b/tests/parallel/test_0760_swarm_cache_migration.py @@ -0,0 +1,103 @@ +""" +Regression test for swarm cache invalidation after migration. + +Verifies that SwarmVariable._canonical_data caches are properly invalidated +after Swarm.migrate() and after bare-bones dm.migrate() in global_evaluate. + +Bug: SwarmVariable caches were only invalidated inside the delete_lost_points +branch of Swarm.migrate(), so caches became stale when particles moved between +ranks without deletion. This caused shape mismatches in global_evaluate. + +See: https://github.com/underworldcode/underworld3/issues/64 + +Run with: + mpirun -n 2 python -m pytest --with-mpi tests/parallel/test_0760_swarm_cache_migration.py +""" + +import pytest +import numpy as np +import underworld3 as uw +from mpi4py import MPI + +pytestmark = [pytest.mark.mpi(min_size=2), pytest.mark.timeout(60)] + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_swarm_cache_valid_after_migration(): + """Swarm variable caches must reflect actual particle count after migration.""" + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0)) + swarm = uw.swarm.Swarm(mesh) + var = uw.swarm.SwarmVariable("test_var", swarm, vtype=uw.VarType.SCALAR, _proxy=False) + + # Add particles at random positions — distribution will be uneven across ranks + np.random.seed(42 + uw.mpi.rank) + coords = np.random.random((200, mesh.dim)) + swarm.add_particles_with_global_coordinates(coords, migrate=False) + var.data[...] = uw.mpi.rank + + pre_count = swarm.dm.getLocalSize() + + # Migrate — particles move to owning rank + swarm.migrate(remove_sent_points=True, delete_lost_points=False) + + post_count = swarm.dm.getLocalSize() + coords_cached = swarm._particle_coordinates.data.shape[0] + var_cached = var.data.shape[0] + + # Cached sizes must match the actual DM particle count + assert coords_cached == post_count, ( + f"Rank {uw.mpi.rank}: coordinate cache ({coords_cached}) != " + f"DM count ({post_count}) after migration" + ) + assert var_cached == post_count, ( + f"Rank {uw.mpi.rank}: variable cache ({var_cached}) != " + f"DM count ({post_count}) after migration" + ) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_global_evaluate_after_migration(): + """global_evaluate must succeed with coordinates that force heavy migration.""" + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0)) + v = uw.discretisation.MeshVariable("u", mesh, mesh.dim, degree=2) + + # Bias coordinates to one side — forces cross-rank particle movement + np.random.seed(42) + N = 300 + coords = np.random.random((N, mesh.dim)) + coords[:, 0] = 0.5 + 0.5 * coords[:, 0] # all in right half + + result = uw.function.global_evaluate(v.sym, coords) + + assert result.shape[0] == N, ( + f"Rank {uw.mpi.rank}: expected {N} results, got {result.shape[0]}" + ) + + +@pytest.mark.mpi(min_size=2) +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_global_evaluate_displaced_nodes(): + """global_evaluate with displaced node coordinates (DDt/SemiLagrangian path).""" + mesh = uw.meshing.UnstructuredSimplexBox(minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0)) + v = uw.discretisation.MeshVariable("u", mesh, mesh.dim, degree=2) + + # Displace node coordinates — simulates semi-Lagrangian departure points + node_coords = mesh.X.coords + np.random.seed(7) + displacement = np.random.random(node_coords.shape) * 0.3 + mid_pt_coords = node_coords - displacement + + # Clamp to domain + mid_pt_coords = np.clip(mid_pt_coords, 0.0, 1.0) + + result = uw.function.evaluate(v.sym, mid_pt_coords) + + assert result.shape[0] == node_coords.shape[0], ( + f"Rank {uw.mpi.rank}: expected {node_coords.shape[0]} results, " + f"got {result.shape[0]}" + )