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(