More work towards supporting mixed-topology meshes

cpp/dolfinx/fem/Form.h

@@ -258,6 +258,35 @@ class Form
     return _function_spaces;
   }
 
+  /// @brief Get the kernel function for integral `i` on given domain
+  /// type.
+  /// @param[in] type Integral type.
+  /// @param[in] i The subdomain ID.
+  /// @param[in] kernel_idx Index of the kernel (we may have multiple kernels
+  /// for a given ID in mixed-topology meshes).
+  /// @return Function to call for `tabulate_tensor`.
+  std::function<void(scalar_type*, const scalar_type*, const scalar_type*,
+                     const geometry_type*, const int*, const uint8_t*)>
+  kernel(IntegralType type, int i, int kernel_idx) const
+  {
+    const std::vector<integral_data<scalar_type, geometry_type>>& integrals
+        = _integrals[static_cast<std::size_t>(type)];
+
+    // Get the range of integrals with a given ID
+    auto get_id = [](const auto& a) { return a.id; };
+    auto start = std::ranges::lower_bound(integrals, i, std::less<>{}, get_id);
+    auto end = std::ranges::upper_bound(integrals, i, std::less<>{}, get_id);
+
+    // Check that the kernel is valid and return it if so
+    if (start == integrals.end() or start->id != i
+        or std::distance(start, end) <= kernel_idx)
+    {
+      throw std::runtime_error("No kernel for requested domain index.");
+    }
+
+    return std::next(start, kernel_idx)->kernel;
+  }
+
   /// @brief Get the kernel function for integral `i` on given domain
   /// type.
   /// @param[in] type Integral type.
@@ -267,13 +296,7 @@ class Form
                      const geometry_type*, const int*, const uint8_t*)>
   kernel(IntegralType type, int i) const
   {
-    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
-    auto it = std::ranges::lower_bound(integrals, i, std::less<>{},
-                                       [](const auto& a) { return a.id; });
-    if (it != integrals.end() and it->id == i)
-      return it->kernel;
-    else
-      throw std::runtime_error("No kernel for requested domain index.");
+    return kernel(type, i, 0);
   }
 
   /// @brief Get types of integrals in the form.
@@ -310,8 +333,7 @@ class Form
   /// @param[in] i Index of the integral.
   std::vector<int> active_coeffs(IntegralType type, std::size_t i) const
   {
-    assert(i < _integrals[static_cast<std::size_t>(type)].size());
-    return _integrals[static_cast<std::size_t>(type)][i].coeffs;
+    return _integrals[static_cast<std::size_t>(type)].at(i).coeffs;
   }
 
   /// @brief Get the IDs for integrals (kernels) for given integral type.
@@ -324,9 +346,15 @@ class Form
   std::vector<int> integral_ids(IntegralType type) const
   {
     std::vector<int> ids;
-    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
+    const std::vector<integral_data<scalar_type, geometry_type>>& integrals
+        = _integrals[static_cast<std::size_t>(type)];
     std::ranges::transform(integrals, std::back_inserter(ids),
                            [](auto& integral) { return integral.id; });
+
+    // IDs may be repeated in mixed-topology meshes, so remove duplicates
+    std::sort(ids.begin(), ids.end());
+    auto it = std::unique(ids.begin(), ids.end());
+    ids.erase(it, ids.end());
     return ids;
   }
 
@@ -344,12 +372,14 @@ class Form
   /// local_facet_index_1)`. Data is flattened with row-major layout,
   /// `shape=(num_facets, 4)`.
   ///
-  /// @param[in] type Integral domain type.
+  /// @param[in] type Integral type.
   /// @param[in] i Integral ID, i.e. (sub)domain index.
   /// @return List of active entities for the given integral (kernel).
   std::span<const std::int32_t> domain(IntegralType type, int i) const
   {
-    const auto& integrals = _integrals[static_cast<std::size_t>(type)];
+    // FIXME This should call domain with kernel_idx=0
+    const std::vector<integral_data<scalar_type, geometry_type>>& integrals
+        = _integrals[static_cast<std::size_t>(type)];
     auto it = std::ranges::lower_bound(integrals, i, std::less<>{},
                                        [](const auto& a) { return a.id; });
     if (it != integrals.end() and it->id == i)
@@ -358,24 +388,52 @@ class Form
       throw std::runtime_error("No mesh entities for requested domain index.");
   }
 
+  /// @brief Get the list of mesh entity indices for the ith integral
+  /// (kernel) of a given type.
+  /// @param[in] type Integral type.
+  /// @param[in] i Integral ID, i.e. (sub)domain index.
+  /// @param[in] kernel_idx Index of the kernel (we may have multiple kernels
+  /// for a given ID in mixed-topology meshes).
+  /// @return List of active entities for the given integral (kernel).
+  std::vector<std::int32_t> domain(IntegralType type, int i,
+                                   int kernel_idx) const
+  {
+    const std::vector<integral_data<scalar_type, geometry_type>>& integrals
+        = _integrals[static_cast<std::size_t>(type)];
+    auto get_id = [](const auto& a) { return a.id; };
+    auto start = std::ranges::lower_bound(integrals, i, std::less<>{}, get_id);
+    auto end = std::ranges::upper_bound(integrals, i, std::less<>{}, get_id);
+
+    // Check that the kernel is valid and return it if so
+    if (start == integrals.end() or start->id != i
+        or std::distance(start, end) <= kernel_idx)
+    {
+      throw std::runtime_error("No kernel for requested domain index.");
+    }
+
+    return std::next(start, kernel_idx)->entities;
+  }
+
   /// @brief Compute the list of entity indices in `mesh` for the ith
   /// integral (kernel) of a given type (i.e. cell, exterior facet, or
   /// interior facet).
   ///
   /// @param type Integral type.
   /// @param i Integral ID, i.e. the (sub)domain index.
+  /// @param kernel_idx Index of the kernel (we may have multiple
+  /// kernels for a given ID in mixed-topology meshes).
   /// @param mesh The mesh the entities are numbered with respect to.
   /// @return List of active entities in `mesh` for the given integral.
-  std::vector<std::int32_t> domain(IntegralType type, int i,
+  std::vector<std::int32_t> domain(IntegralType type, int i, int kernel_idx,
                                    const mesh::Mesh<geometry_type>& mesh) const
   {
     // Hack to avoid passing shared pointer to this function
     std::shared_ptr<const mesh::Mesh<geometry_type>> msh_ptr(
         &mesh, [](const mesh::Mesh<geometry_type>*) {});
 
-    std::span<const std::int32_t> entities = domain(type, i);
+    std::vector<std::int32_t> entities = domain(type, i, kernel_idx);
     if (msh_ptr == _mesh)
-      return std::vector(entities.begin(), entities.end());
+      return entities;
     else
     {
       std::span<const std::int32_t> entity_map = _entity_maps.at(msh_ptr);
@@ -416,6 +474,7 @@ class Form
             // Get the facet index
             const std::int32_t facet
                 = c_to_f->links(entities[i])[entities[i + 1]];
+
             // Add cell and the local facet index
             mapped_entities.insert(mapped_entities.end(),
                                    {entity_map[facet], entities[i + 1]});
@@ -443,16 +502,17 @@ class Form
         }
         else if (codim == 1)
         {
-          // In this case, the entity maps take facets in (`_mesh`) to cells in
-          // `mesh`, so we need to get the facet number from the (cell,
-          // local_facet pair) first.
+          // In this case, the entity maps take facets in (`_mesh`) to
+          // cells in `mesh`, so we need to get the facet number from
+          // the (cell, local_facet pair) first.
           auto c_to_f = _mesh->topology()->connectivity(tdim, tdim - 1);
           assert(c_to_f);
           for (std::size_t i = 0; i < entities.size(); i += 2)
           {
             // Get the facet index
             const std::int32_t facet
                 = c_to_f->links(entities[i])[entities[i + 1]];
+
             // Add cell and the local facet index
             mapped_entities.insert(mapped_entities.end(),
                                    {entity_map[facet], entities[i + 1]});
@@ -468,6 +528,20 @@ class Form
     }
   }
 
+  /// @brief Compute the list of entity indices in `mesh` for the ith
+  /// integral (kernel) of a given type (i.e. cell, exterior facet, or
+  /// interior facet).
+  ///
+  /// @param type Integral type.
+  /// @param i Integral ID, i.e. the (sub)domain index.
+  /// @param mesh The mesh the entities are numbered with respect to.
+  /// @return List of active entities in `mesh` for the given integral.
+  std::vector<std::int32_t> domain(IntegralType type, int i,
+                                   const mesh::Mesh<geometry_type>& mesh) const
+  {
+    return domain(type, i, 0, mesh);
+  }
+
   /// @brief Access coefficients.
   const std::vector<
       std::shared_ptr<const Function<scalar_type, geometry_type>>>&
@@ -531,5 +605,5 @@ class Form
   std::map<std::shared_ptr<const mesh::Mesh<geometry_type>>,
            std::vector<std::int32_t>>
       _entity_maps;
-}; // namespace dolfinx::fem
+};
 } // namespace dolfinx::fem