Let integration entities be specified manually

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/utils.cpp

@@ -266,3 +266,119 @@ 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(const 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();
+  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));
+
+  // Structure to store (meshtag value, entity) pairs
+  std::vector<std::pair<int, std::vector<std::int32_t>>> value_entity_pairs;
+  switch (integral_type)
+  {
+    // TODO: Sort pairs or use std::iota
+  case fem::IntegralType::cell:
+    for (std::size_t i = 0; i < values.size(); ++i)
+      value_entity_pairs.push_back({values[i], {entities[i]}});
+    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);
+        value_entity_pairs.push_back(
+            {values[pos], std::vector(facet.begin(), facet.end())});
+      }
+    }
+    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);
+
+          value_entity_pairs.push_back(
+              {values[j], std::vector(facets.begin(), facets.end())});
+        }
+      }
+    }
+    break;
+    default:
+      throw std::runtime_error(
+          "Cannot compute integration domains. Integral type not supported.");
+    }
+  }
+
+  // Sort pairs by meshtag value so that entities can be grouped
+  std::sort(value_entity_pairs.begin(), value_entity_pairs.end());
+
+  std::vector<std::pair<int, std::vector<std::int32_t>>> integrals;
+  // Iterator to mark the start of the group
+  auto group_start_it = value_entity_pairs.begin();
+  while (group_start_it != value_entity_pairs.end())
+  {
+    // Get iterator pointing to end of group
+    auto val = group_start_it->first;
+    auto group_end_it = std::lower_bound(
+        value_entity_pairs.begin(), value_entity_pairs.end(), val,
+        [](auto& pair, auto val) { return pair.first <= val; });
+
+    // Meshtag value for this group
+    const int id = group_start_it->first;
+    // List to store entities in this group
+    std::vector<std::int32_t> group_entities;
+    // Loop through entities in this group and add
+    for (auto it = group_start_it; it != group_end_it; ++it)
+    {
+      const auto entity = it->second;
+      group_entities.insert(group_entities.end(), entity.begin(), entity.end());
+    }
+    integrals.push_back({id, std::move(group_entities)});
+    group_start_it = group_end_it;
+  }
+
+  return integrals;
+}
+//-----------------------------------------------------------------------------

cpp/dolfinx/fem/utils.h

@@ -103,6 +103,20 @@ using scalar_value_type_t = typename scalar_value_type<T>::value_type;
 
 } // namespace impl
 
+/// @brief Given an integral type and MeshTags compute
+/// the entities that should be integrated over.
+///
+/// I.e., local cell indices for cell integrals, local facets on the
+/// boundary as (cell, local facet index) pairs for exterior facet
+/// integrals etc.)
+///
+/// @param[in] integral_type The integral type
+/// @param[in] meshtags The meshtags
+/// @return A list of (integral id, entities) pairs
+std::vector<std::pair<int, std::vector<std::int32_t>>>
+compute_integration_domains(const IntegralType integral_type,
+                            const mesh::MeshTags<int>& meshtags);
+
 /// @brief Finite element cell kernel concept.
 ///
 /// Kernel functions that can be passed to an assembler for execution
@@ -262,14 +276,18 @@ std::vector<std::string> get_constant_names(const ufcx_form& ufcx_form);
 /// @param[in] coefficients Coefficient fields in the form
 /// @param[in] constants Spatial constants in the form
 /// @param[in] subdomains Subdomain markers
+/// @pre Each value in `subdomains` must be sorted by domain id
 /// @param[in] mesh The mesh of the domain
 template <typename T>
 Form<T> create_form(
     const ufcx_form& ufcx_form,
     const std::vector<std::shared_ptr<const FunctionSpace>>& spaces,
     const std::vector<std::shared_ptr<const Function<T>>>& coefficients,
     const std::vector<std::shared_ptr<const Constant<T>>>& constants,
-    const std::map<IntegralType, const mesh::MeshTags<int>*>& subdomains,
+    const std::map<
+        IntegralType,
+        std::vector<std::pair<std::int32_t, std::vector<std::int32_t>>>>&
+        subdomains,
     std::shared_ptr<const mesh::Mesh> mesh = nullptr)
 {
   if (ufcx_form.rank != (int)spaces.size())
@@ -370,29 +388,25 @@ Form<T> create_form(
       assert(k);
 
       // Build list of entities to assembler over
-      std::vector<std::int32_t> e;
       if (id == -1)
       {
         // Default kernel, operates on all (owned) cells
         assert(topology.index_map(tdim));
+        std::vector<std::int32_t> e;
         e.resize(topology.index_map(tdim)->size_local(), 0);
         std::iota(e.begin(), e.end(), 0);
+        itg.first->second.emplace_back(id, k, std::move(e));
       }
-      else if (sd != subdomains.end() and sd->second)
+      else if (sd != subdomains.end())
       {
-        assert(topology.index_map(tdim));
-        std::span<const std::int32_t> entities = sd->second->indices();
-        std::span<const int> values = sd->second->values();
-        auto it0 = entities.begin();
-        auto it1 = std::lower_bound(it0, entities.end(),
-                                    topology.index_map(tdim)->size_local());
-        entities = entities.first(std::distance(it0, it1));
-        for (std::size_t j = 0; j < entities.size(); ++j)
-          if (values[j] == id)
-            e.push_back(entities[j]);
+        // NOTE: This requires that pairs are sorted
+        auto it = std::lower_bound(sd->second.begin(), sd->second.end(), id,
+                                   [](auto& pair, auto val)
+                                   { return pair.first < val; });
+        if (it != sd->second.end() && (*it).first == id)
+          itg.first->second.emplace_back(id, k, (*it).second);
       }
 
-      itg.first->second.emplace_back(id, k, std::move(e));
       if (integral->needs_facet_permutations)
         needs_facet_permutations = true;
     }
@@ -432,7 +446,6 @@ Form<T> create_form(
       assert(k);
 
       // Build list of entities to assembler over
-      std::vector<std::int32_t> e;
       const std::vector bfacets = mesh::exterior_facet_indices(topology);
       auto f_to_c = topology.connectivity(tdim - 1, tdim);
       assert(f_to_c);
@@ -441,6 +454,7 @@ Form<T> create_form(
       if (id == -1)
       {
         // Default kernel, operates on all (owned) exterior facets
+        std::vector<std::int32_t> e;
         e.reserve(2 * bfacets.size());
         for (std::int32_t f : bfacets)
         {
@@ -449,34 +463,18 @@ Form<T> create_form(
               = impl::get_cell_facet_pairs<1>(f, f_to_c->links(f), *c_to_f);
           e.insert(e.end(), pair.begin(), pair.end());
         }
+        itg.first->second.emplace_back(id, k, std::move(e));
       }
-      else if (sd != subdomains.end() and sd->second)
+      else if (sd != subdomains.end())
       {
-        // Create list of tagged boundary facets
-        std::span<const std::int32_t> entities = sd->second->indices();
-        auto it0 = entities.begin();
-        auto it1 = std::lower_bound(it0, entities.end(),
-                                    topology.index_map(tdim - 1)->size_local());
-        entities = entities.first(std::distance(it0, it1));
-        std::span<const int> values = sd->second->values();
-        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);
-          if (values[pos] == id)
-          {
-            auto facet
-                = impl::get_cell_facet_pairs<1>(f, f_to_c->links(f), *c_to_f);
-            e.insert(e.end(), facet.begin(), facet.end());
-          }
-        }
+        // NOTE: This requires that pairs are sorted
+        auto it = std::lower_bound(sd->second.begin(), sd->second.end(), id,
+                                   [](auto& pair, auto val)
+                                   { return pair.first < val; });
+        if (it != sd->second.end() && (*it).first == id)
+          itg.first->second.emplace_back(id, k, (*it).second);
       }
 
-      itg.first->second.emplace_back(id, k, std::move(e));
       if (integral->needs_facet_permutations)
         needs_facet_permutations = true;
     }
@@ -516,14 +514,14 @@ Form<T> create_form(
       assert(k);
 
       // Build list of entities to assembler over
-      std::vector<std::int32_t> e;
       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);
       if (id == -1)
       {
         // Default kernel, operates on all (owned) interior facets
+        std::vector<std::int32_t> e;
         assert(topology.index_map(tdim - 1));
         std::int32_t num_facets = topology.index_map(tdim - 1)->size_local();
         e.reserve(4 * num_facets);
@@ -536,30 +534,17 @@ Form<T> create_form(
             e.insert(e.end(), pairs.begin(), pairs.end());
           }
         }
+        itg.first->second.emplace_back(id, k, std::move(e));
       }
-      else if (sd != subdomains.end() and sd->second)
+      else if (sd != subdomains.end())
       {
-        std::span<const std::int32_t> entities = sd->second->indices();
-        auto it0 = entities.begin();
-        auto it1 = std::lower_bound(it0, entities.end(),
-                                    topology.index_map(tdim - 1)->size_local());
-        entities = entities.first(std::distance(it0, it1));
-        std::span<const int> values = sd->second->values();
-        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 and values[j] == id)
-          {
-            // 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);
-            e.insert(e.end(), facets.begin(), facets.end());
-          }
-        }
+        auto it = std::lower_bound(sd->second.begin(), sd->second.end(), id,
+                                   [](auto& pair, auto val)
+                                   { return pair.first < val; });
+        if (it != sd->second.end() && (*it).first == id)
+          itg.first->second.emplace_back(id, k, (*it).second);
       }
 
-      itg.first->second.emplace_back(id, k, std::move(e));
       if (integral->needs_facet_permutations)
         needs_facet_permutations = true;
     }
@@ -575,6 +560,7 @@ Form<T> create_form(
 /// @param[in] coefficients Coefficient fields in the form (by name)
 /// @param[in] constants Spatial constants in the form (by name)
 /// @param[in] subdomains Subdomain makers
+/// @pre Each value in `subdomains` must be sorted by domain id
 /// @param[in] mesh The mesh of the domain. This is required if the form
 /// has no arguments, e.g. a functional
 /// @return A Form
@@ -585,7 +571,10 @@ Form<T> create_form(
     const std::map<std::string, std::shared_ptr<const Function<T>>>&
         coefficients,
     const std::map<std::string, std::shared_ptr<const Constant<T>>>& constants,
-    const std::map<IntegralType, const mesh::MeshTags<int>*>& subdomains,
+    const std::map<
+        IntegralType,
+        std::vector<std::pair<std::int32_t, std::vector<std::int32_t>>>>&
+        subdomains,
     std::shared_ptr<const mesh::Mesh> mesh = nullptr)
 {
   // Place coefficients in appropriate order
@@ -622,6 +611,7 @@ Form<T> create_form(
 /// @param[in] coefficients Coefficient fields in the form (by name)
 /// @param[in] constants Spatial constants in the form (by name)
 /// @param[in] subdomains Subdomain markers
+/// @pre Each value in `subdomains` must be sorted by domain id
 /// @param[in] mesh The mesh of the domain. This is required if the form
 /// has no arguments, e.g. a functional.
 /// @return A Form
@@ -632,7 +622,10 @@ Form<T> create_form(
     const std::map<std::string, std::shared_ptr<const Function<T>>>&
         coefficients,
     const std::map<std::string, std::shared_ptr<const Constant<T>>>& constants,
-    const std::map<IntegralType, const mesh::MeshTags<int>*>& subdomains,
+    const std::map<
+        IntegralType,
+        std::vector<std::pair<std::int32_t, std::vector<std::int32_t>>>>&
+        subdomains,
     std::shared_ptr<const mesh::Mesh> mesh = nullptr)
 {
   ufcx_form* form = fptr();