Indexing of Geometry cmaps and dofmaps for multiple cell types (#4227)

* Revert changes to cmap

* Revert more cmap

* Fix up Python layer including deprecation warning for

* Update interface to geometry::dofmaps

* clang-format

* clang-format

* clang-format...

* some fixes

* Fix in demo

* fixes

Author: chrisrichardson

cpp/demo/custom_kernel/main.cpp

@@ -141,9 +141,10 @@ double assemble_matrix1(const mesh::Geometry<T>& g, const fem::DofMap& dofmap,
   common::Timer timer("Assembler1 lambda (matrix)");
   md::mdspan<const T, md::extents<std::size_t, md::dynamic_extent, 3>> x(
       g.x().data(), g.x().size() / 3, 3);
-  fem::impl::assemble_cells_matrix<T>(
-      A.mat_add_values(), g.dofmap(), x, cells, {dofmap.map(), 1, cells}, ident,
-      {dofmap.map(), 1, cells}, ident, {}, {}, kernel, {}, {}, {}, {});
+  fem::impl::assemble_cells_matrix<T>(A.mat_add_values(), g.dofmaps().front(),
+                                      x, cells, {dofmap.map(), 1, cells}, ident,
+                                      {dofmap.map(), 1, cells}, ident, {}, {},
+                                      kernel, {}, {}, {}, {});
   A.scatter_rev();
   return A.squared_norm();
 }
@@ -168,8 +169,8 @@ double assemble_vector1(const mesh::Geometry<T>& g, const fem::DofMap& dofmap,
       g.x().data(), g.x().size() / 3, 3);
   common::Timer timer("Assembler1 lambda (vector)");
   fem::impl::assemble_cells<1>([](auto, auto, auto, auto) {}, b.array(),
-                               g.dofmap(), x, cells, {dofmap.map(), 1, cells},
-                               kernel, {}, {}, {});
+                               g.dofmaps().front(), x, cells,
+                               {dofmap.map(), 1, cells}, kernel, {}, {}, {});
   b.scatter_rev(std::plus<T>());
   return la::squared_norm(b);
 }

cpp/dolfinx/fem/Function.h

@@ -489,10 +489,11 @@ class Function
     auto map = mesh->topology()->index_map(tdim);
 
     // Get coordinate map
-    const CoordinateElement<geometry_type>& cmap = mesh->geometry().cmap();
+    const CoordinateElement<geometry_type>& cmap
+        = mesh->geometry().cmaps().front();
 
     // Get geometry data
-    auto x_dofmap = mesh->geometry().dofmap();
+    auto x_dofmap = mesh->geometry().dofmaps().front();
     const std::size_t num_dofs_g = cmap.dim();
     auto x_g = mesh->geometry().x();
 

cpp/dolfinx/fem/FunctionSpace.h

@@ -287,10 +287,11 @@ class FunctionSpace
       const auto [X, Xshape] = _elements[i]->interpolation_points();
 
       // Get coordinate map
-      const CoordinateElement<geometry_type>& cmap = _mesh->geometry().cmap(i);
+      const CoordinateElement<geometry_type>& cmap
+          = _mesh->geometry().cmaps().at(i);
 
       // Prepare cell geometry
-      auto x_dofmap = _mesh->geometry().dofmap(i);
+      auto x_dofmap = _mesh->geometry().dofmaps().at(i);
       const std::size_t num_dofs_g = cmap.dim();
       std::vector<geometry_type> coordinate_dofs_b(num_dofs_g * gdim);
       mdspan2_t coordinate_dofs(coordinate_dofs_b.data(), num_dofs_g, gdim);

cpp/dolfinx/fem/assemble_expression_impl.h

@@ -185,8 +185,8 @@ void tabulate_expression(
   }
 
   tabulate_expression<T, U>(values, fn, Xshape, value_size, num_argument_dofs,
-                            mesh.geometry().dofmap(), mesh.geometry().x(),
-                            coeffs, constants, entities, cell_info,
-                            post_dof_transform);
+                            mesh.geometry().dofmaps().front(),
+                            mesh.geometry().x(), coeffs, constants, entities,
+                            cell_info, post_dof_transform);
 }
 } // namespace dolfinx::fem::impl

cpp/dolfinx/fem/assemble_matrix_impl.h

@@ -642,7 +642,7 @@ void assemble_matrix(
   for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
   {
     // Geometry dofmap and data
-    mdspan2_t x_dofmap = mesh->geometry().dofmap(cell_type_idx);
+    mdspan2_t x_dofmap = mesh->geometry().dofmaps().at(cell_type_idx);
 
     // Get dofmap data
     std::shared_ptr<const fem::DofMap> dofmap0

cpp/dolfinx/fem/assemble_vector_impl.h

@@ -581,7 +581,7 @@ void assemble_vector(
   for (int cell_type_idx = 0; cell_type_idx < num_cell_types; ++cell_type_idx)
   {
     // Geometry dofmap and data
-    mdspan2_t x_dofmap = mesh->geometry().dofmap(cell_type_idx);
+    mdspan2_t x_dofmap = mesh->geometry().dofmaps().at(cell_type_idx);
 
     // Get dofmap data
     assert(L.function_spaces().at(0));

cpp/dolfinx/fem/assembler.h

@@ -175,7 +175,7 @@ T assemble_scalar(
   {
     // Geometry dofmap and data
     md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> x_dofmap
-        = mesh->geometry().dofmap(cell_type_idx);
+        = mesh->geometry().dofmaps().at(cell_type_idx);
     if constexpr (std::is_same_v<U, scalar_value_t<T>>)
     {
       val += impl::assemble_scalar(M, x_dofmap,

cpp/dolfinx/fem/discreteoperators.h

@@ -151,8 +151,8 @@ void discrete_curl(const FunctionSpace<T>& V0, const FunctionSpace<T>& V1,
   const auto [X, Xshape] = e1->interpolation_points();
 
   // Get/compute geometry map and evaluate at interpolation points
-  const CoordinateElement<T>& cmap = mesh->geometry().cmap();
-  auto x_dofmap = mesh->geometry().dofmap();
+  const CoordinateElement<T>& cmap = mesh->geometry().cmaps().front();
+  auto x_dofmap = mesh->geometry().dofmaps().front();
   const std::size_t num_dofs_g = cmap.dim();
   std::span<const T> x_g = mesh->geometry().x();
   std::array<std::size_t, 4> Phi_g_shape = cmap.tabulate_shape(1, Xshape[0]);
@@ -453,8 +453,8 @@ void interpolation_matrix(const FunctionSpace<U>& V0,
   const std::size_t value_size0 = V0.element()->reference_value_size();
 
   // Get geometry data
-  const CoordinateElement<U>& cmap = mesh->geometry().cmap();
-  auto x_dofmap = mesh->geometry().dofmap();
+  const CoordinateElement<U>& cmap = mesh->geometry().cmaps().front();
+  auto x_dofmap = mesh->geometry().dofmaps().front();
   const std::size_t num_dofs_g = cmap.dim();
   std::span<const U> x_g = mesh->geometry().x();
 

cpp/dolfinx/fem/interpolate.h

@@ -54,9 +54,9 @@ std::vector<T> interpolation_coords(const fem::FiniteElement<T>& element,
   // Find CoordinateElement appropriate to element
   auto cmap_index = [&geometry](mesh::CellType cell_type)
   {
-    for (std::size_t i = 0; i < geometry.num_maps(); ++i)
+    for (std::size_t i = 0; i < geometry.cmaps().size(); ++i)
     {
-      if (geometry.cmap(i).cell_shape() == cell_type)
+      if (geometry.cmaps().at(i).cell_shape() == cell_type)
         return i;
     }
     throw std::runtime_error("Cannot find CoordinateElement for FiniteElement");
@@ -65,10 +65,10 @@ std::vector<T> interpolation_coords(const fem::FiniteElement<T>& element,
 
   // Get geometry data and the element coordinate map
   const std::size_t gdim = geometry.dim();
-  auto x_dofmap = geometry.dofmap(index);
+  auto x_dofmap = geometry.dofmaps().at(index);
   std::span<const T> x_g = geometry.x();
 
-  const CoordinateElement<T>& cmap = geometry.cmap(index);
+  const CoordinateElement<T>& cmap = geometry.cmaps().at(index);
   const std::size_t num_dofs_g = cmap.dim();
 
   // Get the interpolation points on the reference cells
@@ -520,8 +520,8 @@ void interpolate_nonmatching_maps(Function<T, U>& u1,
   const std::size_t value_size_ref0 = element0->reference_value_size();
   const std::size_t value_size0 = V0->element()->reference_value_size();
 
-  const CoordinateElement<U>& cmap = mesh0->geometry().cmap();
-  auto x_dofmap = mesh0->geometry().dofmap();
+  const CoordinateElement<U>& cmap = mesh0->geometry().cmaps().front();
+  auto x_dofmap = mesh0->geometry().dofmaps().front();
   std::span<const U> x_g = mesh0->geometry().x();
 
   // (0) is derivative index, (1) is the point index, (2) is the basis
@@ -927,10 +927,10 @@ void piola_mapped_evaluation(const FiniteElement<U>& element, bool symmetric,
     throw std::runtime_error("Interpolation data has the wrong shape.");
 
   // Get coordinate map
-  const CoordinateElement<U>& cmap = mesh.geometry().cmap();
+  const CoordinateElement<U>& cmap = mesh.geometry().cmaps().front();
 
   // Get geometry data
-  auto x_dofmap = mesh.geometry().dofmap();
+  auto x_dofmap = mesh.geometry().dofmaps().front();
   const int num_dofs_g = cmap.dim();
   std::span<const U> x_g = mesh.geometry().x();
 

cpp/dolfinx/fem/utils.h

@@ -532,12 +532,12 @@ Form<T, U> create_form_factory(
       = [&geo = mesh->geometry()](const ufcx_integral& integral,
                                   std::size_t cell_idx)
   {
-    if (integral.coordinate_element_hash != geo.cmap(cell_idx).hash())
+    if (integral.coordinate_element_hash != geo.cmaps().at(cell_idx).hash())
     {
       throw std::runtime_error(
           "Generated integral geometry element does not match mesh geometry: "
           + std::to_string(integral.coordinate_element_hash) + ", "
-          + std::to_string(geo.cmap(cell_idx).hash()));
+          + std::to_string(geo.cmaps().at(cell_idx).hash()));
     }
   };
 
@@ -962,7 +962,7 @@ Expression<T, U> create_expression(
     assert(coefficients.front()->function_space());
     std::shared_ptr<const mesh::Mesh<U>> mesh
         = coefficients.front()->function_space()->mesh();
-    if (mesh->geometry().cmap().hash() != e.coordinate_element_hash)
+    if (mesh->geometry().cmaps().front().hash() != e.coordinate_element_hash)
     {
       throw std::runtime_error(
           "Expression and mesh geometric maps do not match.");
@@ -1082,7 +1082,7 @@ mesh::Mesh<T> interpolate_geometry(
         const graph::AdjacencyList<std::int32_t>&)>& reorder_fn = nullptr)
 {
   assert(mesh);
-  const CoordinateElement<T>& old_cmap = mesh->geometry().cmap();
+  const CoordinateElement<T>& old_cmap = mesh->geometry().cmaps().front();
   if (new_cmap.cell_shape() != old_cmap.cell_shape())
   {
     throw std::runtime_error(