Let integration entities be specified manually (#2269)

* Start work

* Pass integration entities

* Change Form constructor

* Restructure integral data

* Set domains

* Same for ext and int facets

* Rename new integral data

* Set default domains

* Start work on converting meshtags

* Fix for parallel

* Start work on ext facets

* Move function

* Get ext facets working

* Same for int facets

* Add comment

* Add docs

* Fix docs

* Remove old code

* Add test

* Fix for parallel

* Remove cout

* Add fix for measures without subdomain data

* Flake8

* Bind Form constructor

* Bind Form constructor

* Make const

* Fix and simplify

* Make const reference

* Flake8

* Fix consts

* Update test

* Add FIXMEs

* Format

* Remove commented code

* Remove old code

* Add to test

* Remove duplicate code

* Simplify

* Improve test

* Update error message

* Remove TODO

* FIXMEs

* Add NOTE

* Update docs

* Simplify

* Docs

* Docs

* Specify type explicitly

* Use exterior_facet_indices

* Reduce number of map lookups

* Improve docs

* fix typo

* Don't use isinstance

* Don't use isinstance

* Move subdomain data check to dolfinx

* Flake8

* Flake8

* Flake8

* Mypy

* Use UFL branch

* Put on single line

* Workaround for strange issue

* Flake8

* Add TODOs

* Check all subdomain data is the same

* Use UFL branch

* Pass partitioner

* Fix typo

* Try to use sd.get(domain)

* Change UFL branch

* Change UFL branch

* Use std::size_t

* Fix bug

* Some fixes

* Move meshtag to entity code

* Fix bug

* Fix for parallel

* Fix bug where ghost entites were being added to integration domains by default

* Simplify

* Remove TODO

* Remove unused function

* Avoid map lookup in tight loop

* Rename

* Same for facet integrals

* Remove unnecessary sort

* Remove function

* Remove function

* Tidy

* Remove second map

* Change type

* Update test

* Tidy

* Update type

* Tidy

* Fix seg fault

* Tidy

* Simplify

* Simplify

* Small style fixes

* Readability improvement.

* Improve pybind11 wrapping

* Avoid copies

* Simplifications

* Small fix

* Revert some changes

* Interface fix

* Work on doc

* Avoid costly allocations

* Cleanup

* Use stable sort

* Simplificstions

* Small naming improvements

* Bug fix

* Another fix

* Tidy up

---------

Co-authored-by: Jørgen Schartum Dokken <jsd55@cam.ac.uk>
Co-authored-by: Chris Richardson <chris@bpi.cam.ac.uk>
Co-authored-by: Garth N. Wells <gnw20@cam.ac.uk>

cpp/dolfinx/fem/Form.h

@@ -119,11 +119,8 @@ class Form
       throw std::runtime_error("No mesh could be associated with the Form.");
 
     // Store kernels, looping over integrals by domain type (dimension)
-    for (auto& integral_type : integrals)
+    for (auto& [type, kernels] : integrals)
     {
-      const IntegralType type = integral_type.first;
-      auto& kernels = integral_type.second;
-
       // Loop over integrals kernels and set domains
       switch (type)
       {

cpp/dolfinx/fem/Function.h

@@ -172,7 +172,6 @@ class Function
     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);
-
     interpolate(v, cells, nmm_interpolation_data);
   }
 
@@ -217,6 +216,7 @@ class Function
         throw std::runtime_error(
             "Data returned by callable has wrong shape(0) size");
       }
+
       if (fshape[1] != x.size() / 3)
       {
         throw std::runtime_error(
@@ -285,6 +285,7 @@ class Function
             "Function element interpolation points has different shape to "
             "Expression interpolation points");
       }
+
       for (std::size_t i = 0; i < X0.size(); ++i)
       {
         if (std::abs(X0[i] - X1[i]) > 1.0e-10)
@@ -295,9 +296,8 @@ class Function
       }
     }
 
-    namespace stdex = std::experimental;
-
     // Array to hold evaluated Expression
+    namespace stdex = std::experimental;
     std::size_t num_cells = cells.size();
     std::size_t num_points = e.X().second[0];
     std::vector<T> fdata(num_cells * num_points * value_size);
@@ -312,7 +312,6 @@ class Function
     // value_size), i.e. xyzxyz ordering of dof values per cell per point.
     // The interpolation uses xxyyzz input, ordered for all points of each
     // cell, i.e. (value_size, num_cells*num_points)
-
     std::vector<T> fdata1(num_cells * num_points * value_size);
     stdex::mdspan<T, stdex::dextents<std::size_t, 3>> f1(
         fdata1.data(), value_size, num_cells, num_points);
@@ -371,6 +370,7 @@ class Function
       throw std::runtime_error(
           "Number of points and number of cells must be equal.");
     }
+
     if (xshape[0] != ushape[0])
     {
       throw std::runtime_error(

cpp/dolfinx/fem/assembler.h

@@ -27,7 +27,7 @@ class FunctionSpace;
 
 /// @brief Create a map of `std::span`s from a map of `std::vector`s
 template <typename T>
-std::map<std::pair<fem::IntegralType, int>, std::pair<std::span<const T>, int>>
+std::map<std::pair<IntegralType, int>, std::pair<std::span<const T>, int>>
 make_coefficients_span(const std::map<std::pair<IntegralType, int>,
                                       std::pair<std::vector<T>, int>>& coeffs)
 {
@@ -394,7 +394,7 @@ void set_diagonal(auto set_fn, std::span<const std::int32_t> rows,
 /// @param[in] diagonal The value to add to the diagonal for rows with a
 /// boundary condition applied
 template <typename T>
-void set_diagonal(auto set_fn, const fem::FunctionSpace& V,
+void set_diagonal(auto set_fn, const FunctionSpace& V,
                   const std::vector<std::shared_ptr<const DirichletBC<T>>>& bcs,
                   T diagonal = 1.0)
 {

cpp/dolfinx/fem/interpolate.cpp

@@ -69,40 +69,43 @@ 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
+  auto element_u = Vu.element();
+  assert(element_u);
+  auto mesh = Vu.mesh();
+  assert(mesh);
   const std::vector<double> coords_b
-      = fem::interpolation_coords(*element_u, *mesh, cells);
+      = 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());
+  std::vector<double> x(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);
+  auto mesh_v = Vv.mesh();
+  assert(mesh_v);
+  return 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)
+//-----------------------------------------------------------------------------
+fem::nmm_interpolation_data_t
+fem::create_nonmatching_meshes_interpolation_data(const FunctionSpace& Vu,
+                                                  const FunctionSpace& Vv)
 {
   assert(Vu.mesh());
   int tdim = Vu.mesh()->topology().dim();
@@ -111,7 +114,6 @@ fem::nmm_interpolation_data_t fem::create_nonmatching_meshes_interpolation_data(
   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

@@ -916,28 +916,28 @@ void interpolate(Function<T>& u, std::span<const T> f,
   }
 }
 
-/// Generate data needed to interpolate discrete functions across different
-/// meshes
+/// Generate data needed to interpolate discrete functions across
+/// different meshes.
 ///
-/// @param[out] Vu The function space of the function to interpolate into
+/// @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.
+/// @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);
 
-/// Generate data needed to interpolate discrete functions defined on different
-/// meshes. Interpolate on all cells in the mesh.
+/// 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);
 
-//----------------------------------------------------------------------------
 /// 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

cpp/dolfinx/fem/utils.cpp

@@ -12,6 +12,7 @@
 #include "Function.h"
 #include "FunctionSpace.h"
 #include "dofmapbuilder.h"
+#include <algorithm>
 #include <array>
 #include <dolfinx/common/IndexMap.h>
 #include <dolfinx/common/Timer.h>
@@ -197,8 +198,6 @@ fem::FunctionSpace fem::create_functionspace(
     const std::function<std::vector<int>(
         const graph::AdjacencyList<std::int32_t>&)>& reorder_fn)
 {
-  assert(mesh);
-
   // Create a DOLFINx element
   auto _e = std::make_shared<FiniteElement>(e, bs);
 
@@ -220,6 +219,7 @@ fem::FunctionSpace fem::create_functionspace(
   // Create a dofmap
   ElementDofLayout layout(bs, e.entity_dofs(), e.entity_closure_dofs(), {},
                           sub_doflayout);
+  assert(mesh);
   auto dofmap = std::make_shared<const DofMap>(
       create_dofmap(mesh->comm(), layout, mesh->topology(), reorder_fn, *_e));
 
@@ -266,3 +266,128 @@ fem::FunctionSpace fem::create_functionspace(
           mesh->comm(), layout, mesh->topology(), reorder_fn, *element)));
 }
 //-----------------------------------------------------------------------------
+std::vector<std::pair<int, std::vector<std::int32_t>>>
+fem::compute_integration_domains(fem::IntegralType integral_type,
+                                 const mesh::MeshTags<int>& meshtags)
+{
+  std::shared_ptr<const mesh::Mesh> mesh = meshtags.mesh();
+  assert(mesh);
+  const mesh::Topology& topology = mesh->topology();
+  const int tdim = topology.dim();
+  const int dim = integral_type == IntegralType::cell ? tdim : tdim - 1;
+  if (dim != meshtags.dim())
+  {
+    throw std::runtime_error("Invalid MeshTags dimension: "
+                             + std::to_string(meshtags.dim()));
+  }
+
+  std::span<const std::int32_t> entities = meshtags.indices();
+  std::span<const int> values = meshtags.values();
+  {
+    assert(topology.index_map(dim));
+    auto it0 = entities.begin();
+    auto it1 = std::lower_bound(it0, entities.end(),
+                                topology.index_map(dim)->size_local());
+    entities = entities.first(std::distance(it0, it1));
+    values = values.first(std::distance(it0, it1));
+  }
+
+  std::vector<std::int32_t> entity_data;
+  std::vector<int> values1;
+  switch (integral_type)
+  {
+  case IntegralType::cell:
+    entity_data.insert(entity_data.begin(), entities.begin(), entities.end());
+    values1.insert(values1.begin(), values.begin(), values.end());
+    break;
+  default:
+    mesh->topology_mutable().create_connectivity(dim, tdim);
+    mesh->topology_mutable().create_connectivity(tdim, dim);
+    auto f_to_c = topology.connectivity(tdim - 1, tdim);
+    assert(f_to_c);
+    auto c_to_f = topology.connectivity(tdim, tdim - 1);
+    assert(c_to_f);
+    switch (integral_type)
+    {
+    case IntegralType::exterior_facet:
+    {
+      // Create list of tagged boundary facets
+      const std::vector bfacets = mesh::exterior_facet_indices(topology);
+      std::vector<std::int32_t> facets;
+      std::set_intersection(entities.begin(), entities.end(), bfacets.begin(),
+                            bfacets.end(), std::back_inserter(facets));
+      for (auto f : facets)
+      {
+        auto index_it = std::lower_bound(entities.begin(), entities.end(), f);
+        assert(index_it != entities.end() and *index_it == f);
+        std::size_t pos = std::distance(entities.begin(), index_it);
+        auto facet
+            = impl::get_cell_facet_pairs<1>(f, f_to_c->links(f), *c_to_f);
+        entity_data.insert(entity_data.end(), facet.begin(), facet.end());
+        values1.push_back(values[pos]);
+      }
+    }
+    break;
+    case IntegralType::interior_facet:
+    {
+      for (std::size_t j = 0; j < entities.size(); ++j)
+      {
+        const std::int32_t f = entities[j];
+        if (f_to_c->num_links(f) == 2)
+        {
+          // Get the facet as a pair of (cell, local facet) pairs, one
+          // for each cell
+          auto facets
+              = impl::get_cell_facet_pairs<2>(f, f_to_c->links(f), *c_to_f);
+          entity_data.insert(entity_data.end(), facets.begin(), facets.end());
+          values1.push_back(values[j]);
+        }
+      }
+    }
+    break;
+    default:
+      throw std::runtime_error(
+          "Cannot compute integration domains. Integral type not supported.");
+    }
+  }
+
+  // Build permutation that sorts by meshtag value
+  std::vector<std::int32_t> perm(values1.size());
+  std::iota(perm.begin(), perm.end(), 0);
+  std::stable_sort(perm.begin(), perm.end(),
+                   [&values1](auto p0, auto p1)
+                   { return values1[p0] < values1[p1]; });
+
+  std::size_t shape = 1;
+  if (integral_type == IntegralType::exterior_facet)
+    shape = 2;
+  else if (integral_type == IntegralType::interior_facet)
+    shape = 4;
+  std::vector<std::pair<int, std::vector<std::int32_t>>> integrals;
+  {
+    // Iterator to mark the start of the group
+    auto p0 = perm.begin();
+    while (p0 != perm.end())
+    {
+      auto id0 = values1[*p0];
+      auto p1 = std::find_if_not(p0, perm.end(),
+                                 [id0, &values1](auto idx)
+                                 { return id0 == values1[idx]; });
+
+      std::vector<std::int32_t> data;
+      data.reserve(shape * std::distance(p0, p1));
+      for (auto it = p0; it != p1; ++it)
+      {
+        auto e_it0 = std::next(entity_data.begin(), (*it) * shape);
+        auto e_it1 = std::next(e_it0, shape);
+        data.insert(data.end(), e_it0, e_it1);
+      }
+
+      integrals.push_back({id0, std::move(data)});
+      p0 = p1;
+    }
+  }
+
+  return integrals;
+}
+//-----------------------------------------------------------------------------