introduce PointEvaluator object

docs/source/point-evaluation.rst

@@ -214,6 +214,7 @@ be a new vertex: this is true for both ``redundant = True`` and
 and switch from ``redundant = True`` to ``redundant = False`` we will get point
 duplication.
 
+.. _missing_points:
 
 Points outside the domain
 ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -241,6 +242,51 @@ warning or switched off entirely:
    :start-after: [test_vom_manual_points_outside_domain 5]
    :end-before: [test_vom_manual_points_outside_domain 6]
 
+.. _point_evaluator:
+
+``PointEvaluator`` convenience object
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The :py:class:`~.PointEvaluator` class performs point evaluation using vertex-only meshes,
+as described above. This is a convenience object for users who want to evaluate
+functions at points without writing the vertex-only mesh boilerplate code each time.
+
+First, create a :py:class:`~.PointEvaluator` object by passing the 
+parent mesh and the points to evaluate at:
+
+.. code-block:: python3
+
+   point_evaluator = PointEvaluator(mesh, points)
+
+Internally, this creates a vertex-only mesh at the given points, immersed in the given mesh.
+To evaluate a :py:class:`~.Function` defined on the parent mesh at the given points,
+we use :meth:`~.PointEvaluator.evaluate`:
+
+.. code-block:: python3
+
+   f_at_points = point_evaluator.evaluate(f)
+
+Under the hood, this creates the appropriate P0DG function space on the vertex-only mesh
+and performs the interpolation. The points are then reordered to match the input ordering,
+as described in :ref:`the section on the input ordering property <input_ordering>`. The result
+is a Numpy array containing the values of ``f`` at the given points, in the order the points were
+provided to the :py:class:`~.PointEvaluator` constructor.
+
+If ``redundant=True`` (the default) was used when creating the :py:class:`~.PointEvaluator`,
+then only the points given to the constructor on rank 0 will be evaluated. The result is then
+broadcast to all ranks. Use this option if the same points are given on all ranks.
+If ``redundant=False`` was used when creating the :py:class:`~.PointEvaluator`, then
+each rank will evaluate the points it was given. Use this option if different points are given
+on different ranks, for example when using external point data.
+
+The parameters ``missing_points_behaviour`` and ``tolerance`` (discussed :ref:`here <missing_points>` 
+and :ref:`here <tolerance>` respectively) can be set when creating the :py:class:`~.PointEvaluator` 
+and will be passed to the :func:`~.VertexOnlyMesh` it creates internally.
+
+If the :ref:`coordinates <changing_coordinate_fs>` or the :ref:`tolerance <tolerance>` of the parent mesh
+are changed after creating the :py:class:`~.PointEvaluator`, then the vertex-only mesh
+will be reconstructed on the new mesh the next time :meth:`~.PointEvaluator.evaluate` is called.
+
 
 Expressions with point evaluations
 ----------------------------------

firedrake/function.py

@@ -11,6 +11,8 @@
 from collections.abc import Collection
 from numbers import Number
 from pathlib import Path
+from functools import partial
+from typing import Tuple
 
 from pyop2 import op2, mpi
 from pyop2.exceptions import DataTypeError, DataValueError
@@ -23,9 +25,10 @@
 from firedrake import utils
 from firedrake.adjoint_utils import FunctionMixin
 from firedrake.petsc import PETSc
+from firedrake.mesh import MeshGeometry, VertexOnlyMesh
+from firedrake.functionspace import FunctionSpace, VectorFunctionSpace, TensorFunctionSpace
 
-
-__all__ = ['Function', 'PointNotInDomainError', 'CoordinatelessFunction']
+__all__ = ['Function', 'PointNotInDomainError', 'CoordinatelessFunction', 'PointEvaluator']
 
 
 class _CFunction(ctypes.Structure):
@@ -698,6 +701,121 @@ def __str__(self):
         return "domain %s does not contain point %s" % (self.domain, self.point)
 
 
+class PointEvaluator:
+    r"""Convenience class for evaluating a :class:`Function` at a set of points."""
+
+    def __init__(self, mesh: MeshGeometry, points: np.ndarray | list, tolerance: float | None = None,
+                 missing_points_behaviour: str = "error", redundant: bool = True) -> None:
+        r"""
+        Parameters
+        ----------
+        mesh : MeshGeometry
+            The mesh on which to embed the points.
+        points : numpy.ndarray | list
+            Array or list of points to evaluate at.
+        tolerance : float | None
+            Tolerance to use when checking if a point is in a cell.
+            If ``None`` (the default), the ``tolerance`` of the ``mesh`` is used.
+        missing_points_behaviour : str
+            Behaviour when a point is not found in the mesh. Options are:
+            "error": raise a :class:`~.VertexOnlyMeshMissingPointsError` if a point is not found in the mesh.
+            "warn": warn if a point is not found in the mesh, but continue.
+            "ignore": ignore points not found in the mesh.
+        redundant : bool
+            If True, only the points given to the constructor on rank 0 are evaluated, and the result is broadcast to all ranks.
+            If False, each rank evaluates the points it has been given. False is useful if you are inputting
+            external data that is already distributed across ranks. Default is True.
+        """
+        self.points = np.asarray(points, dtype=utils.ScalarType)
+        if not self.points.shape:
+            self.points = self.points.reshape(-1)
+        gdim = mesh.geometric_dimension()
+        if self.points.shape[-1] != gdim and (len(self.points.shape) != 1 or gdim != 1):
+            raise ValueError(f"Point dimension ({self.points.shape[-1]}) does not match geometric dimension ({gdim}).")
+        self.points = self.points.reshape(-1, gdim)
+
+        self.mesh = mesh
+
+        self.redundant = redundant
+        self.missing_points_behaviour = missing_points_behaviour
+        self.tolerance = tolerance
+        self.vom = VertexOnlyMesh(
+            mesh, self.points, missing_points_behaviour=missing_points_behaviour,
+            redundant=redundant, tolerance=tolerance
+        )
+
+    def evaluate(self, function: Function) -> np.ndarray | Tuple[np.ndarray, ...]:
+        r"""Evaluate the given :class:`Function`.
+        Points that were not found in the mesh will be evaluated to np.nan.
+
+        Parameters
+        ----------
+        function :
+            The :class:`Function` to evaluate.
+
+        Returns
+        -------
+        numpy.ndarray | Tuple[numpy.ndarray, ...]
+            A Numpy array of values at the points. If the function is scalar-valued, the Numpy array
+            has shape ``(len(points),)``. If the function is vector-valued with shape ``(n,)``, the Numpy array has shape
+            ``(len(points), n)``. If the function is tensor-valued with shape ``(n, m)``, the Numpy array has shape
+            ``(len(points), n, m)``. If the function is a mixed function, a tuple of Numpy arrays is returned,
+            one for each subfunction.
+
+
+        .. warning::
+
+            This method returns a numpy array and hence isn't taped for use with firedrake-adjoint.
+            If you want to use point evaluation with the adjoint, create a :class:`~.VertexOnlyMesh`
+            as described in the manual.
+        """
+        from firedrake import assemble, interpolate
+        if not isinstance(function, Function):
+            raise TypeError(f"Expected a Function, got {type(function).__name__}")
+        if function.function_space().ufl_element().family() == "Real":
+            return function.dat.data_ro
+
+        function_mesh = function.function_space().mesh()
+        if function_mesh is not self.mesh:
+            raise ValueError("Function mesh must be the same Mesh object as the PointEvaluator mesh.")
+        if coord_changed := function_mesh.coordinates.dat.dat_version != self.mesh._saved_coordinate_dat_version:
+            # TODO: This is here until https://github.com/firedrakeproject/firedrake/issues/4540 is solved
+            self.mesh = function_mesh
+        if tol_changed := self.mesh.tolerance != self.tolerance:
+            self.tolerance = self.mesh.tolerance
+        if coord_changed or tol_changed:
+            self.vom = VertexOnlyMesh(
+                self.mesh, self.points, missing_points_behaviour=self.missing_points_behaviour,
+                redundant=self.redundant, tolerance=self.tolerance
+            )
+
+        subfunctions = function.subfunctions
+        if len(subfunctions) > 1:
+            return tuple(self.evaluate(subfunction) for subfunction in subfunctions)
+
+        shape = function.ufl_function_space().value_shape
+        if len(shape) == 0:
+            fs = FunctionSpace
+        elif len(shape) == 1:
+            fs = partial(VectorFunctionSpace, dim=shape[0])
+        else:
+            fs = partial(TensorFunctionSpace, shape=shape)
+        P0DG = fs(self.vom, "DG", 0)
+        P0DG_io = fs(self.vom.input_ordering, "DG", 0)
+
+        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
+
+        # 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
+
+
 @PETSc.Log.EventDecorator()
 def make_c_evaluate(function, c_name="evaluate", ldargs=None, tolerance=None):
     r"""Generates, compiles and loads a C function to evaluate the

tests/firedrake/regression/test_point_eval_api.py

@@ -3,6 +3,7 @@
 import pytest
 
 from firedrake import *
+from firedrake.mesh import VertexOnlyMeshMissingPointsError
 
 cwd = abspath(dirname(__file__))
 
@@ -161,3 +162,182 @@ def test_tolerance():
     with pytest.raises(PointNotInDomainError):
         f.at((-0.1, 0.4))
     assert np.allclose([-1e-12, 0.4], f.at((-1e-12, 0.4)))
+
+
+@pytest.fixture(scope="module")
+def mesh_and_points():
+    points = [[0.1, 0.1], [0.2, 0.2], [0.3, 0.3]]
+    mesh = UnitSquareMesh(10, 10)
+    evaluator = PointEvaluator(mesh, points)
+    return mesh, evaluator
+
+
+@pytest.mark.parallel([1, 3])
+def test_point_evaluator_scalar(mesh_and_points):
+    mesh, evaluator = mesh_and_points
+    points = evaluator.points
+    V = FunctionSpace(mesh, "CG", 1)
+    f = Function(V)
+    x, y = SpatialCoordinate(mesh)
+    f.interpolate(x + y)
+
+    # 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])
+
+    # 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])
+
+
+@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_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)
+
+    V_tensor = TensorFunctionSpace(mesh, "CG", 1, shape=(2, 3))
+    f_tensor = Function(V_tensor)
+    f_tensor.interpolate(as_matrix([[x, y, x*y], [y, x, x*y]]))
+    f_tensor_at_points = evaluator.evaluate(f_tensor)
+    tensor_expected = np.array([[[0.1, 0.1, 0.01], [0.1, 0.1, 0.01]],
+                                [[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)
+
+    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 np.allclose(f_mixed_at_points[0], vec_expected)
+    assert np.allclose(f_mixed_at_points[1], tensor_expected)
+
+
+@pytest.mark.parallel(3)
+def test_point_evaluator_nonredundant(mesh_and_points):
+    mesh = mesh_and_points[0]
+    if mesh.comm.rank == 0:
+        points = [[0.1, 0.1]]
+    elif mesh.comm.rank == 1:
+        points = [[0.4, 0.4], [0.5, 0.5]]
+    else:
+        points = [[0.8, 0.8]]
+    V = FunctionSpace(mesh, "CG", 1)
+    f = Function(V)
+    x, y = SpatialCoordinate(mesh)
+    f.interpolate(x + y)
+    evaluator = PointEvaluator(mesh, points, redundant=False)
+    f_at_points = evaluator.evaluate(f)
+    if mesh.comm.rank == 0:
+        assert np.allclose(f_at_points, [0.2])
+    elif mesh.comm.rank == 1:
+        assert np.allclose(f_at_points, [0.8, 1.0])
+    else:
+        assert np.allclose(f_at_points, [1.6])
+
+
+def test_point_evaluator_moving_mesh(mesh_and_points):
+    mesh, evaluator = mesh_and_points
+    V = FunctionSpace(mesh, "CG", 1)
+    f = Function(V)
+    x, y = SpatialCoordinate(mesh)
+    f.interpolate(x + y)
+
+    mesh.coordinates.dat.data[:, 0] += 1.0
+
+    with pytest.raises(VertexOnlyMeshMissingPointsError):
+        # The VOM is reimmersed, but the points
+        # are now outside of the mesh.
+        f_at_points = evaluator.evaluate(f)
+
+    mesh.coordinates.dat.data[:, 0] -= 1.0
+    f_at_points = evaluator.evaluate(f)
+    assert np.allclose(f_at_points, [0.2, 0.4, 0.6])
+
+
+def test_point_evaluator_tolerance():
+    mesh = UnitSquareMesh(1, 1)
+    old_tol = mesh.tolerance
+    f = Function(VectorFunctionSpace(mesh, "CG", 1)).interpolate(SpatialCoordinate(mesh))
+    ev = PointEvaluator(mesh, [[0.2, 0.4]])
+    assert np.allclose([0.2, 0.4], ev.evaluate(f))
+    # tolerance of mesh is not changed
+    assert mesh.tolerance == old_tol
+    # Outside mesh, but within tolerance
+    ev = PointEvaluator(mesh, [[-0.1, 0.4]], tolerance=0.2)
+    assert np.allclose([-0.1, 0.4], ev.evaluate(f))
+    # tolerance of mesh is changed
+    assert mesh.tolerance == 0.2
+    # works if mesh tolerance is changed
+    mesh.tolerance = 1e-11
+    with pytest.raises(VertexOnlyMeshMissingPointsError):
+        ev.evaluate(f)
+    ev = PointEvaluator(mesh, [[-1e-12, 0.4]])
+    assert np.allclose([-1e-12, 0.4], ev.evaluate(f))
+
+
+def test_point_evaluator_inputs_1d():
+    mesh = UnitIntervalMesh(1)
+    f = mesh.coordinates
+
+    # 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))
+
+    # 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))
+        e3 = PointEvaluator(mesh, list(input))
+        assert np.allclose([[0.2, 0.3]], e3.evaluate(f))
+
+    # 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))
+
+    # 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):
+            PointEvaluator(mesh, input)
+
+
+def test_point_evaluator_inputs_2d():
+    mesh = UnitSquareMesh(1, 1)
+    f = mesh.coordinates
+
+    # 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))
+
+    # 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))
+        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))
+
+    # test incorrect inputs
+    for input in [0.2, [0.2]]:
+        with pytest.raises(ValueError):
+            PointEvaluator(mesh, input)