Cell location on simplicial higher-order meshes (#4484)

* cache bounding box calculation. Recalculate if mesh coords are changed

Author: leo-collins

firedrake/interpolation.py

@@ -465,6 +465,13 @@ def __init__(
         V_dest = V.function_space() if isinstance(V, firedrake.Function) else V
         src_mesh = extract_unique_domain(expr)
         dest_mesh = as_domain(V_dest)
+        if (
+            ufl.cell.simplex(src_mesh.topological_dimension()) != src_mesh.ufl_cell()
+            and numpy.any(numpy.asarray(src_mesh.ufl_coordinate_element().degree()) > 1)
+        ):
+            raise NotImplementedError(
+                "Cannot yet interpolate from non-simplicial higher-order meshes into other meshes."
+            )
         src_mesh_gdim = src_mesh.geometric_dimension()
         dest_mesh_gdim = dest_mesh.geometric_dimension()
         if src_mesh_gdim != dest_mesh_gdim:
@@ -473,15 +480,6 @@ def __init__(
             )
         self.src_mesh = src_mesh
         self.dest_mesh = dest_mesh
-        if numpy.any(
-            numpy.asarray(src_mesh.coordinates.function_space().ufl_element().degree())
-            > 1
-        ):
-            # Need to implement vertex-only mesh immersion in high order meshes
-            # for this to work.
-            raise NotImplementedError(
-                "Cannot yet interpolate from high order meshes to other meshes."
-            )
 
         self.sub_interpolators = []
 

firedrake/mesh.py

@@ -6,6 +6,7 @@
 import finat.ufl
 import FIAT
 import weakref
+from typing import Tuple
 from collections import OrderedDict, defaultdict
 from collections.abc import Sequence
 from ufl.classes import ReferenceGrad
@@ -2289,6 +2290,7 @@ def __init__(self, coordinates):
         # submesh
         self.submesh_parent = None
 
+        self._bounding_box_coords = None
         self._spatial_index = None
         self._saved_coordinate_dat_version = coordinates.dat.dat_version
 
@@ -2457,54 +2459,71 @@ def clear_spatial_index(self):
         the coordinate field)."""
         self._spatial_index = None
 
-    @property
-    def spatial_index(self):
-        """Spatial index to quickly find which cell contains a given point.
+    @utils.cached_property
+    def bounding_box_coords(self) -> Tuple[np.ndarray, np.ndarray] | None:
+        """Calculates bounding boxes for spatial indexing.
 
-        Notes
-        -----
+        Returns
+        -------
+        Tuple of arrays of shape (num_cells, gdim) containing
+        the minimum and maximum coordinates of each cell's bounding box.
 
-        If this mesh has a :attr:`tolerance` property, which
-        should be a float, this tolerance is added to the extrama of the
-        spatial index so that points just outside the mesh, within tolerance,
-        can be found.
+        None if the geometric dimension is 1, since libspatialindex
+        does not support 1D.
 
+        Notes
+        -----
+        If we have a higher-order (bendy) mesh we project the mesh coordinates into
+        a Bernstein finite element space. Functions on a Bernstein element are
+        Bezier curves and are completely contained in the convex hull of the mesh nodes.
+        Hence the bounding box will contain the entire element.
         """
         from firedrake import function, functionspace
         from firedrake.parloops import par_loop, READ, MIN, MAX
 
-        if (
-            self._spatial_index
-            and self.coordinates.dat.dat_version == self._saved_coordinate_dat_version
-        ):
-            return self._spatial_index
-
         gdim = self.geometric_dimension()
         if gdim <= 1:
             info_red("libspatialindex does not support 1-dimension, falling back on brute force.")
             return None
 
+        coord_element = self.ufl_coordinate_element()
+        coord_degree = coord_element.degree()
+        if ufl.cell.simplex(self.topological_dimension()) != self.ufl_cell():
+            # Non-simplex element, e.g. quad or tensor product
+            mesh = self
+        elif coord_degree == 1:
+            mesh = self
+        elif coord_element.family() == "Bernstein":
+            # Already have Bernstein coordinates, no need to project
+            mesh = self
+        else:
+            # For bendy meshes we project the coordinate function onto Bernstein
+            bernstein_fs = functionspace.VectorFunctionSpace(self, "Bernstein", coord_degree)
+            f = function.Function(bernstein_fs)
+            f.interpolate(self.coordinates)
+            mesh = Mesh(f)
+
         # Calculate the bounding boxes for all cells by running a kernel
-        V = functionspace.VectorFunctionSpace(self, "DG", 0, dim=gdim)
+        V = functionspace.VectorFunctionSpace(mesh, "DG", 0, dim=gdim)
         coords_min = function.Function(V, dtype=RealType)
         coords_max = function.Function(V, dtype=RealType)
 
         coords_min.dat.data.fill(np.inf)
         coords_max.dat.data.fill(-np.inf)
 
         if utils.complex_mode:
-            if not np.allclose(self.coordinates.dat.data_ro.imag, 0):
+            if not np.allclose(mesh.coordinates.dat.data_ro.imag, 0):
                 raise ValueError("Coordinate field has non-zero imaginary part")
-            coords = function.Function(self.coordinates.function_space(),
-                                       val=self.coordinates.dat.data_ro_with_halos.real.copy(),
+            coords = function.Function(mesh.coordinates.function_space(),
+                                       val=mesh.coordinates.dat.data_ro_with_halos.real.copy(),
                                        dtype=RealType)
         else:
-            coords = self.coordinates
+            coords = mesh.coordinates
 
-        cell_node_list = self.coordinates.function_space().cell_node_list
+        cell_node_list = mesh.coordinates.function_space().cell_node_list
         _, nodes_per_cell = cell_node_list.shape
 
-        domain = "{{[d, i]: 0 <= d < {0} and 0 <= i < {1}}}".format(gdim, nodes_per_cell)
+        domain = f"{{[d, i]: 0 <= d < {gdim} and 0 <= i < {nodes_per_cell}}}"
         instructions = """
         for d, i
             f_min[0, d] = fmin(f_min[0, d], f[i, d])
@@ -2518,21 +2537,51 @@ def spatial_index(self):
 
         # Reorder bounding boxes according to the cell indices we use
         column_list = V.cell_node_list.reshape(-1)
-        coords_min = self._order_data_by_cell_index(column_list, coords_min.dat.data_ro_with_halos)
-        coords_max = self._order_data_by_cell_index(column_list, coords_max.dat.data_ro_with_halos)
+        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
+    def spatial_index(self):
+        """Builds spatial index from bounding box coordinates, expanding
+        the bounding box by the mesh tolerance.
+
+        Returns
+        -------
+        :class:`~.spatialindex.SpatialIndex` or None if the mesh is
+        one-dimensional.
+
+        Notes
+        -----
+        If this mesh has a :attr:`tolerance` property, which
+        should be a float, this tolerance is added to the extrema of the
+        spatial index so that points just outside the mesh, within tolerance,
+        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._spatial_index:
+                return self._spatial_index
         # 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.
         # This aids point evaluation on immersed manifolds and other cases
         # where points may be just off the mesh but should be evaluated.
         # TODO: This is perhaps unnecessary when we aren't in these special
         # cases.
-
         # We also push max and min out so we can find points on the boundary
         # within the mesh tolerance.
         # NOTE: getattr doesn't work here due to the inheritance games that are
         # going on in getattr.
+        if self.bounding_box_coords is None:
+            # This happens in 1D meshes
+            return None
+        else:
+            coords_min, coords_max = self.bounding_box_coords
         tolerance = self.tolerance if hasattr(self, "tolerance") else 0.0
         coords_mid = (coords_max + coords_min)/2
         d = np.max(coords_max - coords_min, axis=1)[:, None]
@@ -3360,11 +3409,13 @@ def VertexOnlyMesh(mesh, vertexcoords, reorder=None, missing_points_behaviour='e
     _, pdim = vertexcoords.shape
     if not np.isclose(np.sum(abs(vertexcoords.imag)), 0):
         raise ValueError("Point coordinates must have zero imaginary part")
-    # Bendy meshes require a smarter bounding box algorithm at partition and
-    # (especially) cell level. Projecting coordinates to Bernstein may be
-    # sufficient.
-    if np.any(np.asarray(mesh.coordinates.function_space().ufl_element().degree()) > 1):
-        raise NotImplementedError("Only straight edged meshes are supported")
+    if (
+        ufl.cell.simplex(mesh.topological_dimension()) != mesh.ufl_cell()
+        and np.any(np.asarray(mesh.ufl_coordinate_element().degree()) > 1)
+    ):
+        raise NotImplementedError(
+            "Cannot yet immerse a VertexOnlyMesh in non-simplicial higher-order meshes."
+        )
     # Currently we take responsibility for locating the mesh cells in which the
     # vertices lie.
     #

tests/firedrake/regression/test_interpolate_cross_mesh.py

@@ -42,14 +42,7 @@ def make_high_order(m_low_order, degree):
         "unitsquare",
         "circlemanifold",
         "circlemanifold_to_high_order",
-        pytest.param(
-            "unitsquare_from_high_order",
-            marks=pytest.mark.xfail(
-                # CalledProcessError is so the parallel tests correctly xfail
-                raises=(subprocess.CalledProcessError, NotImplementedError),
-                reason="Cannot yet interpolate from high order meshes to other meshes.",
-            ),
-        ),
+        "unitsquare_from_high_order",
         "unitsquare_to_high_order",
         "extrudedcube",
         "unitsquare_vfs",

tests/firedrake/regression/test_locate_cell.py

@@ -1,6 +1,12 @@
 import pytest
 import numpy as np
 from firedrake import *
+from functools import reduce
+from operator import mul
+
+
+def warp(x, p):
+    return p * x * (2 * x - 1)
 
 
 @pytest.fixture(scope="module", params=[False, True])
@@ -76,3 +82,67 @@ def test_locate_cells_ref_coords_and_dists(meshdata):
     fcells, ref_coords, l1_dists = m.locate_cells_ref_coords_and_dists(points[:2], cells_ignore=np.array([cells[:4], cells[1:5]]))
     assert fcells[0] == -1 or fcells[0] in cells[4:]
     assert fcells[1] == -1 or fcells[1] in cells[5:] or fcells[1] in cells[:1]
+
+
+def test_high_order_location():
+    mesh = UnitSquareMesh(2, 2)
+    V = VectorFunctionSpace(mesh, "CG", 3, variant="equispaced")
+    f = Function(V)
+    f.interpolate(mesh.coordinates)
+
+    warp_indices = np.where((f.dat.data[:, 0] > 0.0) & (f.dat.data[:, 0] < 0.5) & (f.dat.data[:, 1] == 0.0))[0]
+    f.dat.data[warp_indices, 1] = warp(f.dat.data[warp_indices, 0], 5.0)
+    mesh = Mesh(f)
+
+    # The point (0.25, -0.6) *is* in the mesh, but falls outside the Lagrange bounding box
+    # The below used to return (None, None), but projecting to Bernstein coordinates
+    # allows us to locate the cell.
+    assert mesh.locate_cell([0.25, -0.6], tolerance=0.001) is not None
+    # The point (0.25, -0.7) is outside the mesh, but inside the Bernstein bounding box.
+    # This should return (None, None).
+    assert mesh.locate_cell([0.25, -0.7], tolerance=0.001) is None
+
+    # Change mesh coordinates to check that the bounding box is recalculated
+    mesh.coordinates.dat.data_wo[warp_indices, 1] = warp(mesh.coordinates.dat.data_ro[warp_indices, 0], 8.0)
+    assert mesh.locate_cell([0.25, -0.6], tolerance=0.0001) is not None
+    assert mesh.locate_cell([0.25, -0.7], tolerance=0.0001) is not None
+    assert mesh.locate_cell([0.25, -0.95], tolerance=0.0001) is not None
+    assert mesh.locate_cell([0.25, -1.05], tolerance=0.0001) is None
+
+
+def test_high_order_location_warped_interior_facet():
+    # Here we bend an interior facet and check the right cell is located.
+    mesh = UnitSquareMesh(2, 2)
+    V = VectorFunctionSpace(mesh, "CG", 3, variant="equispaced")
+    f = Function(V)
+    f.interpolate(mesh.coordinates)
+
+    warp_indices = np.where((f.dat.data[:, 0] > 0.0) & (f.dat.data[:, 0] < 0.5) & np.isclose(f.dat.data[:, 1], 0.5))[0]
+    f.dat.data[warp_indices, 1] += 0.1
+    mesh = Mesh(f)
+
+    assert mesh.locate_cell([0.25, 0.605], tolerance=0.0001) == 1
+    assert mesh.locate_cell([0.25, 0.62], tolerance=0.0001) == 3
+
+
+@pytest.mark.parallel([1, 3])
+def test_parallel_high_order_location():
+    mesh = UnitSquareMesh(2, 2)
+    V = VectorFunctionSpace(mesh, "CG", 3, variant="equispaced")
+    f = Function(V)
+    f.interpolate(mesh.coordinates)
+
+    warp_indices = np.where((f.dat.data[:, 0] > 0.0) & (f.dat.data[:, 0] < 0.5) & (f.dat.data[:, 1] == 0.0))[0]
+    f.dat.data[warp_indices, 1] = warp(f.dat.data[warp_indices, 0], 5.0)
+
+    mesh = Mesh(f)
+    V = FunctionSpace(mesh, "CG", 3)
+    f = Function(V).interpolate(reduce(mul, SpatialCoordinate(mesh)))
+
+    vom = VertexOnlyMesh(mesh, [[0.25, -0.6]], tolerance=0.0001, redundant=False)
+    P0DG = FunctionSpace(vom, "DG", 0)
+    P0DG_io = FunctionSpace(vom.input_ordering, "DG", 0)
+    f_at = assemble(interpolate(f, P0DG))
+    f_at_correct_order = assemble(interpolate(f_at, P0DG_io))
+
+    assert np.allclose(f_at_correct_order.dat.data_ro, [-0.6 * 0.25], atol=0.002)