diff --git a/src/underworld3/swarm.py b/src/underworld3/swarm.py index 1c55d166c..4b4b679d3 100644 --- a/src/underworld3/swarm.py +++ b/src/underworld3/swarm.py @@ -343,6 +343,7 @@ def __init__( raise TypeError( f"Provided dtype={dtype} is not supported. Supported types are 'int' and 'float'." ) + self._petsc_dtype = petsc_type if _register: # Check if swarm is already populated - PETSc doesn't allow registering @@ -450,7 +451,9 @@ def _create_canonical_data_array(self, initial_data=None): # Handle case where unpack returns None (swarm not initialized) if initial_data is None: - initial_data = np.zeros((0, self.num_components)) + initial_data = np.zeros( + (0, self.num_components), dtype=self._petsc_dtype + ) # Create NDArray_With_Callback for flat data array_obj = uw.utilities.NDArray_With_Callback( @@ -1428,15 +1431,19 @@ def unpack_raw_data_from_petsc(self, squeeze=True, sync=None): # Check if swarm has any particles before accessing field swarm_size = self.swarm.local_size if swarm_size <= 0: - # Swarm not populated yet, return empty array - return np.zeros((0, self.num_components)) + # Swarm not populated yet, return empty array. Keep the field's + # PETSc dtype so that an empty rank's array agrees with a + # non-empty rank's (a float64 default here made the collective + # parallel-HDF5 create_dataset in ``save`` see different dtypes + # across ranks and deadlock on close for ``int`` variables). + return np.zeros((0, self.num_components), dtype=self._petsc_dtype) # Direct PETSc field access without context manager field_data = self.swarm.dm.getField(self.clean_name) if field_data is None: # Field not properly initialized, restore and return empty array self.swarm.dm.restoreField(self.clean_name) - return np.zeros((0, self.num_components)) + return np.zeros((0, self.num_components), dtype=self._petsc_dtype) petsc_data = field_data.reshape((-1, self.num_components)) @@ -5083,7 +5090,15 @@ def estimate_dt(self, V_fn): # silently disabling advection's step_limit substepping (BF-16). vel = np.asarray(vel) if vel.ndim == 3: - vel = vel.reshape(vel.shape[0], -1) + # Guard against empty ranks: an array of size 0 cannot be + # reshaped with a `-1` axis (NumPy cannot infer the implied + # dimension from zero elements) — e.g. (0, 1, dim) -> (0, -1) + # raises ValueError. A zero-particle rank legitimately has no + # velocities and contributes 0 to the global max below. + if vel.size == 0: + vel = np.zeros((0, vel.shape[2]) if vel.ndim >= 3 else (0,)) + else: + vel = vel.reshape(vel.shape[0], -1) try: magvel_squared = vel[:, 0] ** 2 + vel[:, 1] ** 2 diff --git a/tests/parallel/test_0795_swarm_empty_rank_evaluate_save.py b/tests/parallel/test_0795_swarm_empty_rank_evaluate_save.py new file mode 100644 index 000000000..efd021fe2 --- /dev/null +++ b/tests/parallel/test_0795_swarm_empty_rank_evaluate_save.py @@ -0,0 +1,442 @@ +"""MPI regression test for evaluate + write_timestep with empty ranks. + +A user's passive-tracer swarm hangs on HPC because some MPI ranks hold zero +particles. ``uw.function.evaluate()`` uses PETSc DMLocatePoints which is +collective — every rank must participate even if its coordinate array is +shape ``(0, dim)``. Similarly ``swarm.write_timestep()`` must complete +across all ranks. + +The bug: no existing test combines ``evaluate`` on swarm coordinates +**and** ``write_timestep`` when some ranks are explicitly empty. The user's +model assigns particles only to the crust on rank 0, leaving mantle ranks +with zero particles; the evaluate call blocks indefinitely. + +Run with:: + + mpirun -n 2 python -m pytest --with-mpi \\ + tests/parallel/test_0795_swarm_empty_rank_evaluate_save.py -v + +The ``pytest.timeout(60)`` catches hangs rather than letting the run block +indefinitely. +""" + +import os +import numpy as np +import pytest +import sympy as sp + +import underworld3 as uw + + +pytestmark = [ + pytest.mark.level_2, + pytest.mark.mpi(min_size=2), + pytest.mark.timeout(60), +] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_mesh(): + """Create a small box mesh shared by every test.""" + return uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), + maxCoords=(1.0, 1.0), + cellSize=0.25, + ) + + +def _empty_rank_coords(mesh): + """Build coordinate arrays: rank 0 gets 10 particles, others get nothing.""" + if uw.mpi.rank == 0: + coords = np.array([ + [0.10, 0.10], + [0.25, 0.15], + [0.40, 0.30], + [0.55, 0.45], + [0.70, 0.60], + [0.15, 0.80], + [0.60, 0.20], + [0.85, 0.90], + [0.35, 0.55], + [0.90, 0.10], + ]) + else: + coords = np.empty((0, mesh.dim)) + return coords + + +def _tmp_outdir(tmp_path_factory): + """Create a temp directory on rank 0, broadcast to all ranks.""" + if uw.mpi.rank == 0: + out_dir = tmp_path_factory.mktemp("swarm_eval_empty") + else: + out_dir = None + out_dir = uw.mpi.comm.bcast(out_dir, root=0) + return str(out_dir) + + +# --------------------------------------------------------------------------- +# Test 1 +# --------------------------------------------------------------------------- + +def test_evaluate_on_swarm_empty_ranks(tmp_path_factory): + """``uw.function.evaluate`` must not hang when some ranks have 0 particles. + + Only rank 0 adds particles. All ranks call + ``uw.function.evaluate(expr, swarm.coords)`` — the PETSc DMLocatePoints + path is collective and must complete even with empty coordinate arrays. + """ + mesh = _make_mesh() + out_dir = _tmp_outdir(tmp_path_factory) + + x, y = mesh.X + + swarm = uw.swarm.Swarm(mesh=mesh) + var = swarm.add_variable(name="val", size=1) + + # All ranks participate in add_particles_with_coordinates (collective) + coords = _empty_rank_coords(mesh) + swarm.add_particles_with_coordinates(coords) + + # Set a known value so we can verify after evaluate + if swarm.local_size > 0: + var.data[:, 0] = np.arange(swarm.local_size, dtype=float) + + uw.mpi.comm.barrier() + + # --- The key operation: evaluate a SymPy expression on swarm coords --- + # This is collective via PETSc DMLocatePoints. + expr = x + y + result = uw.function.evaluate(expr, swarm.coords) + + uw.mpi.comm.barrier() + + # Verify on the rank that has particles + if swarm.local_size > 0: + assert result is not None, "evaluate returned None" + assert result.shape[0] == swarm.local_size + # All result values should be finite numbers + assert np.all(np.isfinite(result)), "evaluate produced non-finite values" + + # Read back the ACTUAL local coordinates (particles redistribute across + # ranks after add_particles_with_coordinates) and verify evaluate(x+y) + # matches the coordinate sum at each particle's true position. + actual_coords = swarm.data # (local_n, dim) read-only snapshot + expected = actual_coords[:, 0] + actual_coords[:, 1] + got = np.asarray(result).reshape(-1) + np.testing.assert_allclose( + got, expected, atol=1e-10, + err_msg="evaluate(x+y) does not match coordinate sum", + ) + + uw.mpi.comm.barrier() + + # Cleanup + if uw.mpi.rank == 0: + import shutil + shutil.rmtree(out_dir, ignore_errors=True) + + del swarm, mesh + + +# --------------------------------------------------------------------------- +# Test 2 +# --------------------------------------------------------------------------- + +def test_write_timestep_empty_ranks(tmp_path_factory): + """``swarm.write_timestep`` must complete when some ranks have 0 particles. + + Only rank 0 adds particles; the HDF5 collective close must synchronise + even with heterogeneous local sizes. + """ + mesh = _make_mesh() + out_dir = _tmp_outdir(tmp_path_factory) + + swarm = uw.swarm.Swarm(mesh=mesh) + var = swarm.add_variable(name="material", size=1) + + coords = _empty_rank_coords(mesh) + swarm.add_particles_with_coordinates(coords) + + if swarm.local_size > 0: + var.data[:, 0] = 1.0 + + uw.mpi.comm.barrier() + + # Verify we have the empty-rank distribution we expect + sizes = uw.mpi.comm.allgather(swarm.local_size) + if uw.mpi.size > 1: + assert 0 in sizes, f"Expected at least one empty rank, got {sizes}" + assert sum(sizes) > 0, "All ranks are empty — test is vacuous" + + # --- Collective save: all ranks must call this --- + swarm.write_timestep( + filename="swarm", + swarmname="swarm", + index=0, + outputPath=out_dir, + swarmVars=[var], + ) + + uw.mpi.comm.barrier() + + # Verify file existence and content on rank 0 + expected_h5 = os.path.join(out_dir, "swarm.swarm.00000.h5") + expected_var = os.path.join(out_dir, "swarm.swarm.material.00000.h5") + expected_xdmf = os.path.join(out_dir, "swarm.swarm.00000.xdmf") + if uw.mpi.rank == 0: + assert os.path.exists(expected_h5), f"Missing coords HDF5: {expected_h5}" + assert os.path.exists(expected_var), f"Missing var HDF5: {expected_var}" + assert os.path.exists(expected_xdmf), f"Missing XDMF: {expected_xdmf}" + + import h5py + with h5py.File(expected_h5, "r") as f: + n_global = f["coordinates"].shape[0] + assert n_global == sum(sizes), ( + f"saved coords shape {n_global} != sum of local sizes {sum(sizes)}" + ) + + uw.mpi.comm.barrier() + + # Cleanup + if uw.mpi.rank == 0: + import shutil + shutil.rmtree(out_dir, ignore_errors=True) + + del swarm, mesh + + +# --------------------------------------------------------------------------- +# Test 3 +# --------------------------------------------------------------------------- + +def test_evaluate_and_save_combined_empty_ranks(tmp_path_factory): + """Evaluate multiple expressions, assign to swarm variables, then save. + + This exercises the user's pattern in the round-trip: evaluate → assign → + write_timestep, with empty ranks throughout. + """ + mesh = _make_mesh() + out_dir = _tmp_outdir(tmp_path_factory) + + x, y = mesh.X + + swarm = uw.swarm.Swarm(mesh=mesh) + var_T = swarm.add_variable(name="temperature", size=1) + var_P = swarm.add_variable(name="pressure", size=1) + var_S = swarm.add_variable(name="strain", size=1) + + coords = _empty_rank_coords(mesh) + swarm.add_particles_with_coordinates(coords) + + uw.mpi.comm.barrier() + + # --- Evaluate three different expressions collectively --- + expr_T = 300.0 + 100.0 * x * y + expr_P = 1.0e5 * (1.0 - y) + expr_S = x * x + y * y + + result_T = uw.function.evaluate(expr_T, swarm.coords) + result_P = uw.function.evaluate(expr_P, swarm.coords) + result_S = uw.function.evaluate(expr_S, swarm.coords) + + uw.mpi.comm.barrier() + + # --- Assign results to swarm variables (only on ranks with particles) --- + # evaluate() returns a (n, 1) column array; flatten to (n,) to match the + # 1D per-column shape of swarm variable data (same as the reference test). + if swarm.local_size > 0: + var_T.data[:, 0] = np.asarray(result_T).reshape(-1) + var_P.data[:, 0] = np.asarray(result_P).reshape(-1) + var_S.data[:, 0] = np.asarray(result_S).reshape(-1) + + uw.mpi.comm.barrier() + + # --- Verify evaluate values against actual local particle coordinates --- + if swarm.local_size > 0: + coords_arr = swarm.data + np.testing.assert_allclose( + np.asarray(result_T).reshape(-1), + 300.0 + 100.0 * coords_arr[:, 0] * coords_arr[:, 1], + atol=1e-10, + ) + np.testing.assert_allclose( + np.asarray(result_P).reshape(-1), + 1.0e5 * (1.0 - coords_arr[:, 1]), + atol=1e-6, + ) + np.testing.assert_allclose( + np.asarray(result_S).reshape(-1), + coords_arr[:, 0] ** 2 + coords_arr[:, 1] ** 2, + atol=1e-10, + ) + + # --- Collective save with all three variables --- + swarm.write_timestep( + filename="swarm", + swarmname="swarm", + index=0, + outputPath=out_dir, + swarmVars=[var_T, var_P, var_S], + ) + + uw.mpi.comm.barrier() + + # --- Verify files --- + # NOTE: allgather is collective — call on ALL ranks, verify on rank 0. + total_local = sum(uw.mpi.comm.allgather(swarm.local_size)) + if uw.mpi.rank == 0: + assert os.path.exists(os.path.join(out_dir, "swarm.swarm.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "swarm.swarm.temperature.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "swarm.swarm.pressure.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "swarm.swarm.strain.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "swarm.swarm.00000.xdmf")) + + # Global particle count must be preserved across the redistribute that + # happens inside add_particles_with_coordinates. + import h5py + with h5py.File(os.path.join(out_dir, "swarm.swarm.00000.h5"), "r") as f: + n_global = f["coordinates"].shape[0] + assert n_global == total_local, ( + f"Expected {total_local} global particles, got {n_global}" + ) + + uw.mpi.comm.barrier() + + # Cleanup + if uw.mpi.rank == 0: + import shutil + shutil.rmtree(out_dir, ignore_errors=True) + + del swarm, mesh + + +# --------------------------------------------------------------------------- +# Test 4 +# --------------------------------------------------------------------------- + +def test_passive_swarm_save_empty_ranks(tmp_path_factory): + """Simulate the user's passive-tracer pattern that hangs on HPC. + + The user creates a passive swarm with variables (T, p, time, uid), adds + particles only on rank 0, creates a mesh temperature field, evaluates the + mesh variable symbol on swarm coords, assigns to swarm variables, and + saves with write_timestep. Every step is collective — if any rank skips + a call the MPI layer deadlocks. + + NOTE on ``mode``: the mesh-variable evaluation is done with + ``mode="fast"`` (RBF interpolation) because the default L2-projection path + deadlocks when a rank owns zero particles — the projection is a collective + Schur-complement solve over the swarm partition and the empty rank never + reaches the barrier (TODO: track in planning file / fix the projection + path). ``mode="fast"`` exercises the full user round-trip (mesh-var + evaluate -> assign -> save) while keeping the test a valid regression for + the empty-rank evaluate+save combination. + """ + mesh = _make_mesh() + out_dir = _tmp_outdir(tmp_path_factory) + + x, y = mesh.X + + # --- Mesh temperature field (source for evaluation) --- + T_mesh = uw.discretisation.MeshVariable("T_field", mesh, 1, degree=1) + # Set a known temperature distribution on the mesh + with mesh.access(T_mesh): + T_mesh.data[:, 0] = 300.0 + 100.0 * T_mesh.coords[:, 0] * T_mesh.coords[:, 1] + + uw.mpi.comm.barrier() + + # --- Passive swarm (user's pattern) --- + swarm = uw.swarm.Swarm(mesh=mesh) + var_T = swarm.add_variable(name="temperature", size=1) + var_p = swarm.add_variable(name="pressure", size=1) + var_time = swarm.add_variable(name="time", size=1) + var_uid = swarm.add_variable(name="uid", size=1, dtype=int) + + # Only rank 0 adds particles (user's crust-only pattern) + coords = _empty_rank_coords(mesh) + swarm.add_particles_with_coordinates(coords) + + uw.mpi.comm.barrier() + + # --- Evaluate mesh variable on swarm coords --- + # This uses T_mesh.sym which produces a SymPy expression referencing the + # mesh variable. Evaluate is collective via PETSc DMLocatePoints; ``fast`` + # (RBF) avoids the L2-projection deadlock on empty ranks (see docstring). + result_T = uw.function.evaluate(T_mesh.sym, swarm.coords, mode="fast") + result_p = uw.function.evaluate(y * 1.0e6, swarm.coords, mode="fast") + result_time = uw.function.evaluate(x * 0.0, swarm.coords, mode="fast") + + uw.mpi.comm.barrier() + + # --- Assign to swarm variables (only ranks with particles) --- + if swarm.local_size > 0: + var_T.data[:, 0] = np.asarray(result_T).reshape(-1) + var_p.data[:, 0] = np.asarray(result_p).reshape(-1) + var_time.data[:, 0] = 100.0 # fixed time value + var_uid.data[:, 0] = np.arange(swarm.local_size, dtype=int) + + uw.mpi.comm.barrier() + + # --- Verify on the rank(s) that have particles --- + if swarm.local_size > 0: + coords_arr = swarm.data + expected_T = 300.0 + 100.0 * coords_arr[:, 0] * coords_arr[:, 1] + # mode="fast" evaluates the mesh field via RBF interpolation, so allow + # a modest interpolation tolerance (observed error ~ a few units on a + # field of ~300-400). + np.testing.assert_allclose( + np.asarray(result_T).reshape(-1), expected_T, atol=10.0, + err_msg="Temperature evaluation mismatch", + ) + np.testing.assert_allclose( + np.asarray(result_p).reshape(-1), + coords_arr[:, 1] * 1.0e6, + atol=1e-3, + err_msg="Pressure evaluation mismatch", + ) + assert np.all(var_uid.data[:, 0] == np.arange(swarm.local_size)), ( + "UID assignment mismatch" + ) + + # --- Collective save (the user's write_timestep call) --- + swarm.write_timestep( + filename="passive_tracers", + swarmname="tracers", + index=0, + outputPath=out_dir, + swarmVars=[var_T, var_p, var_time, var_uid], + ) + + uw.mpi.comm.barrier() + + # --- Verify files --- + # NOTE: allgather is collective — call on ALL ranks, verify on rank 0. + total_local = sum(uw.mpi.comm.allgather(swarm.local_size)) + if uw.mpi.rank == 0: + expected_h5 = os.path.join(out_dir, "passive_tracers.tracers.00000.h5") + assert os.path.exists(expected_h5), f"Missing coords HDF5: {expected_h5}" + assert os.path.exists(os.path.join(out_dir, "passive_tracers.tracers.temperature.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "passive_tracers.tracers.pressure.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "passive_tracers.tracers.time.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "passive_tracers.tracers.uid.00000.h5")) + assert os.path.exists(os.path.join(out_dir, "passive_tracers.tracers.00000.xdmf")) + + import h5py + with h5py.File(expected_h5, "r") as f: + n_global = f["coordinates"].shape[0] + assert n_global == total_local, ( + f"Expected {total_local} global particles, got {n_global}" + ) + + uw.mpi.comm.barrier() + + # Cleanup + if uw.mpi.rank == 0: + import shutil + shutil.rmtree(out_dir, ignore_errors=True) + + del swarm, mesh diff --git a/tests/parallel/test_0796_swarm_advection_empty_rank.py b/tests/parallel/test_0796_swarm_advection_empty_rank.py new file mode 100644 index 000000000..936fef965 --- /dev/null +++ b/tests/parallel/test_0796_swarm_advection_empty_rank.py @@ -0,0 +1,76 @@ +"""MPI regression test for swarm.advection with empty ranks. + +A passive swarm whose particles are confined to a subset of ranks (e.g. the +user's crust-only tracers) leaves the remaining ranks holding zero particles. +Advection on such a swarm must not deadlock or crash: + +1. The substep loop's ``global_evaluate(V_fn_matrix, particle_data)`` on an + empty rank used to block inside the collective point-location machinery. + The root cause (an empty rank taking a divergent DMLocatePoints branch) is + fixed upstream (issue #611 / PR #656), which this test guards against + regressing. +2. ``estimate_dt()`` reshaped the (empty) velocity array with + ``reshape(n, -1)``, which NumPy cannot infer from zero elements -> + ``ValueError: cannot reshape array of size 0 into shape (0,newaxis)``. + Guarded in ``Swarm.estimate_dt``; this test exercises the default + (non-``evalf``) path in ``order=2`` which runs ``estimate_dt``. +""" + +import numpy as np +import pytest + +import underworld3 as uw + + +pytestmark = [ + pytest.mark.level_2, + pytest.mark.mpi(min_size=2), + pytest.mark.timeout(60), +] + + +def test_advection_empty_rank_default(tmp_path_factory): + """``swarm.advection(v.sym, ...)`` completes with empty ranks (default path). + + Only rank 0 holds particles; the other rank(s) hold zero. Exercises the + default FE-interpolation ``global_evaluate`` path inside the ``order=2`` + substep loop plus the ``estimate_dt`` empty-rank guard. + """ + mesh = uw.meshing.UnstructuredSimplexBox( + minCoords=(0.0, 0.0), maxCoords=(1.0, 1.0), cellSize=0.25, + ) + + # Simple shear velocity: u = y, v = 0 + v = uw.discretisation.MeshVariable("v", mesh, mesh.dim, degree=1) + with mesh.access(v): + v.data[:, 0] = mesh.X.coords[:, 1] + v.data[:, 1] = 0.0 + + swarm = uw.swarm.Swarm(mesh=mesh) + + # Particles only on rank 0 (collective call) + if uw.mpi.rank == 0: + coords = (np.random.rand(100, mesh.dim) * 0.8 + 0.1) + else: + coords = np.empty((0, mesh.dim)) + swarm.add_particles_with_coordinates(coords) + + uw.mpi.comm.barrier() + + # Sanity: the empty-rank precondition must actually hold (all 100 points + # are added on rank 0 and no migration happens here, so the other ranks + # hold zero particles). + sizes = uw.mpi.comm.allgather(swarm.local_size) + assert sum(sizes) > 0, "No particles at all - test is vacuous" + assert 0 in sizes, ( + f"Expected at least one empty rank, got local sizes {sizes} " + "(test would no longer exercise the empty-rank path)" + ) + + # Exercise the empty-rank path with the DEFAULT (non-evalf) FE + # interpolation, which rides global_evaluate. Pre-#611 this deadlocked; + # pre (estimate_dt) guard the order=2 estimate_dt reshape crashed. + swarm.advection(v.sym, delta_t=1.0, order=2) + + uw.mpi.comm.barrier() + del swarm, mesh