Add functionalities to cache data when interpolating on nonmatching meshes

cpp/demo/interpolation_different_meshes/main.cpp

@@ -64,7 +64,10 @@ int main(int argc, char* argv[])
     u_tet->interpolate(fun);
 
     // Interpolate from u_tet to u_hex
-    u_hex->interpolate(*u_tet);
+    auto nmm_interpolation_data
+        = fem::create_nonmatching_meshes_interpolation_data(
+            *u_hex->function_space(), *u_tet->function_space());
+    u_hex->interpolate(*u_tet, nmm_interpolation_data);
 
 #ifdef HAS_ADIOS2
     io::VTXWriter write_tet(mesh_tet->comm(), "u_tet.vtx", {u_tet});

cpp/dolfinx/fem/Function.h

@@ -1,4 +1,4 @@
-// Copyright (C) 2003-2022 Anders Logg and Garth N. Wells
+// Copyright (C) 2003-2022 Anders Logg, Garth N. Wells and Massimiliano Leoni
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -145,14 +145,24 @@ class Function
   /// Interpolate a Function
   /// @param[in] v The function to be interpolated
   /// @param[in] cells The cells to interpolate on
-  void interpolate(const Function<T>& v, std::span<const std::int32_t> cells)
+  /// @param[in] nmm_interpolation_data Auxiliary data to interpolate on
+  /// nonmatching meshes. This data can be generated with
+  /// generate_nonmatching_meshes_interpolation_data (optional).
+  void interpolate(const Function<T>& v, std::span<const std::int32_t> cells,
+                   const nmm_interpolation_data_t& nmm_interpolation_data
+                   = nmm_interpolation_data_t{})
   {
-    fem::interpolate(*this, v, cells);
+    fem::interpolate(*this, v, cells, nmm_interpolation_data);
   }
 
   /// Interpolate a Function
   /// @param[in] v The function to be interpolated
-  void interpolate(const Function<T>& v)
+  /// @param[in] nmm_interpolation_data Auxiliary data to interpolate on
+  /// nonmatching meshes. This data can be generated with
+  /// generate_nonmatching_meshes_interpolation_data (optional).
+  void interpolate(const Function<T>& v,
+                   const nmm_interpolation_data_t& nmm_interpolation_data
+                   = nmm_interpolation_data_t{})
   {
     assert(_function_space);
     assert(_function_space->mesh());
@@ -163,7 +173,7 @@ class Function
     std::vector<std::int32_t> cells(num_cells, 0);
     std::iota(cells.begin(), cells.end(), 0);
 
-    fem::interpolate(*this, v, cells);
+    interpolate(v, cells, nmm_interpolation_data);
   }
 
   /// Interpolate an expression function on a list of cells

cpp/dolfinx/fem/interpolate.cpp

@@ -1,4 +1,5 @@
-// Copyright (C) 2021 Garth N. Wells, Jørgen S. Dokken, Igor A. Baratta
+// Copyright (C) 2021 Garth N. Wells, Jørgen S. Dokken, Igor A. Baratta,
+// Massimiliano Leoni
 //
 // This file is part of DOLFINx (https://www.fenicsproject.org)
 //
@@ -68,4 +69,49 @@ fem::interpolation_coords(const FiniteElement& element, const mesh::Mesh& mesh,
 
   return x;
 }
+
+fem::nmm_interpolation_data_t fem::create_nonmatching_meshes_interpolation_data(
+    const fem::FunctionSpace& Vu, const fem::FunctionSpace& Vv,
+    std::span<const std::int32_t> cells)
+{
+  std::vector<double> x;
+  auto mesh = Vu.mesh();
+  auto mesh_v = Vv.mesh();
+  auto element_u = Vu.element();
+
+  // Collect all the points at which values are needed to define the
+  // interpolating function
+  const std::vector<double> coords_b
+      = fem::interpolation_coords(*element_u, *mesh, cells);
+
+  namespace stdex = std::experimental;
+  using cmdspan2_t
+      = stdex::mdspan<const double, stdex::dextents<std::size_t, 2>>;
+  using mdspan2_t = stdex::mdspan<double, stdex::dextents<std::size_t, 2>>;
+  cmdspan2_t coords(coords_b.data(), 3, coords_b.size() / 3);
+  // Transpose interpolation coords
+  x.resize(coords.size());
+  mdspan2_t _x(x.data(), coords_b.size() / 3, 3);
+  for (std::size_t j = 0; j < coords.extent(1); ++j)
+    for (std::size_t i = 0; i < 3; ++i)
+      _x(j, i) = coords(i, j);
+
+  // Determine ownership of each point
+  return dolfinx::geometry::determine_point_ownership(*mesh_v, x);
+}
+
+fem::nmm_interpolation_data_t fem::create_nonmatching_meshes_interpolation_data(
+    const dolfinx::fem::FunctionSpace& Vu,
+    const dolfinx::fem::FunctionSpace& Vv)
+{
+  assert(Vu.mesh());
+  int tdim = Vu.mesh()->topology().dim();
+  auto cell_map = Vu.mesh()->topology().index_map(tdim);
+  assert(cell_map);
+  std::int32_t num_cells = cell_map->size_local() + cell_map->num_ghosts();
+  std::vector<std::int32_t> cells(num_cells, 0);
+  std::iota(cells.begin(), cells.end(), 0);
+
+  return create_nonmatching_meshes_interpolation_data(Vu, Vv, cells);
+}
 //-----------------------------------------------------------------------------

cpp/dolfinx/fem/interpolate.h

@@ -38,6 +38,18 @@ class Function;
 std::vector<double> interpolation_coords(const fem::FiniteElement& element,
                                          const mesh::Mesh& mesh,
                                          std::span<const std::int32_t> cells);
+
+/// Helper type for the data that can be cached to speed up repeated
+/// interpolation of discrete functions on nonmatching meshes
+using nmm_interpolation_data_t = decltype(std::function{
+    dolfinx::geometry::determine_point_ownership})::result_type;
+
+/// Forward declaration
+template <typename T>
+void interpolate(Function<T>& u, std::span<const T> f,
+                 std::array<std::size_t, 2> fshape,
+                 std::span<const std::int32_t> cells);
+
 namespace impl
 {
 /// @brief Scatter data into non-contiguous memory
@@ -541,6 +553,75 @@ void interpolate_nonmatching_maps(Function<T>& u1, const Function<T>& u0,
         array1[dof_bs1 * dofs1[i] + k] = local1[dof_bs1 * i + k];
   }
 }
+
+template <typename T>
+void interpolate_nonmatching_meshes(
+    Function<T>& u, const Function<T>& v, std::span<const std::int32_t> cells,
+    const nmm_interpolation_data_t& nmm_interpolation_data)
+{
+  int result;
+  auto mesh = u.function_space()->mesh();
+  auto mesh_v = v.function_space()->mesh();
+  MPI_Comm_compare(mesh->comm(), mesh_v->comm(), &result);
+
+  if (result == MPI_UNEQUAL)
+    throw std::runtime_error("Interpolation on different meshes is only "
+                             "supported with the same communicator.");
+
+  MPI_Comm comm = mesh->comm();
+  const int tdim = mesh->topology().dim();
+  const auto cell_map = mesh->topology().index_map(tdim);
+
+  std::shared_ptr<const FiniteElement> element_u
+      = u.function_space()->element();
+  const std::size_t value_size = element_u->value_size();
+
+  if (std::get<0>(nmm_interpolation_data).empty())
+  {
+    throw std::runtime_error(
+        "In order to interpolate on nonmatching meshes, the user needs to "
+        "provide the necessary interpolation data. This can be computed "
+        "with fem::create_nonmatching_meshes_interpolation_data.");
+  }
+
+  const std::tuple_element_t<0, nmm_interpolation_data_t>& dest_ranks
+      = std::get<0>(nmm_interpolation_data);
+  const std::tuple_element_t<1, nmm_interpolation_data_t>& src_ranks
+      = std::get<1>(nmm_interpolation_data);
+  const std::tuple_element_t<2, nmm_interpolation_data_t>& received_points
+      = std::get<2>(nmm_interpolation_data);
+  const std::tuple_element_t<3, nmm_interpolation_data_t>& evaluation_cells
+      = std::get<3>(nmm_interpolation_data);
+
+  // Evaluate the interpolating function where possible
+  std::vector<T> send_values(received_points.size() / 3 * value_size);
+  v.eval(received_points, {received_points.size() / 3, (std::size_t)3},
+         evaluation_cells, send_values,
+         {received_points.size() / 3, (std::size_t)value_size});
+
+  // Send values back to owning process
+  std::array<std::size_t, 2> v_shape = {src_ranks.size(), value_size};
+  std::vector<T> values_b(dest_ranks.size() * value_size);
+  impl::scatter_values(comm, src_ranks, dest_ranks,
+                       std::span<const T>(send_values), v_shape,
+                       std::span<T>(values_b));
+
+  // Transpose received data
+  namespace stdex = std::experimental;
+  stdex::mdspan<const T, stdex::dextents<std::size_t, 2>> values(
+      values_b.data(), dest_ranks.size(), value_size);
+
+  std::vector<T> valuesT_b(value_size * dest_ranks.size());
+  stdex::mdspan<T, stdex::dextents<std::size_t, 2>> valuesT(
+      valuesT_b.data(), value_size, dest_ranks.size());
+  for (std::size_t i = 0; i < values.extent(0); ++i)
+    for (std::size_t j = 0; j < values.extent(1); ++j)
+      valuesT(j, i) = values(i, j);
+
+  // Call local interpolation operator
+  fem::interpolate<T>(u, std::span(valuesT_b.data(), valuesT_b.size()),
+                      {valuesT.extent(0), valuesT.extent(1)}, cells);
+}
 //----------------------------------------------------------------------------
 } // namespace impl
 
@@ -835,103 +916,40 @@ void interpolate(Function<T>& u, std::span<const T> f,
   }
 }
 
-namespace impl
-{
-
-/// Interpolate from one finite element Function to another
-/// @param[in,out] u The function to interpolate into
-/// @param[in] v The function to be interpolated
-/// @param[in] cells List of cell indices to interpolate on
-template <typename T>
-void interpolate_nonmatching_meshes(Function<T>& u, const Function<T>& v,
-                                    std::span<const std::int32_t> cells)
-{
-  assert(u.function_space());
-  assert(v.function_space());
-
-  std::shared_ptr<const mesh::Mesh> mesh = u.function_space()->mesh();
-  std::shared_ptr<const mesh::Mesh> mesh_v = v.function_space()->mesh();
-
-  assert(mesh);
-
-  int result;
-  MPI_Comm_compare(mesh->comm(), mesh_v->comm(), &result);
-
-  if (result == MPI_UNEQUAL)
-    throw std::runtime_error("Interpolation on different meshes is only "
-                             "supported with the same communicator.");
-
-  MPI_Comm comm = mesh->comm();
-  const int tdim = mesh->topology().dim();
-  const auto cell_map = mesh->topology().index_map(tdim);
-
-  std::shared_ptr<const FiniteElement> element_u
-      = u.function_space()->element();
-  const std::size_t value_size = element_u->value_size();
-
-  // Collect all the points at which values are needed to define the
-  // interpolating function
-  std::vector<double> x;
-  {
-    const std::vector<double> coords_b
-        = fem::interpolation_coords(*element_u, *mesh, cells);
-
-    namespace stdex = std::experimental;
-    using cmdspan2_t
-        = stdex::mdspan<const double, stdex::dextents<std::size_t, 2>>;
-    using mdspan2_t = stdex::mdspan<double, stdex::dextents<std::size_t, 2>>;
-    cmdspan2_t coords(coords_b.data(), 3, coords_b.size() / 3);
-    // Transpose interpolation coords
-    x.resize(coords.size());
-    mdspan2_t _x(x.data(), coords_b.size() / 3, 3);
-    for (std::size_t j = 0; j < coords.extent(1); ++j)
-      for (std::size_t i = 0; i < 3; ++i)
-        _x(j, i) = coords(i, j);
-  }
-
-  // Determine ownership of each point
-  auto [dest_ranks, src_ranks, received_points, evaluation_cells]
-      = geometry::determine_point_ownership(*mesh_v, x);
-
-  // Evaluate the interpolating function where possible
-  std::vector<T> send_values(received_points.size() / 3 * value_size);
-  v.eval(received_points, {received_points.size() / 3, (std::size_t)3},
-         evaluation_cells, send_values,
-         {received_points.size() / 3, (std::size_t)value_size});
-
-  // Send values back to owning process
-  std::array<std::size_t, 2> v_shape = {src_ranks.size(), value_size};
-  std::vector<T> values_b(dest_ranks.size() * value_size);
-  impl::scatter_values(comm, src_ranks, dest_ranks,
-                       std::span<const T>(send_values), v_shape,
-                       std::span<T>(values_b));
-
-  // Transpose received data
-  namespace stdex = std::experimental;
-  stdex::mdspan<const T, stdex::dextents<std::size_t, 2>> values(
-      values_b.data(), dest_ranks.size(), value_size);
-
-  std::vector<T> valuesT_b(value_size * dest_ranks.size());
-  stdex::mdspan<T, stdex::dextents<std::size_t, 2>> valuesT(
-      valuesT_b.data(), value_size, dest_ranks.size());
-  for (std::size_t i = 0; i < values.extent(0); ++i)
-    for (std::size_t j = 0; j < values.extent(1); ++j)
-      valuesT(j, i) = values(i, j);
+/// Generate data needed to interpolate discrete functions across different
+/// meshes
+///
+/// @param[out] Vu The function space of the function to interpolate into
+/// @param[in] Vv The function space of the function to interpolate from
+/// @param[in] cells Indices of the cells in the destination mesh on which to
+/// interpolate. Should be the same as the list used when calling
+/// fem::interpolation_coords.
+nmm_interpolation_data_t create_nonmatching_meshes_interpolation_data(
+    const FunctionSpace& Vu, const FunctionSpace& Vv,
+    std::span<const std::int32_t> cells);
 
-  // Call local interpolation operator
-  fem::interpolate<T>(u, std::span(valuesT_b.data(), valuesT_b.size()),
-                      {valuesT.extent(0), valuesT.extent(1)}, cells);
-}
+/// Generate data needed to interpolate discrete functions defined on different
+/// meshes. Interpolate on all cells in the mesh.
+///
+/// @param[out] Vu The function space of the function to interpolate into
+/// @param[in] Vv The function space of the function to interpolate from
+nmm_interpolation_data_t
+create_nonmatching_meshes_interpolation_data(const FunctionSpace& Vu,
+                                             const FunctionSpace& Vv);
 
-} // namespace impl
 //----------------------------------------------------------------------------
 /// Interpolate from one finite element Function to another one
 /// @param[out] u The function to interpolate into
 /// @param[in] v The function to be interpolated
 /// @param[in] cells List of cell indices to interpolate on
+/// @param[in] nmm_interpolation_data Auxiliary data to interpolate on
+/// nonmatching meshes. This data can be generated with
+/// create_nonmatching_meshes_interpolation_data (optional).
 template <typename T>
 void interpolate(Function<T>& u, const Function<T>& v,
-                 std::span<const std::int32_t> cells)
+                 std::span<const std::int32_t> cells,
+                 const nmm_interpolation_data_t& nmm_interpolation_data
+                 = nmm_interpolation_data_t{})
 {
   assert(u.function_space());
   assert(v.function_space());
@@ -952,7 +970,10 @@ void interpolate(Function<T>& u, const Function<T>& v,
   {
     // Get mesh and check that functions share the same mesh
     if (auto mesh_v = v.function_space()->mesh(); mesh != mesh_v)
-      impl::interpolate_nonmatching_meshes(u, v, cells);
+    {
+      impl::interpolate_nonmatching_meshes(u, v, cells, nmm_interpolation_data);
+      return;
+    }
     else
     {
 
@@ -1022,5 +1043,4 @@ void interpolate(Function<T>& u, const Function<T>& v,
     }
   }
 }
-
 } // namespace dolfinx::fem

python/dolfinx/fem/__init__.py

@@ -6,6 +6,7 @@
 """Tools for assembling and manipulating finite element forms."""
 
 from dolfinx.cpp.fem import transpose_dofmap  # noqa
+from dolfinx.cpp.fem import create_nonmatching_meshes_interpolation_data
 from dolfinx.cpp.fem import IntegralType
 from dolfinx.cpp.fem import create_sparsity_pattern as _create_sparsity_pattern
 from dolfinx.fem import petsc
@@ -44,4 +45,4 @@ def create_sparsity_pattern(a: FormMetaClass):
     "assemble_scalar", "assemble_matrix", "assemble_vector", "apply_lifting", "set_bc",
     "DirichletBCMetaClass", "dirichletbc", "bcs_by_block", "DofMap", "FormMetaClass", "form", "IntegralType",
     "locate_dofs_geometrical", "locate_dofs_topological",
-    "extract_function_spaces", "petsc"]
+    "extract_function_spaces", "petsc", "create_nonmatching_meshes_interpolation_data"]

python/dolfinx/fem/function.py

@@ -332,7 +332,8 @@ def eval(self, x: npt.ArrayLike, cells: npt.ArrayLike, u=None) -> np.ndarray:
         return u
 
     def interpolate(self, u: typing.Union[typing.Callable, Expression, Function],
-                    cells: typing.Optional[np.ndarray] = None) -> None:
+                    cells: typing.Optional[np.ndarray] = None,
+                    nmm_interpolation_data=((), (), (), ())) -> None:
         """Interpolate an expression
 
         Args:
@@ -344,12 +345,12 @@ def interpolate(self, u: typing.Union[typing.Callable, Expression, Function],
         @singledispatch
         def _interpolate(u, cells: typing.Optional[np.ndarray] = None):
             """Interpolate a cpp.fem.Function"""
-            self._cpp_object.interpolate(u, cells)
+            self._cpp_object.interpolate(u, cells, nmm_interpolation_data)
 
         @_interpolate.register(Function)
         def _(u: Function, cells: typing.Optional[np.ndarray] = None):
             """Interpolate a fem.Function"""
-            self._cpp_object.interpolate(u._cpp_object, cells)
+            self._cpp_object.interpolate(u._cpp_object, cells, nmm_interpolation_data)
 
         @_interpolate.register(int)
         def _(u_ptr: int, cells: typing.Optional[np.ndarray] = None):