diff --git a/docs/source/point-evaluation.rst b/docs/source/point-evaluation.rst index 731a0081f5..03e612cf67 100644 --- a/docs/source/point-evaluation.rst +++ b/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 `. 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 ` +and :ref:`here ` 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 ` or the :ref:`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 ---------------------------------- diff --git a/firedrake/function.py b/firedrake/function.py index 4cce6b9edb..a508dddc8b 100644 --- a/firedrake/function.py +++ b/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 diff --git a/tests/firedrake/regression/test_point_eval_api.py b/tests/firedrake/regression/test_point_eval_api.py index e953a423db..22d177cb04 100644 --- a/tests/firedrake/regression/test_point_eval_api.py +++ b/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)