Add interface to mesh::Geometry for multiple cmaps and dofmaps (#3035)

* Update interface to dofmapbuilder

* tidying up in Geometry.h

* Improve documentation

* Improve documentation

* Apply suggestions from code review

Co-authored-by: Joe Dean <[email protected]>

* Change signature

* Update wrapper

* Use front rather than [0].

Makes clearer where updates will be needed.

* More updates

* More updates

---------

Co-authored-by: Joe Dean <[email protected]>
Co-authored-by: Garth N. Wells <[email protected]>

cpp/dolfinx/fem/DofMap.cpp

@@ -225,11 +225,11 @@ std::pair<DofMap, std::vector<std::int32_t>> DofMap::collapse(
       // Create new element dof layout and reset parent
       ElementDofLayout collapsed_dof_layout = layout.copy();
 
-      auto [_index_map, bs, dofmap]
-          = build_dofmap_data(comm, topology, collapsed_dof_layout, reorder_fn);
+      auto [_index_map, bs, dofmaps] = build_dofmap_data(
+          comm, topology, {collapsed_dof_layout}, reorder_fn);
       auto index_map
           = std::make_shared<common::IndexMap>(std::move(_index_map));
-      return DofMap(layout, index_map, bs, std::move(dofmap), bs);
+      return DofMap(layout, index_map, bs, std::move(dofmaps.front()), bs);
     }
     else
     {

cpp/dolfinx/fem/dofmapbuilder.cpp

@@ -541,10 +541,10 @@ std::pair<std::vector<std::int64_t>, std::vector<int>> get_global_indices(
 } // namespace
 
 //-----------------------------------------------------------------------------
-std::tuple<common::IndexMap, int, std::vector<std::int32_t>>
+std::tuple<common::IndexMap, int, std::vector<std::vector<std::int32_t>>>
 fem::build_dofmap_data(
     MPI_Comm comm, const mesh::Topology& topology,
-    const ElementDofLayout& element_dof_layout,
+    const std::vector<ElementDofLayout>& element_dof_layouts,
     const std::function<std::vector<int>(
         const graph::AdjacencyList<std::int32_t>&)>& reorder_fn)
 {
@@ -557,12 +557,12 @@ fem::build_dofmap_data(
   // pair {dimension, mesh entity index} giving the mesh entity that dof
   // i is associated with.
   const auto [node_graph0, local_to_global0, dof_entity0]
-      = build_basic_dofmap(topology, element_dof_layout);
+      = build_basic_dofmap(topology, element_dof_layouts.front());
 
   std::vector<std::shared_ptr<const common::IndexMap>> index_maps(D + 1);
   for (int d = 0; d <= D; ++d)
   {
-    if (element_dof_layout.num_entity_dofs(d) > 0)
+    if (element_dof_layouts.front().num_entity_dofs(d) > 0)
     {
       assert(topology.index_map(d));
       index_maps[d] = topology.index_map(d);
@@ -578,10 +578,10 @@ fem::build_dofmap_data(
   std::int64_t offset = 0;
   for (int d = 0; d <= D; ++d)
   {
-    if (element_dof_layout.num_entity_dofs(d) > 0)
+    if (element_dof_layouts.front().num_entity_dofs(d) > 0)
     {
       offset += index_maps[d]->local_range()[0]
-                * element_dof_layout.num_entity_dofs(d);
+                * element_dof_layouts.front().num_entity_dofs(d);
     }
   }
 
@@ -596,11 +596,12 @@ fem::build_dofmap_data(
                              local_to_global_owner);
 
   // Build re-ordered dofmap
-  std::vector<std::int32_t> dofmap(node_graph0.array.size());
-  for (std::size_t i = 0; i < dofmap.size(); ++i)
-    dofmap[i] = old_to_new[node_graph0.array[i]];
+  std::vector<std::vector<std::int32_t>> dofmaps(1);
+  dofmaps.front().resize(node_graph0.array.size());
+  for (std::size_t i = 0; i < dofmaps.front().size(); ++i)
+    dofmaps.front()[i] = old_to_new[node_graph0.array[i]];
 
-  return {std::move(index_map), element_dof_layout.block_size(),
-          std::move(dofmap)};
+  return {std::move(index_map), element_dof_layouts.front().block_size(),
+          std::move(dofmaps)};
 }
 //-----------------------------------------------------------------------------

cpp/dolfinx/fem/dofmapbuilder.h

@@ -35,13 +35,14 @@ class ElementDofLayout;
 /// Build dofmap data for elements on a mesh topology
 /// @param[in] comm MPI communicator
 /// @param[in] topology The mesh topology
-/// @param[in] element_dof_layout The element dof layout
+/// @param[in] element_dof_layouts The element dof layouts for each cell type in
+/// @p topology
 /// @param[in] reorder_fn Graph reordering function that is applied to
 /// the dofmaps
-/// @return The index map, block size, and dofmap for element type 0
-std::tuple<common::IndexMap, int, std::vector<std::int32_t>>
+/// @return The index map, block size, and dofmaps for each element type
+std::tuple<common::IndexMap, int, std::vector<std::vector<std::int32_t>>>
 build_dofmap_data(MPI_Comm comm, const mesh::Topology& topology,
-                  const ElementDofLayout& element_dof_layout,
+                  const std::vector<ElementDofLayout>& element_dof_layouts,
                   const std::function<std::vector<int>(
                       const graph::AdjacencyList<std::int32_t>&)>& reorder_fn);
 

cpp/dolfinx/fem/utils.cpp

@@ -110,8 +110,8 @@ fem::DofMap fem::create_dofmap(
       topology.create_entities(d);
   }
 
-  auto [_index_map, bs, dofmap]
-      = build_dofmap_data(comm, topology, layout, reorder_fn);
+  auto [_index_map, bs, dofmaps]
+      = build_dofmap_data(comm, topology, {layout}, reorder_fn);
   auto index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
 
   // If the element's DOF transformations are permutations, permute the
@@ -126,12 +126,12 @@ fem::DofMap fem::create_dofmap(
     int dim = layout.num_dofs();
     for (std::int32_t cell = 0; cell < num_cells; ++cell)
     {
-      std::span<std::int32_t> dofs(dofmap.data() + cell * dim, dim);
+      std::span<std::int32_t> dofs(dofmaps.front().data() + cell * dim, dim);
       unpermute_dofs(dofs, cell_info[cell]);
     }
   }
 
-  return DofMap(layout, index_map, bs, std::move(dofmap), bs);
+  return DofMap(layout, index_map, bs, std::move(dofmaps.front()), bs);
 }
 //-----------------------------------------------------------------------------
 std::vector<std::string> fem::get_coefficient_names(const ufcx_form& ufcx_form)