Fix evaluate_function in parallel: keep nan, nan outside the mesh, one reduction

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

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(