Indexing of Geometry cmaps and dofmaps for multiple cell types

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(

cpp/dolfinx/geometry/utils.h

@@ -57,7 +57,7 @@ std::vector<T> shortest_vector(const mesh::Mesh<T>& mesh, int dim,
   const mesh::Geometry<T>& geometry = mesh.geometry();
 
   std::span<const T> geom_dofs = geometry.x();
-  auto x_dofmap = geometry.dofmap();
+  auto x_dofmap = geometry.dofmaps().front();
   std::vector<T> shortest_vectors;
   shortest_vectors.reserve(3 * entities.size());
   if (dim == tdim)
@@ -106,7 +106,7 @@ std::vector<T> shortest_vector(const mesh::Mesh<T>& mesh, int dim,
       // Tabulate geometry dofs for the entity
       auto dofs = md::submdspan(x_dofmap, c, md::full_extent);
       const std::vector<int> entity_dofs
-          = geometry.cmap().create_dof_layout().entity_closure_dofs(
+          = geometry.cmaps().front().create_dof_layout().entity_closure_dofs(
               dim, local_cell_entity);
       std::vector<T> nodes(3 * entity_dofs.size());
       for (std::size_t i = 0; i < entity_dofs.size(); i++)
@@ -528,7 +528,7 @@ std::int32_t compute_first_colliding_cell(const mesh::Mesh<T>& mesh,
   {
     const mesh::Geometry<T>& geometry = mesh.geometry();
     std::span<const T> geom_dofs = geometry.x();
-    auto x_dofmap = geometry.dofmap();
+    auto x_dofmap = geometry.dofmaps().front();
     const std::size_t num_nodes = x_dofmap.extent(1);
     std::vector<T> coordinate_dofs(num_nodes * 3);
     for (auto cell : cells)
@@ -779,7 +779,7 @@ determine_point_ownership(const mesh::Mesh<T>& mesh, std::span<const T> points,
   // Get mesh geometry for closest entity
   const mesh::Geometry<T>& geometry = mesh.geometry();
   std::span<const T> geom_dofs = geometry.x();
-  auto x_dofmap = geometry.dofmap();
+  auto x_dofmap = geometry.dofmaps().front();
 
   // Compute candidate cells for collisions (and extrapolation)
   const graph::AdjacencyList<std::int32_t> candidate_collisions

cpp/dolfinx/io/ADIOS2Writers.h

@@ -313,8 +313,8 @@ void vtx_write_mesh(adios2::IO& io, adios2::Engine& engine,
       io, "NumberOfNodes", {adios2::LocalValueDim});
   engine.Put<std::uint32_t>(vertices, num_vertices);
 
-  auto [vtkcells, shape]
-      = io::extract_vtk_connectivity(geometry.dofmap(), topology->cell_type());
+  auto [vtkcells, shape] = io::extract_vtk_connectivity(
+      geometry.dofmaps().front(), topology->cell_type());
 
   // Add cell metadata
   int tdim = topology->dim();

cpp/dolfinx/io/VTKFile.cpp

@@ -408,7 +408,7 @@ void write_function(
   {
     std::vector<std::int64_t> tmp;
     std::tie(tmp, cshape) = io::extract_vtk_connectivity(
-        mesh0->geometry().dofmap(), topology0->cell_type());
+        mesh0->geometry().dofmaps().front(), topology0->cell_type());
     cells.assign(tmp.begin(), tmp.end());
     const mesh::Geometry<U>& geometry = mesh0->geometry();
     x.assign(geometry.x().begin(), geometry.x().end());
@@ -800,7 +800,7 @@ void io::VTKFile::write(const mesh::Mesh<U>& mesh, double time)
 
   // Add mesh data to "Piece" node
   const auto [cells, cshape]
-      = extract_vtk_connectivity(mesh.geometry().dofmap(), cell_type);
+      = extract_vtk_connectivity(mesh.geometry().dofmaps().front(), cell_type);
   std::array<std::size_t, 2> xshape = {geometry.x().size() / 3, 3};
   std::vector<std::uint8_t> x_ghost(xshape[0], 0);
   std::fill(std::next(x_ghost.begin(), xmap->size_local()), x_ghost.end(), 1);

cpp/dolfinx/io/VTKHDF.h

@@ -70,7 +70,7 @@ void write_mesh(const std::filesystem::path& filename,
   for (std::size_t i = 0; i < cell_index_maps.size(); ++i)
   {
     md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> g_dofmap
-        = mesh.geometry().dofmap(i);
+        = mesh.geometry().dofmaps().at(i);
 
     std::vector<std::uint16_t> perm
         = cells::perm_vtk(cell_types[i], g_dofmap.extent(1));
@@ -102,7 +102,7 @@ void write_mesh(const std::filesystem::path& filename,
   std::vector<std::int64_t> cell_stop_pos;
   for (std::size_t i = 0; i < cell_index_maps.size(); ++i)
   {
-    num_nodes_per_cell.push_back(mesh.geometry().cmap(i).dim());
+    num_nodes_per_cell.push_back(mesh.geometry().cmaps().at(i).dim());
     std::array<std::int64_t, 2> r = cell_index_maps[i]->local_range();
     cell_start_pos.push_back(r[0]);
     cell_stop_pos.push_back(r[1]);

cpp/dolfinx/io/XDMFFile.cpp

@@ -390,8 +390,9 @@ XDMFFile::read_meshtags(const mesh::Mesh<double>& mesh, const std::string& name,
       entities_values = io::distribute_entity_data<std::int32_t>(
           *mesh.topology(), mesh.geometry().input_global_indices(),
           mesh.geometry().index_map()->size_global(),
-          mesh.geometry().cmap().create_dof_layout(), mesh.geometry().dofmap(),
-          mesh::cell_dim(cell_type), entities_span, values);
+          mesh.geometry().cmaps().front().create_dof_layout(),
+          mesh.geometry().dofmaps().front(), mesh::cell_dim(cell_type),
+          entities_span, values);
 
   spdlog::info("XDMF create meshtags");
   std::size_t num_vertices_per_entity = mesh::cell_num_entities(

cpp/dolfinx/io/vtk_utils.h

@@ -86,10 +86,10 @@ tabulate_lagrange_dof_coordinates(const fem::FunctionSpace<T>& V)
   // Get the dof coordinates on the reference element and the  mesh
   // coordinate map
   const auto [X, Xshape] = element->interpolation_points();
-  const fem::CoordinateElement<T>& cmap = mesh->geometry().cmap();
+  const fem::CoordinateElement<T>& cmap = mesh->geometry().cmaps().front();
 
   // Prepare cell geometry
-  auto dofmap_x = mesh->geometry().dofmap();
+  auto dofmap_x = mesh->geometry().dofmaps().front();
   std::span<const T> x_g = mesh->geometry().x();
   const std::size_t num_dofs_g = cmap.dim();
 

cpp/dolfinx/io/xdmf_function.cpp

@@ -111,7 +111,7 @@ void xdmf_function::add_function(MPI_Comm comm, const fem::Function<T, U>& u,
   {
     // Get number of geometry nodes per cell
     auto& geometry = mesh->geometry();
-    auto& cmap = geometry.cmap();
+    auto& cmap = geometry.cmaps().front();
     int cmap_dim = cmap.dim();
     int cell_dim = element->space_dimension() / element->block_size();
     if (cmap_dim != cell_dim)
@@ -138,7 +138,7 @@ void xdmf_function::add_function(MPI_Comm comm, const fem::Function<T, U>& u,
     std::int32_t num_local_points = map_x->size_local();
 
     // Get dof array and pack into array (padded where appropriate)
-    auto dofmap_x = geometry.dofmap();
+    auto dofmap_x = geometry.dofmaps().front();
     data_values.resize(num_local_points * num_components, 0);
     for (std::int32_t c = 0; c < num_cells; ++c)
     {

cpp/dolfinx/io/xdmf_mesh.cpp

@@ -43,7 +43,7 @@ void xdmf_mesh::add_topology_data(MPI_Comm comm, pugi::xml_node& xml_node,
       = mesh::cell_entity_type(cell_type, dim, 0);
 
   const fem::ElementDofLayout cmap_dof_layout
-      = geometry.cmap().create_dof_layout();
+      = geometry.cmaps().front().create_dof_layout();
 
   // Get number of nodes per entity
   const int num_nodes_per_entity
@@ -61,7 +61,7 @@ void xdmf_mesh::add_topology_data(MPI_Comm comm, pugi::xml_node& xml_node,
   // Pack topology data
   std::vector<std::int64_t> topology_data;
 
-  auto x_dofmap = geometry.dofmap();
+  auto x_dofmap = geometry.dofmaps().front();
   auto map_g = geometry.index_map();
   assert(map_g);
   const std::int64_t offset_g = map_g->local_range()[0];