Fix point evaluation API change (#4675)

Author: leo-collins

firedrake/function.py

@@ -549,8 +549,12 @@ def evaluate(self, coord, mapping, component, index_values):
         # Called by UFL when evaluating expressions at coordinates
         if component or index_values:
             raise NotImplementedError("Unsupported arguments when attempting to evaluate Function.")
+        coord = np.asarray(coord, dtype=utils.ScalarType)
         evaluator = PointEvaluator(self.function_space().mesh(), coord)
-        return evaluator.evaluate(self)
+        result = evaluator.evaluate(self)
+        if len(coord.shape) == 1:
+            result = result.squeeze(axis=0)
+        return result
 
     def at(self, arg, *args, **kwargs):
         warnings.warn(
@@ -809,7 +813,7 @@ def evaluate(self, function: Function) -> np.ndarray | Tuple[np.ndarray, ...]:
         f_at_points = assemble(interpolate(function, P0DG))
         f_at_points_io = Function(P0DG_io).assign(np.nan)
         f_at_points_io.interpolate(f_at_points)
-        result = f_at_points_io.dat.data_ro
+        result = f_at_points_io.dat.data_ro.copy()
 
         # If redundant, all points are now on rank 0, so we broadcast the result
         if self.redundant and self.mesh.comm.size > 1:

tests/firedrake/regression/test_point_eval_api.py

@@ -1,4 +1,5 @@
 from os.path import abspath, dirname
+from numbers import Number
 import numpy as np
 import pytest
 
@@ -184,23 +185,33 @@ def test_point_evaluator_scalar(mesh_and_points):
     # Test standard scalar function evaluation at points
     f_at_points = evaluator.evaluate(f)
     assert np.allclose(f_at_points, [0.2, 0.4, 0.6])
+    assert isinstance(f_at_points, np.ndarray)
+    assert f_at_points.shape == (len(evaluator.points),) + f.ufl_shape
+    assert isinstance(f_at_points[0], Number)
 
     # Test standard scalar function with missing points
     eval_missing = PointEvaluator(mesh, np.append(points, [[1.5, 1.5]], axis=0), missing_points_behaviour="ignore")
     f_at_points_missing = eval_missing.evaluate(f)
     assert np.isnan(f_at_points_missing[-1])
 
+    # Can modify result
+    f_at_points *= 2.0
+    assert np.allclose(f_at_points, [0.4, 0.8, 1.2])
+
 
 @pytest.mark.parallel([1, 3])
 def test_point_evaluator_vector_tensor_mixed(mesh_and_points):
     mesh, evaluator = mesh_and_points
     V_vec = VectorFunctionSpace(mesh, "CG", 1)
-    f_vec = Function(V_vec)
     x, y = SpatialCoordinate(mesh)
-    f_vec.interpolate(as_vector([x, y]))
+    f_vec = Function(V_vec).interpolate(as_vector([x, y]))
     f_vec_at_points = evaluator.evaluate(f_vec)
     vec_expected = np.array([[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]])
     assert np.allclose(f_vec_at_points, vec_expected)
+    assert isinstance(f_vec_at_points, np.ndarray)
+    assert f_vec_at_points.shape == (len(evaluator.points),) + f_vec.ufl_shape
+    assert isinstance(f_vec_at_points[0, 0], Number)
+    assert isinstance(f_vec_at_points[0, :], np.ndarray)
 
     V_tensor = TensorFunctionSpace(mesh, "CG", 1, shape=(2, 3))
     f_tensor = Function(V_tensor)
@@ -210,15 +221,25 @@ def test_point_evaluator_vector_tensor_mixed(mesh_and_points):
                                 [[0.2, 0.2, 0.04], [0.2, 0.2, 0.04]],
                                 [[0.3, 0.3, 0.09], [0.3, 0.3, 0.09]]])
     assert np.allclose(f_tensor_at_points, tensor_expected)
+    assert f_tensor_at_points.shape == (len(evaluator.points),) + f_tensor.ufl_shape
+    assert isinstance(f_tensor_at_points[0, 0, 0], Number)
+    assert isinstance(f_tensor_at_points[0, 0, :], np.ndarray)
+    assert isinstance(f_tensor_at_points[0, :, :], np.ndarray)
 
     V_mixed = V_vec * V_tensor
     f_mixed = Function(V_mixed)
     f_vec, f_tensor = f_mixed.subfunctions
     f_vec.interpolate(as_vector([x, y]))
     f_tensor.interpolate(as_matrix([[x, y, x*y], [y, x, x*y]]))
     f_mixed_at_points = evaluator.evaluate(f_mixed)
+    assert isinstance(f_mixed_at_points, tuple)
+    assert len(f_mixed_at_points) == 2
     assert np.allclose(f_mixed_at_points[0], vec_expected)
     assert np.allclose(f_mixed_at_points[1], tensor_expected)
+    assert isinstance(f_mixed_at_points[0], np.ndarray)
+    assert isinstance(f_mixed_at_points[1], np.ndarray)
+    assert f_mixed_at_points[0].shape == (len(evaluator.points),) + f_vec.ufl_shape
+    assert f_mixed_at_points[1].shape == (len(evaluator.points),) + f_tensor.ufl_shape
 
 
 @pytest.mark.parallel(3)
@@ -286,58 +307,98 @@ def test_point_evaluator_tolerance():
 
 def test_point_evaluator_inputs_1d():
     mesh = UnitIntervalMesh(1)
-    f = mesh.coordinates
+    f = mesh.coordinates  # ufl_shape (1,)
 
     # one point
     for input in [0.2, (0.2,), [0.2], np.array([0.2])]:
         e = PointEvaluator(mesh, input)
-        assert np.allclose(0.2, e.evaluate(f))
+        res = e.evaluate(f)
+        assert np.allclose(0.2, res)
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (1,)
+        assert isinstance(res[0], Number)
 
     # multiple points as tuples/list
     for input in [
         (0.2, 0.3), ((0.2,), (0.3,)), ([0.2], [0.3]),
         (np.array(0.2), np.array(0.3)), (np.array([0.2]), np.array([0.3]))
     ]:
         e2 = PointEvaluator(mesh, input)
-        assert np.allclose([[0.2, 0.3]], e2.evaluate(f))
+        res = e2.evaluate(f)
+        assert np.allclose([[0.2, 0.3]], res)
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (len(input),)
+        assert isinstance(res[0], Number)
+
         e3 = PointEvaluator(mesh, list(input))
-        assert np.allclose([[0.2, 0.3]], e3.evaluate(f))
+        res2 = e3.evaluate(f)
+        assert np.allclose([[0.2, 0.3]], res2)
+        assert isinstance(res2, np.ndarray)
+        assert res2.shape == (len(input),)
+        assert isinstance(res2[0], Number)
 
     # multiple points as numpy array
     for input in [np.array([0.2, 0.3]), np.array([[0.2], [0.3]])]:
         e = PointEvaluator(mesh, input)
-        assert np.allclose([[0.2, 0.3]], e.evaluate(f))
+        res = e.evaluate(f)
+        assert np.allclose([[0.2, 0.3]], res)
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (len(input),)
+        assert isinstance(res[0], Number)
 
     # test incorrect inputs
     for input in [[[0.2, 0.3]], ([0.2, 0.3], [0.4, 0.5]), np.array([[0.2, 0.3]])]:
-        with pytest.raises(ValueError):
+        with pytest.raises(ValueError, match=r"Point dimension \(2\) does not match geometric dimension \(1\)."):
             PointEvaluator(mesh, input)
 
 
 def test_point_evaluator_inputs_2d():
     mesh = UnitSquareMesh(1, 1)
-    f = mesh.coordinates
+    f = mesh.coordinates  # ufl_shape (2,)
 
     # one point
     for input in [(0.2, 0.4), [0.2, 0.4], [[0.2, 0.4]], np.array([0.2, 0.4])]:
         e = PointEvaluator(mesh, input)
-        assert np.allclose([0.2, 0.4], e.evaluate(f))
+        res = e.evaluate(f)
+        assert np.allclose([0.2, 0.4], res)
+        assert isinstance(res, np.ndarray)
+        assert res.shape == (1,) + f.ufl_shape
+        assert isinstance(res[0], np.ndarray)
+        assert isinstance(res[0, 0], Number)
 
     # multiple points as tuple
     for input in [
         ((0.2, 0.4), (0.3, 0.5)), ([0.2, 0.4], [0.3, 0.5]),
         (np.array([0.2, 0.4]), np.array([0.3, 0.5]))
     ]:
         e1 = PointEvaluator(mesh, input)
-        assert np.allclose([[0.2, 0.4], [0.3, 0.5]], e1.evaluate(f))
+        res1 = e1.evaluate(f)
+        assert np.allclose([[0.2, 0.4], [0.3, 0.5]], res1)
+        assert isinstance(res1, np.ndarray)
+        assert res1.shape == (len(input),) + f.ufl_shape
+        assert isinstance(res1[0], np.ndarray)
+        assert isinstance(res1[0, 0], Number)
+
         e2 = PointEvaluator(mesh, list(input))
         assert np.allclose([[0.2, 0.4], [0.3, 0.5]], e2.evaluate(f))
 
     # multiple points as numpy array
-    e = PointEvaluator(mesh, np.array([[0.2, 0.4], [0.3, 0.5]]))
-    assert np.allclose([[0.2, 0.4], [0.3, 0.5]], e.evaluate(f))
+    points = np.array([[0.2, 0.4], [0.3, 0.5]])
+    e = PointEvaluator(mesh, points)
+    res = e.evaluate(f)
+    assert np.allclose([[0.2, 0.4], [0.3, 0.5]], res)
+    assert isinstance(res, np.ndarray)
+    assert res.shape == (len(points),) + f.ufl_shape
+    assert isinstance(res[0], np.ndarray)
+    assert isinstance(res[0, 0], Number)
+
+    res2 = e.evaluate(f)
+    res3 = res + res2
+    assert np.allclose([[0.4, 0.8], [0.6, 1.0]], res3)
+    assert isinstance(res3, np.ndarray)
+    assert res3.shape == (len(points),) + f.ufl_shape
 
     # test incorrect inputs
     for input in [0.2, [0.2]]:
-        with pytest.raises(ValueError):
+        with pytest.raises(ValueError, match=r"Point dimension \(1\) does not match geometric dimension \(2\)."):
             PointEvaluator(mesh, input)

tests/firedrake/regression/test_point_eval_fs.py

@@ -1,4 +1,5 @@
 from os.path import abspath, dirname
+from numbers import Number
 import numpy as np
 import pytest
 
@@ -104,6 +105,11 @@ def test_triangle_tensor(mesh_triangle, family, degree):
 
     assert np.allclose([[0.4, 0.8], [0.48, 0.08]], f([0.6, 0.4]))
     assert np.allclose([[0.9, 0.2], [0.00, 0.18]], f([0.0, 0.9]))
+    res = f([0.1, 0.2])
+    assert isinstance(res, np.ndarray)
+    assert res.shape == (2, 2)
+    assert isinstance(res[0, :], np.ndarray)
+    assert isinstance(res[0, 0], Number)
 
 
 def test_triangle_mixed(mesh_triangle):
@@ -143,6 +149,10 @@ def test_quadrilateral(mesh_quadrilateral, family, degree):
     f = Function(V).interpolate((x[0] - 0.5)*(x[1] - 0.2))
     assert np.allclose(+0.02, f([0.6, 0.4]))
     assert np.allclose(-0.35, f([0.0, 0.9]))
+    res = f([0.1, 0.2])
+    assert isinstance(res, np.ndarray)
+    assert len(res.shape) == 0
+    assert isinstance(res.item(), Number)
 
 
 @pytest.mark.parametrize(('family', 'degree'),
@@ -161,6 +171,10 @@ def test_quadrilateral_vector(mesh_quadrilateral, family, degree):
 
     assert np.allclose([0.6, 0.56], f([0.6, 0.4]))
     assert np.allclose([1.1, 0.18], f([0.0, 0.9]))
+    res = f([0.1, 0.2])
+    assert isinstance(res, np.ndarray)
+    assert len(res.shape) == 1
+    assert isinstance(res[0], Number)
 
 
 @pytest.mark.parametrize(('family', 'degree'),
@@ -172,6 +186,10 @@ def test_tetrahedron(mesh_tetrahedron, family, degree):
     f = Function(V).interpolate((x[0] - 0.5)*(x[1] - x[2]))
     assert np.allclose(+0.01, f([0.6, 0.4, 0.3]))
     assert np.allclose(-0.06, f([0.4, 0.7, 0.1]))
+    res = f([0.2, 0.3, 0.4])
+    assert isinstance(res, np.ndarray)
+    assert len(res.shape) == 0
+    assert isinstance(res.item(), Number)
 
 
 @pytest.mark.parametrize(('family', 'degree'),
@@ -192,6 +210,10 @@ def test_tetrahedron_vector(mesh_tetrahedron, family, degree):
 
     assert np.allclose([0.6, 0.54, 0.4], f([0.6, 0.4, 0.3]))
     assert np.allclose([0.9, 0.34, 0.7], f([0.4, 0.7, 0.1]))
+    res = f([0.2, 0.3, 0.4])
+    assert isinstance(res, np.ndarray)
+    assert len(res.shape) == 1
+    assert isinstance(res[0], Number)
 
 
 def test_point_eval_forces_writes():