diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 6c785af0..21fcf641 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -347,6 +347,16 @@ def _lambdify_and_evaluate(expr, coords, interpolated_results, coord_sys=None, m return results.reshape(-1, *shape) +def _global_fallback_indices(return_value, return_mask): + """Indices requiring the parallel best-claim fallback. + + A point needs recovery when migration/location marks it extrapolated or + when interpolation returned a non-finite value despite a located flag. + """ + nonfinite = ~np.isfinite(return_value).all(axis=(1, 2)) + return np.where(return_mask[:, 0, 0] | nonfinite)[0] + + def global_evaluate_nd( expr, coords=None, coord_sys=None, @@ -589,7 +599,12 @@ def global_evaluate_nd( expr, from mpi4py import MPI comm = uw.mpi.comm - ext_idx = np.where(return_mask[:, 0, 0])[0] + # A failed interpolation can occasionally return NaN while reporting + # the point as located. Treat that exactly like an extrapolated/lost + # point so a finite value from the globally nearest rank replaces it. + # This is required by SLCN midpoint tracing: one silent NaN here makes + # the departure point and then the transported history non-finite. + ext_idx = _global_fallback_indices(return_value, return_mask) ext_coords = np.ascontiguousarray(coords_array[ext_idx], dtype=np.float64) counts = np.array(comm.allgather(ext_coords.shape[0]), dtype=int) diff --git a/tests/parallel/test_0760_swarm_cache_migration.py b/tests/parallel/test_0760_swarm_cache_migration.py index d27475af..b8b02953 100644 --- a/tests/parallel/test_0760_swarm_cache_migration.py +++ b/tests/parallel/test_0760_swarm_cache_migration.py @@ -17,6 +17,7 @@ import pytest import numpy as np import underworld3 as uw +from underworld3.function._function import _global_fallback_indices from mpi4py import MPI pytestmark = [pytest.mark.mpi(min_size=2), pytest.mark.timeout(60)] @@ -101,3 +102,18 @@ def test_global_evaluate_displaced_nodes(): f"Rank {uw.mpi.rank}: expected {node_coords.shape[0]} results, " f"got {result.shape[0]}" ) + + +@pytest.mark.level_1 +@pytest.mark.tier_a +def test_global_fallback_includes_nonfinite_located_values(): + """Located NaN values must enter the same recovery path as lost points.""" + values = np.ones((4, 1, 2)) + values[1, 0, 0] = np.nan + mask = np.zeros((4, 1, 1), dtype=bool) + mask[2, 0, 0] = True + + assert np.array_equal( + _global_fallback_indices(values, mask), + np.array([1, 2]), + )