diff --git a/firedrake/function.py b/firedrake/function.py index a628ac6599..57a21e03cf 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -550,7 +550,10 @@ def evaluate(self, coord, mapping, component, index_values): if component or index_values: raise NotImplementedError("Unsupported arguments when attempting to evaluate Function.") 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,14 +812,14 @@ 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: if self.mesh.comm.rank != 0: result = np.empty((len(self.points),) + shape, dtype=utils.ScalarType) self.mesh.comm.Bcast(result) - return result + return result.reshape((-1,) + shape) @PETSc.Log.EventDecorator() diff --git a/tests/firedrake/regression/test_point_eval_api.py b/tests/firedrake/regression/test_point_eval_api.py index 9c908d2cd6..5c99b1e055 100644 --- a/tests/firedrake/regression/test_point_eval_api.py +++ b/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,6 +221,10 @@ 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) @@ -217,8 +232,14 @@ def test_point_evaluator_vector_tensor_mixed(mesh_and_points): 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,12 +307,17 @@ 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, 1) + assert isinstance(res[0], np.ndarray) + assert isinstance(res[0, 0], Number) # multiple points as tuples/list for input in [ @@ -299,29 +325,50 @@ def test_point_evaluator_inputs_1d(): (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),) + f.ufl_shape + assert isinstance(res[0], np.ndarray) + assert isinstance(res[0, 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),) + f.ufl_shape + assert isinstance(res2[0], np.ndarray) + assert isinstance(res2[0, 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),) + f.ufl_shape + assert isinstance(res[0], np.ndarray) + assert isinstance(res[0, 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 [ @@ -329,15 +376,33 @@ def test_point_evaluator_inputs_2d(): (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) diff --git a/tests/firedrake/regression/test_point_eval_fs.py b/tests/firedrake/regression/test_point_eval_fs.py index d0a6674f1b..f9aef6603e 100644 --- a/tests/firedrake/regression/test_point_eval_fs.py +++ b/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():