Cache mesh bounding boxes correctly

firedrake/function.py

@@ -20,11 +20,10 @@
 from pyop2.exceptions import DataTypeError, DataValueError
 
 from finat.ufl import MixedElement
-from firedrake.utils import ScalarType, IntType, as_ctypes
+from firedrake.utils import ScalarType, IntType, as_ctypes, cached_property_until, _new_uid, complex_mode, ScalarType_c
 
 from firedrake import functionspaceimpl
 from firedrake.cofunction import Cofunction, RieszMap
-from firedrake import utils
 from firedrake.adjoint_utils import FunctionMixin
 from firedrake.petsc import PETSc
 from firedrake.mesh import MeshGeometry, VertexOnlyMesh
@@ -75,7 +74,7 @@ def __init__(self, function_space, val=None, name=None, dtype=ScalarType):
         # User comm
         self.comm = function_space.comm
         self._function_space = function_space
-        self.uid = utils._new_uid(self.comm)
+        self.uid = _new_uid(self.comm)
         self._name = name or 'function_%d' % self.uid
         self._label = "a function"
 
@@ -564,7 +563,7 @@ 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)
+        coord = np.asarray(coord, dtype=ScalarType)
         evaluator = PointEvaluator(self.function_space().mesh(), coord)
         result = evaluator.evaluate(self)
         if len(coord.shape) == 1:
@@ -602,8 +601,8 @@ def _at(self, arg, *args, **kwargs):
 
         if args:
             arg = (arg,) + args
-        arg = np.asarray(arg, dtype=utils.ScalarType)
-        if utils.complex_mode:
+        arg = np.asarray(arg, dtype=ScalarType)
+        if complex_mode:
             if not np.allclose(arg.imag, 0):
                 raise ValueError("Provided points have non-zero imaginary part")
             arg = arg.real.copy()
@@ -734,7 +733,7 @@ def __init__(self, mesh: MeshGeometry, points: np.ndarray | list, tolerance: flo
             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)
+        self.points = np.asarray(points, dtype=ScalarType)
         if not self.points.shape:
             self.points = self.points.reshape(-1)
         gdim = mesh.geometric_dimension
@@ -743,13 +742,20 @@ def __init__(self, mesh: MeshGeometry, points: np.ndarray | list, tolerance: flo
         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
+        if tolerance is not None:
+            mesh.tolerance = tolerance
+
+    @cached_property_until(
+        lambda self: (self.mesh.coordinates.dat.dat_version, self.mesh.tolerance)
+    )
+    def vom(self) -> MeshGeometry:
+        """The VOM used for point evaluation. This is cached until the mesh coordinates or tolerance change."""
+        # We invalidate when the parent mesh moves until https://github.com/firedrakeproject/firedrake/issues/4540 is fixed.
+        return VertexOnlyMesh(
+            self.mesh, self.points, missing_points_behaviour=self.missing_points_behaviour,
+            redundant=self.redundant, tolerance=None
         )
 
     def evaluate(self, function: Function) -> np.ndarray | Tuple[np.ndarray, ...]:
@@ -787,19 +793,8 @@ def evaluate(self, function: Function) -> np.ndarray | Tuple[np.ndarray, ...]:
         if function.function_space().ufl_element().family() == "Real":
             return function.dat.data_ro
 
-        function_mesh = function.function_space().mesh().unique()
-        if function_mesh is not self.mesh:
+        if function.function_space().mesh().unique() 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:
@@ -823,7 +818,7 @@ def evaluate(self, function: Function) -> np.ndarray | Tuple[np.ndarray, ...]:
         # 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)
+                result = np.empty((len(self.points),) + shape, dtype=ScalarType)
             self.mesh.comm.Bcast(result)
         return result
 
@@ -851,7 +846,7 @@ def make_c_evaluate(function, c_name="evaluate", ldargs=None, tolerance=None):
     arg = function.dat(op2.READ, function.cell_node_map())
     args.append(arg)
 
-    p_ScalarType_c = f"{utils.ScalarType_c}*"
+    p_ScalarType_c = f"{ScalarType_c}*"
     src.append(generate_single_cell_wrapper(mesh.cell_set, args,
                                             forward_args=[p_ScalarType_c,
                                                           p_ScalarType_c],

firedrake/mesh.py

@@ -37,7 +37,7 @@
 import firedrake.extrusion_utils as eutils
 import firedrake.cython.rtree as rtree
 import firedrake.utils as utils
-from firedrake.utils import as_cstr, IntType, RealType
+from firedrake.utils import as_cstr, IntType, RealType, cached_property_until
 from firedrake.logging import logger
 from firedrake.parameters import parameters
 from firedrake.petsc import PETSc, DEFAULT_PARTITIONER
@@ -2380,10 +2380,6 @@ def __init__(self, coordinates):
         # submesh
         self.submesh_parent = None
 
-        self._bounding_box_coords = None
-        self._rtree = None
-        self._saved_coordinate_dat_version = coordinates.dat.dat_version
-
         # Cache mesh object on the coordinateless coordinates function
         coordinates._as_mesh_geometry = weakref.ref(self)
 
@@ -2477,9 +2473,11 @@ def clear_rtree(self):
 
         Use this if you move the mesh (for example by reassigning to
         the coordinate field)."""
-        self._rtree = None
+        # `cached_property_until` stores the cached rtree in self._rtree_cache
+        # setting it to None will force the rtree to be rebuilt on next access.
+        self._rtree_cache = None
 
-    @cached_property
+    @cached_property_until(lambda self: self.coordinates.dat.dat_version)
     @PETSc.Log.EventDecorator()
     def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray]:
         """Calculates bounding boxes for the mesh rtree.
@@ -2558,10 +2556,9 @@ def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray]:
         column_list = V.cell_node_list.reshape(-1)
         coords_min = mesh._order_data_by_cell_index(column_list, coords_min.dat.data_ro_with_halos)
         coords_max = mesh._order_data_by_cell_index(column_list, coords_max.dat.data_ro_with_halos)
-
         return coords_min, coords_max
 
-    @property
+    @cached_property_until(lambda self: (self.coordinates.dat.dat_version, self.tolerance))
     @PETSc.Log.EventDecorator()
     def rtree(self):
         """Builds an rtree from bounding box coordinates, expanding
@@ -2579,12 +2576,6 @@ def rtree(self):
         can be found.
 
         """
-        if self.coordinates.dat.dat_version != self._saved_coordinate_dat_version:
-            if "bounding_box_coords" in self.__dict__:
-                del self.bounding_box_coords
-        else:
-            if self._rtree:
-                return self._rtree
         # Change min and max to refer to an n-hypercube, where n is the
         # geometric dimension of the mesh, centred on the midpoint of the
         # bounding box. Its side length is the L1 diameter of the bounding box.
@@ -2609,7 +2600,6 @@ def rtree(self):
 
         with PETSc.Log.Event("rtree_build"):
             self._rtree = rtree.build_from_aabb(coords_min, coords_max)
-        self._saved_coordinate_dat_version = self.coordinates.dat.dat_version
         return self._rtree
 
     @PETSc.Log.EventDecorator()

firedrake/utils.py

@@ -1,6 +1,8 @@
 # Some generic python utilities not really specific to our work.
 import collections.abc
 import warnings
+import functools
+from typing import Callable, Self, Hashable
 from decorator import decorator
 from pyop2.datatypes import ScalarType, as_cstr
 from pyop2.datatypes import RealType     # noqa: F401
@@ -264,3 +266,38 @@ def check_netgen_installed() -> None:
             "are installed and available to Firedrake (see "
             "https://www.firedrakeproject.org/install.html#netgen)."
         )
+
+
+def cached_property_until(key: Callable[[Self], Hashable]):
+    """Decorator for a property that is cached until some value changes.
+
+    For example, the ``expensive_property`` below will be cached until
+    ``self.value`` changes, and will be recomputed with the new ``self.value``
+    and cached when accessed again.
+
+    .. code-block:: python
+
+        class MyClass:
+
+            def __init__(self, value):
+                self.value = value
+
+            @cached_property_until(lambda self: self.value)
+            def expensive_property(self):
+                # Some expensive computation that depends on self.value
+                ...
+    """
+    def decorator(func):
+        @property
+        @functools.wraps(func)
+        def wrapper(self):
+            cache_attribute = f"_{func.__name__}_cache"
+            current_value = key(self)
+            cached_value = getattr(self, cache_attribute, None)
+            if cached_value is None or cached_value[0] != current_value:
+                result = func(self)
+                setattr(self, cache_attribute, (current_value, result))
+                return result
+            return cached_value[1]
+        return wrapper
+    return decorator

tests/firedrake/regression/test_point_eval_fs.py

@@ -239,7 +239,23 @@ def test_point_reset_works():
 
 def test_changing_coordinates_invalidates_rtree():
     mesh = UnitSquareMesh(2, 2)
-
+    assert mesh.rtree is mesh.rtree
     saved_rtree = mesh.rtree
     mesh.coordinates.assign(mesh.coordinates * 2)
     assert mesh.rtree != saved_rtree
+
+
+def test_changing_coordinates_invalidates_bounding_box():
+    mesh = UnitSquareMesh(2, 2)
+    assert mesh.bounding_box_coords is mesh.bounding_box_coords
+    saved_bounding_box_coords = mesh.bounding_box_coords
+    mesh.coordinates.assign(mesh.coordinates * 2)
+    assert not np.allclose(mesh.bounding_box_coords, saved_bounding_box_coords)
+
+
+def test_changing_tolerance_invalidates_rtree():
+    mesh = UnitSquareMesh(2, 2)
+    assert mesh.rtree is mesh.rtree
+    saved_rtree = mesh.rtree
+    mesh.tolerance = 1e-5
+    assert mesh.rtree != saved_rtree