Refactor dolfinx::fem::Form to use local indexing of forms, rather than integral ids

cpp/demo/custom_kernel/main.cpp

@@ -75,7 +75,7 @@ double assemble_matrix0(std::shared_ptr<const fem::FunctionSpace<T>> V,
 {
   // Kernel data (ID, kernel function, cell indices to execute over)
   std::map integrals{
-      std::pair{std::tuple{fem::IntegralType::cell, -1, 0},
+      std::pair{std::tuple{fem::IntegralType::cell, 0, 0},
                 fem::integral_data<T>(kernel, cells, std::vector<int>{})}};
 
   fem::Form<T, T> a({V, V}, integrals, V->mesh(), {}, {}, false, {});
@@ -105,7 +105,7 @@ double assemble_vector0(std::shared_ptr<const fem::FunctionSpace<T>> V,
 {
   auto mesh = V->mesh();
   std::map integrals{
-      std::pair{std::tuple{fem::IntegralType::cell, -1, 0},
+      std::pair{std::tuple{fem::IntegralType::cell, 0, 0},
                 fem::integral_data<T>(kernel, cells, std::vector<int>{})}};
   fem::Form<T> L({V}, integrals, mesh, {}, {}, false, {});
   auto dofmap = V->dofmap();

cpp/dolfinx/fem/Form.h

@@ -129,9 +129,12 @@ class Form
   /// @param[in] V Function spaces for the form arguments, e.g. test and
   /// trial function spaces.
   /// @param[in] integrals Integrals in the form, where
-  /// `integrals[IntegralType, domain ID, kernel index]` returns the
-  /// integral (`integral_data`) of type `IntegralType` over domain `ID`
-  /// with kernel index `kernel index`.
+  /// `integrals[IntegralType, i, kernel index]` returns the `i`th integral
+  /// (`integral_data`) of type `IntegralType` with kernel index `kernel index`.
+  /// The `i`-index refers to the position of a kernel when flattened by
+  /// sorted subdomain ids, sorted by subdomain ids. The subdomain ids can
+  /// contain duplicate entries referring to different kernels over the same
+  /// subdomain.
   /// @param[in] coefficients Coefficients in the form.
   /// @param[in] constants Constants in the form.
   /// @param[in] mesh Mesh of the domain to integrate over (the
@@ -247,7 +250,7 @@ class Form
 
     for (auto& space : _function_spaces)
     {
-      // Working map: [integral type, domain ID, kernel_idx]->entities
+      // Working map: [integral type, integral_idx, kernel_idx]->entities
       std::map<std::tuple<IntegralType, int, int>,
                std::variant<std::vector<std::int32_t>,
                             std::span<const std::int32_t>>>
@@ -269,7 +272,7 @@ class Form
         bool inverse = emap.sub_topology() == mesh0->topology();
         for (auto& [key, itg] : _integrals)
         {
-          auto [type, id, kernel_idx] = key;
+          auto [type, idx, kernel_idx] = key;
           std::vector<std::int32_t> e;
           if (type == IntegralType::cell)
             e = emap.sub_topology_to_topology(itg.entities, inverse);
@@ -299,13 +302,13 @@ class Form
 
     for (auto& [key, integral] : _integrals)
     {
-      auto [type, id, kernel_idx] = key;
+      auto [type, idx, kernel_idx] = key;
       for (int c : integral.coeffs)
       {
         if (auto mesh0 = coefficients.at(c)->function_space()->mesh();
             mesh0 == _mesh)
         {
-          _cdata.insert({{type, id, c}, std::span(integral.entities)});
+          _cdata.insert({{type, idx, c}, std::span(integral.entities)});
         }
         else
         {
@@ -331,8 +334,7 @@ class Form
           }
           else
             throw std::runtime_error("Integral type not supported.");
-
-          _cdata.insert({{type, id, c}, std::move(e)});
+          _cdata.insert({{type, idx, c}, std::move(e)});
         }
       }
     }
@@ -370,6 +372,9 @@ class Form
   }
 
   /// @brief Get the kernel function for an integral.
+  ///
+  ///
+  ///
   /// @param[in] type Integral type.
   /// @param[in] id Integral subdomain ID.
   /// @param[in] kernel_idx Index of the kernel (we may have multiple
@@ -418,31 +423,29 @@ class Form
     return it->second.coeffs;
   }
 
-  /// @brief Get the IDs for integrals (kernels) for given integral
-  /// domain type.
-  ///
-  /// The IDs correspond to the domain IDs which the integrals are
-  /// defined for in the form. `ID=-1` is the default integral over the
-  /// whole domain.
+  /// @brief Get number of integrals (kernels) for a given integral type and
+  /// kernel index.
   ///
+  /// For a form containing two integrals of `integral_a` and `integral_b`
+  /// with subdomain-ids `(1, 4)` and `(3, 4, 5)` respectively, the integrals
+  /// are stored as a flattened list, sorted by sudomain-ids
+  /// ```cpp
+  /// auto form_integrals = {integral_a, integral_b, integral_a, integral_b,
+  /// integral_b}; auto form_integral_ids = {1, 3, 4, 4, 5}.
+  /// ```
   /// @param[in] type Integral type.
-  /// @return List of IDs for given integral type.
-  std::vector<int> integral_ids(IntegralType type) const
+  /// @param[in] kernel_idx Index of the kernel (we may have multiple
+  /// kernels for a integral type in mixed-topology meshes).
+  int num_integrals(IntegralType type, int kernel_idx) const
   {
-    std::vector<int> ids;
-    for (auto& [key, integral] : _integrals)
-    {
-      auto [t, id, kernel_idx] = key;
-      if (t == type)
-        ids.push_back(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;
+    int count = std::count_if(_integrals.begin(), _integrals.end(),
+                              [type, kernel_idx](auto& x)
+                              {
+                                auto [t, id, k_idx] = x.first;
+                                return t == type and k_idx == kernel_idx;
+                              });
+    return count;
   }
 
   /// @brief Mesh entity indices to integrate over for a given integral
@@ -463,16 +466,23 @@ class Form
   /// is row-major.
   ///
   /// @param[in] type Integral type.
-  /// @param[in] id Integral domain identifier.
+  /// @param[in] idx Integral index in flattened list of integral kernels.
+  /// For a form containing two integrals of `integral_a` and `integral_b`
+  /// with subdomain-ids `(1, 4)` and `(3, 4, 5)` respectively, the integrals
+  /// are stored as a flattened list, sorted by sudomain-ids
+  /// ```cpp
+  /// auto form_integrals = {integral_a, integral_b, integral_a, integral_b,
+  /// integral_b}; auto form_integral_ids = {1, 3, 4, 4, 5}.
+  /// ```
   /// @param[in] kernel_idx Index of the kernel with in the domain (we
   /// may have multiple kernels for a given ID in mixed-topology
   /// meshes).
   /// @return Entity indices in the mesh::Mesh returned by mesh() to
   /// integrate over.
-  std::span<const std::int32_t> domain(IntegralType type, int id,
+  std::span<const std::int32_t> domain(IntegralType type, int idx,
                                        int kernel_idx) const
   {
-    auto it = _integrals.find({type, id, kernel_idx});
+    auto it = _integrals.find({type, idx, kernel_idx});
     if (it == _integrals.end())
       throw std::runtime_error("Requested domain not found.");
     return it->second.entities;
@@ -501,21 +511,21 @@ class Form
   /// may exist in one domain but not another. In this case, the entity
   /// is marked with -1.
   ///
-  /// @param type Integral type.
-  /// @param rank Argument index, e.g. `0` for the test function space, `1`
+  /// @param[in] type Integral type.
+  /// @param[in] rank Argument index, e.g. `0` for the test function space, `1`
   /// for the trial function space.
-  /// @param id Integral domain identifier.
-  /// @param kernel_idx Kernel index (cell type).
+  /// @param[in] idx Integral identifier.
+  /// @param[in] kernel_idx Kernel index (cell type).
   /// @return Entity indices in the argument function space mesh that is
   /// integrated over.
   /// - For cell integrals it has shape `(num_cells,)`.
   /// - For exterior/interior facet integrals, it has shape `(num_facts, 2)`
   /// (row-major storage), where `[i, 0]` is the index of a cell and
   /// `[i, 1]` is the local index of the facet relative to the cell.
-  std::span<const std::int32_t> domain_arg(IntegralType type, int rank, int id,
+  std::span<const std::int32_t> domain_arg(IntegralType type, int rank, int idx,
                                            int kernel_idx) const
   {
-    auto it = _edata.at(rank).find({type, id, kernel_idx});
+    auto it = _edata.at(rank).find({type, idx, kernel_idx});
     if (it == _edata.at(rank).end())
       throw std::runtime_error("Requested domain for argument not found.");
     try
@@ -533,19 +543,19 @@ class Form
   /// This method is equivalent to ::domain_arg, but returns mesh entity
   /// indices for coefficient \link Function Functions. \endlink
   ///
-  /// @param type Integral type.
-  /// @param id Integral identifier index.
-  /// @param c Coefficient index.
+  /// @param[in] type Integral type.
+  /// @param[in] idx Integral identifier.
+  /// @param[in] c Coefficient index.
   /// @return Entity indices in the coefficient function space mesh that
   /// is integrated over.
   /// - For cell integrals it has shape `(num_cells,)`.
   /// - For exterior/interior facet integrals, it has shape `(num_facts, 2)`
   /// (row-major storage), where `[i, 0]` is the index of a cell and
   /// `[i, 1]` is the local index of the facet relative to the cell.
-  std::span<const std::int32_t> domain_coeff(IntegralType type, int id,
+  std::span<const std::int32_t> domain_coeff(IntegralType type, int idx,
                                              int c) const
   {
-    auto it = _cdata.find({type, id, c});
+    auto it = _cdata.find({type, idx, c});
     if (it == _cdata.end())
       throw std::runtime_error("No domain for requested integral.");
     try

cpp/dolfinx/fem/assemble_matrix_impl.h

@@ -589,7 +589,7 @@ void assemble_matrix(
       cell_info1 = std::span(mesh1->topology()->get_cell_permutation_info());
     }
 
-    for (int i : a.integral_ids(IntegralType::cell))
+    for (int i = 0; i < a.num_integrals(IntegralType::cell, cell_type_idx); ++i)
     {
       auto fn = a.kernel(IntegralType::cell, i, cell_type_idx);
       assert(fn);
@@ -618,7 +618,8 @@ void assemble_matrix(
                          num_facets_per_cell);
     }
 
-    for (int i : a.integral_ids(IntegralType::exterior_facet))
+    for (int i = 0;
+         i < a.num_integrals(IntegralType::exterior_facet, cell_type_idx); ++i)
     {
       if (num_cell_types > 1)
       {
@@ -649,7 +650,8 @@ void assemble_matrix(
           cell_info0, cell_info1, perms);
     }
 
-    for (int i : a.integral_ids(IntegralType::interior_facet))
+    for (int i = 0;
+         i < a.num_integrals(IntegralType::interior_facet, cell_type_idx); ++i)
     {
       if (num_cell_types > 1)
       {

cpp/dolfinx/fem/assemble_scalar_impl.h

@@ -165,7 +165,7 @@ T assemble_scalar(
   assert(mesh);
 
   T value = 0;
-  for (int i : M.integral_ids(IntegralType::cell))
+  for (int i = 0; i < M.num_integrals(IntegralType::cell, 0); ++i)
   {
     auto fn = M.kernel(IntegralType::cell, i, 0);
     assert(fn);
@@ -190,7 +190,7 @@ T assemble_scalar(
                        num_facets_per_cell);
   }
 
-  for (int i : M.integral_ids(IntegralType::exterior_facet))
+  for (int i = 0; i < M.num_integrals(IntegralType::exterior_facet, 0); ++i)
   {
     auto fn = M.kernel(IntegralType::exterior_facet, i, 0);
     assert(fn);
@@ -208,7 +208,7 @@ T assemble_scalar(
         perms);
   }
 
-  for (int i : M.integral_ids(IntegralType::interior_facet))
+  for (int i = 0; i < M.num_integrals(IntegralType::interior_facet, 0); ++i)
   {
     auto fn = M.kernel(IntegralType::interior_facet, i, 0);
     assert(fn);

cpp/dolfinx/fem/assemble_vector_impl.h

@@ -1012,7 +1012,7 @@ void lift_bc(V&& b, const Form<T, U>& a, mdspan2_t x_dofmap,
       = element1->template dof_transformation_right_fn<T>(
           doftransform::transpose);
 
-  for (int i : a.integral_ids(IntegralType::cell))
+  for (int i = 0; i < a.num_integrals(IntegralType::cell, 0); ++i)
   {
     auto kernel = a.kernel(IntegralType::cell, i, 0);
     assert(kernel);
@@ -1057,7 +1057,7 @@ void lift_bc(V&& b, const Form<T, U>& a, mdspan2_t x_dofmap,
                        num_facets_per_cell);
   }
 
-  for (int i : a.integral_ids(IntegralType::exterior_facet))
+  for (int i = 0; i < a.num_integrals(IntegralType::exterior_facet, 0); ++i)
   {
     auto kernel = a.kernel(IntegralType::exterior_facet, i, 0);
     assert(kernel);
@@ -1081,7 +1081,7 @@ void lift_bc(V&& b, const Form<T, U>& a, mdspan2_t x_dofmap,
         cell_info1, bc_values1, bc_markers1, x0, alpha, perms);
   }
 
-  for (int i : a.integral_ids(IntegralType::interior_facet))
+  for (int i = 0; i < a.num_integrals(IntegralType::interior_facet, 0); ++i)
   {
     auto kernel = a.kernel(IntegralType::interior_facet, i, 0);
     assert(kernel);
@@ -1252,7 +1252,7 @@ void assemble_vector(
       cell_info0 = std::span(mesh0->topology()->get_cell_permutation_info());
     }
 
-    for (int i : L.integral_ids(IntegralType::cell))
+    for (int i = 0; i < L.num_integrals(IntegralType::cell, 0); ++i)
     {
       auto fn = L.kernel(IntegralType::cell, i, cell_type_idx);
       assert(fn);
@@ -1297,7 +1297,7 @@ void assemble_vector(
         = md::mdspan<const std::int32_t,
                      md::extents<std::size_t, md::dynamic_extent, 2>>;
 
-    for (int i : L.integral_ids(IntegralType::exterior_facet))
+    for (int i = 0; i < L.num_integrals(IntegralType::exterior_facet, 0); ++i)
     {
       auto fn = L.kernel(IntegralType::exterior_facet, i, 0);
       assert(fn);
@@ -1331,7 +1331,7 @@ void assemble_vector(
       }
     }
 
-    for (int i : L.integral_ids(IntegralType::interior_facet))
+    for (int i = 0; i < L.num_integrals(IntegralType::interior_facet, 0); ++i)
     {
       using mdspanx22_t
           = md::mdspan<const std::int32_t,

cpp/dolfinx/fem/pack.h

@@ -200,10 +200,10 @@ allocate_coefficient_storage(const Form<T, U>& form)
   std::map<std::pair<IntegralType, int>, std::pair<std::vector<T>, int>> coeffs;
   for (fem::IntegralType type : form.integral_types())
   {
-    for (int id : form.integral_ids(type))
+    for (int i = 0; i < form.num_integrals(type, 0); ++i)
     {
-      coeffs.emplace_hint(coeffs.end(), std::pair{type, id},
-                          allocate_coefficient_storage(form, type, id));
+      coeffs.emplace_hint(coeffs.end(), std::pair{type, i},
+                          allocate_coefficient_storage(form, type, i));
     }
   }