diff --git a/firedrake/function.py b/firedrake/function.py index 9d8a219fb7..886a4db5ac 100644 --- a/firedrake/function.py +++ b/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], diff --git a/firedrake/mesh.py b/firedrake/mesh.py index b6aa2e3f15..57dc5771c5 100644 --- a/firedrake/mesh.py +++ b/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() diff --git a/firedrake/utils.py b/firedrake/utils.py index fa58f1f7a0..068e557bfc 100644 --- a/firedrake/utils.py +++ b/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 diff --git a/tests/firedrake/regression/test_point_eval_fs.py b/tests/firedrake/regression/test_point_eval_fs.py index 2a4a2c718a..e2ba9553ce 100644 --- a/tests/firedrake/regression/test_point_eval_fs.py +++ b/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