From f9188c023885b45269435efb074a9fecf529e0f5 Mon Sep 17 00:00:00 2001 From: Henrik Finsberg Date: Sat, 14 Nov 2026 11:00:49 +0000 Subject: [PATCH] Fix evaluate_function in parallel: keep nan, nan outside the mesh, one reduction With broadcast=True, every rank filled the points it did not find with -inf and reduced each point with comm.allreduce(value, op=MPI.MAX). That is Python's max on floats, which is order dependent for nan (max(nan, -inf) is nan, max(-inf, nan) is -inf), so a nan value came back as -inf unless rank 0 found the point. Points outside the mesh came back as -inf, and the per-point allreduce made many points slow. Each point is now owned by the lowest rank that finds it (one Allreduce MIN), only the owner contributes its value (one Allreduce SUM), and points that no rank finds are nan. All ranks return the same values, also for discontinuous functions (the owner's value instead of the maximum). The output is sized by the value shape instead of the dofmap block size, which also fixes spaces such as Raviart-Thomas (block size 1, value size 2). Co-Authored-By: Claude Opus 5.5 --- src/scifem/eval.py | 35 ++++++++++------------ tests/test_eval.py | 74 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+), 19 deletions(-) diff --git a/src/scifem/eval.py b/src/scifem/eval.py index 5c5fa5a..e724fc8 100644 --- a/src/scifem/eval.py +++ b/src/scifem/eval.py @@ -58,9 +58,10 @@ def evaluate_function( be called on all active processes synchronously. Note: - If the function is discontinuous, different processes may return - different values for the same point. - In this case, the value returned is the maximum value across all processes. + A point found on several processes (e.g. on a process boundary) + takes its value from the lowest-ranked of them, so all processes + return the same values, also for discontinuous functions. + Points that are not found on any process are ``nan``. Returns: The values of the function evaluated at the points. @@ -92,22 +93,18 @@ def evaluate_function( values = u.eval(points_on_proc, cells) if broadcast: - bs = u.function_space.dofmap.index_map_bs - # Create array to store values and fill with -inf - # to ensure that all points are included in the allreduce - # with op=MPI.MAX - u_out = np.ones((num_points, bs), dtype=np.float64) * -np.inf - # Fill in values for points on this process - u_out[indices, :] = values - # Now loop over all processes and find the maximum value - for i in range(num_points): - if bs > 1: - # If block size is larger than 1, loop over blocks - for j in range(bs): - u_out[i, j] = comm.allreduce(u_out[i, j], op=MPI.MAX) - else: - u_out[i] = comm.allreduce(u_out[i], op=MPI.MAX) - + value_size = int(np.prod(u.ufl_shape)) + # Each point is owned by the lowest rank that finds it (comm.size if none does) + owner = np.full(num_points, comm.size, dtype=np.int32) + owner[indices] = comm.rank + comm.Allreduce(MPI.IN_PLACE, owner, op=MPI.MIN) + # Only the owner contributes the value of a point, so summing over ranks + # passes every value (including nan) through unchanged + owned = owner[indices] == comm.rank + u_out = np.zeros((num_points, value_size), dtype=np.result_type(values.dtype, np.float64)) + u_out[indices[owned], :] = np.reshape(values, (len(indices), value_size))[owned] + comm.Allreduce(MPI.IN_PLACE, u_out, op=MPI.SUM) + u_out[owner == comm.size, :] = np.nan return u_out else: return values diff --git a/tests/test_eval.py b/tests/test_eval.py index 26d3c20..c99f6b6 100644 --- a/tests/test_eval.py +++ b/tests/test_eval.py @@ -117,6 +117,80 @@ def test_evaluate_vector_function_3D(cell_type): assert np.allclose(u_values, exact) +@pytest.mark.parametrize("shape", [(), (2,)]) +def test_evaluate_function_keeps_nan(shape): + # NaN must come back as NaN on any number of ranks, wherever the point is owned + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("P", 1, shape)) + u = dolfinx.fem.Function(V) + u.x.array[:] = np.nan + + x = np.linspace(0.05, 0.95, 5) + points = np.array([[xi, yi] for xi in x for yi in x]) + u_values = evaluate_function(u, points) + + assert u_values.shape == (len(points), int(np.prod(shape))) + assert np.isnan(u_values).all() + + +def test_evaluate_function_value_size_differs_from_block_size(): + # Raviart-Thomas: block size 1, value size 2 + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 4, 4) + V = dolfinx.fem.functionspace(mesh, ("RT", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: (np.ones_like(x[0]), 2 * np.ones_like(x[0]))) + + u_values = evaluate_function(u, np.array([[0.3, 0.4], [0.8, 0.1]])) + + assert np.allclose(u_values, [[1.0, 2.0], [1.0, 2.0]]) + + +def test_evaluate_function_nan_next_to_values(): + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("P", 1)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: np.where(x[0] < 0.49, np.nan, x[0] + 2 * x[1])) + + x = np.linspace(0.05, 0.95, 7) + points = np.array([[xi, yi] for xi in x for yi in x]) + u_values = evaluate_function(u, points)[:, 0] + + left = points[:, 0] < 0.37 # every vertex of the containing cell is NaN + right = points[:, 0] > 0.63 # no vertex of the containing cell is NaN + assert np.isnan(u_values[left]).all() + assert np.allclose(u_values[right], points[right, 0] + 2 * points[right, 1]) + + +def test_evaluate_function_point_outside_mesh_is_nan(): + mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("P", 1)) + u = dolfinx.fem.Function(V) + u.x.array[:] = 1.0 + + u_values = evaluate_function(u, np.array([[0.5, 0.5], [2.0, 2.0]]))[:, 0] + + assert np.isclose(u_values[0], 1.0) + assert np.isnan(u_values[1]) + + +def test_evaluate_function_same_result_on_all_ranks_for_discontinuous_function(): + # Points on facets lie in cells of several ranks; every rank must return the same value + comm = MPI.COMM_WORLD + mesh = dolfinx.mesh.create_unit_square(comm, 8, 8) + V = dolfinx.fem.functionspace(mesh, ("DG", 0)) + u = dolfinx.fem.Function(V) + u.interpolate(lambda x: np.sin(10 * x[0]) + x[1]) + + x = np.linspace(0.0, 1.0, 9) # the vertices of the mesh + points = np.array([[xi, yi] for xi in x for yi in x]) + u_values = evaluate_function(u, points) + + all_values = comm.allgather(u_values) + for other in all_values: + np.testing.assert_array_equal(other, all_values[0]) + assert np.isfinite(u_values).all() + + @pytest.mark.parametrize("degree", [5]) @pytest.mark.parametrize("extrema", [min, max]) @pytest.mark.parametrize(