Geometry access to dofmap(i) with mixed topology

cpp/dolfinx/mesh/Geometry.h

@@ -108,28 +108,22 @@ class Geometry
   int dim() const { return _dim; }
 
   /// @brief DofMap for the geometry.
+  /// @param i Index for the dofmap associated with the ith cell type.
   /// @return A 2D array with shape `(num_cells, dofs_per_cell)`.
-  md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>> dofmap() const
+  md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>
+  dofmap(std::optional<int> i = std::nullopt) const
   {
-    if (_dofmaps.size() != 1)
+    if (i.has_value())
+    {
+      std::size_t ndofs = _cmaps.at(*i).dim();
+      return md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>(
+          _dofmaps.at(*i).data(), _dofmaps.at(*i).size() / ndofs, ndofs);
+    }
+    else if (_dofmaps.size() != 1)
       throw std::runtime_error("Multiple dofmaps");
     return this->dofmap(0);
   }
 
-  /// @brief Degree-of-freedom map associated with the `i`th coordinate
-  /// map element in the geometry.
-  /// @param[in] i Index of the requested degree-of-freedom map. The
-  /// degree-of-freedom map corresponds to the geometry element
-  /// `cmaps()[i]`.
-  /// @return A dofmap array, with shape `(num_cells, dofs_per_cell)`.
-  md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>
-  dofmap(std::size_t i) const
-  {
-    std::size_t ndofs = _cmaps.at(i).dim();
-    return md::mdspan<const std::int32_t, md::dextents<std::size_t, 2>>(
-        _dofmaps.at(i).data(), _dofmaps.at(i).size() / ndofs, ndofs);
-  }
-
   /// @brief Index map for the geometry 'degrees-of-freedom'.
   /// @return Index map for the geometry dofs.
   std::shared_ptr<const common::IndexMap> index_map() const

python/demo/demo_mixed-topology.py

@@ -216,7 +216,7 @@ def marker(x):
 
 for j in range(2):
     vtk_topology = []
-    geom_dm = mesh.geometry.dofmaps(j)
+    geom_dm = mesh.geometry.dofmap(j)
     for c in geom_dm:
         vtk_topology += list(c[perm[j]])
     topology_type = topologies[j]

python/dolfinx/io/utils.py

@@ -324,7 +324,7 @@ def distribute_entity_data(
         mesh.geometry.input_global_indices,
         mesh.geometry.index_map().size_global,
         mesh.geometry.cmap().create_dof_layout(),
-        mesh.geometry.dofmap,
+        mesh.geometry.dofmap(),
         entity_dim,
         entities,
         values,

python/dolfinx/mesh.py

@@ -296,13 +296,12 @@ def dim(self):
         """Dimension of the Euclidean coordinate system."""
         return self._cpp_object.dim
 
-    @property
-    def dofmap(self) -> npt.NDArray[np.int32]:
-        """Dofmap for the geometry.
+    def dofmap(self, i=None) -> npt.NDArray[np.int32]:
+        """Dofmap for the geometry on the ith cell type.
 
         Shape is ``(num_cells, dofs_per_cell)``.
         """
-        return self._cpp_object.dofmap
+        return self._cpp_object.dofmap(i)
 
     def index_map(self) -> _IndexMap:
         """Index map for the geometry points (nodes) distribution."""

python/dolfinx/wrappers/dolfinx_wrappers/mesh.h

@@ -190,24 +190,24 @@ void declare_mesh(nb::module_& m, std::string type)
           nb::arg("x"), nb::arg("input_global_indices"))
       .def_prop_ro("dim", &dolfinx::mesh::Geometry<T>::dim,
                    "Geometric dimension")
-      .def_prop_ro(
+      .def(
           "dofmap",
           [](dolfinx::mesh::Geometry<T>& self)
           {
             auto dofs = self.dofmap();
             return nb::ndarray<const std::int32_t, nb::numpy>(
                 dofs.data_handle(), {dofs.extent(0), dofs.extent(1)});
           },
-          nb::rv_policy::reference_internal)
+          nb::rv_policy::reference_internal, "Get the geometry dofmap.")
       .def(
-          "dofmaps",
-          [](dolfinx::mesh::Geometry<T>& self, int i)
+          "dofmap",
+          [](dolfinx::mesh::Geometry<T>& self, std::optional<int> i)
           {
             auto dofs = self.dofmap(i);
             return nb::ndarray<const std::int32_t, nb::numpy>(
                 dofs.data_handle(), {dofs.extent(0), dofs.extent(1)});
           },
-          nb::rv_policy::reference_internal, nb::arg("i"),
+          nb::rv_policy::reference_internal, nb::arg("i").none(),
           "Get the geometry dofmap associated with coordinate element i (mixed "
           "topology)")
       .def("index_map", &dolfinx::mesh::Geometry<T>::index_map)

python/test/unit/fem/test_custom_assembler.py

@@ -144,7 +144,7 @@ def test_custom_mesh_loop_rank1(dtype):
 
     # Unpack mesh and dofmap data
     num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
-    x_dofs = mesh.geometry.dofmap
+    x_dofs = mesh.geometry.dofmap()
     x = mesh.geometry.x
     dofmap = V.dofmap.list
 

python/test/unit/fem/test_dof_permuting.py

@@ -157,7 +157,7 @@ def test_dof_positions(cell_type, space_type):
 
     # Get coordinates of dofs and edges and check that they are the same
     # for each global dof number
-    coord_dofs = mesh.geometry.dofmap
+    coord_dofs = mesh.geometry.dofmap()
     x_g = mesh.geometry.x
     cmap = mesh.geometry.cmap()
     tdim = mesh.topology.dim

python/test/unit/fem/test_dofmap.py

@@ -330,7 +330,7 @@ def test_higher_order_coordinate_map(points, celltype, order):
 
     V = functionspace(mesh, ("Lagrange", 2))
     X = V.element.interpolation_points
-    coord_dofs = mesh.geometry.dofmap
+    coord_dofs = mesh.geometry.dofmap()
     x_g = mesh.geometry.x
     cmap = mesh.geometry.cmap()
 
@@ -405,7 +405,7 @@ def test_higher_order_tetra_coordinate_map(order):
     mesh = create_mesh(MPI.COMM_WORLD, cells, domain, points)
     V = functionspace(mesh, ("Lagrange", order))
     X = V.element.interpolation_points
-    x_dofs = mesh.geometry.dofmap
+    x_dofs = mesh.geometry.dofmap()
     x_g = mesh.geometry.x
 
     x_coord_new = np.zeros([len(points), mesh.geometry.dim])

python/test/unit/fem/test_element_integrals.py

@@ -371,10 +371,10 @@ def test_plus_minus_simple_vector(cell_type, pm, dtype):
 
     # Check that the above vectors all have the same values as the first
     # one, but permuted due to differently ordered dofs
-    dofmap0 = spaces[0].mesh.geometry.dofmap
+    dofmap0 = spaces[0].mesh.geometry.dofmap()
     for result, space in zip(results[1:], spaces[1:]):
         # Get the data relating to two results
-        dofmap1 = space.mesh.geometry.dofmap
+        dofmap1 = space.mesh.geometry.dofmap()
 
         # For each cell
         for cell in range(2):
@@ -426,10 +426,10 @@ def test_plus_minus_vector(cell_type, pm1, pm2, dtype):
 
     # Check that the above vectors all have the same values as the first
     # one, but permuted due to differently ordered dofs
-    dofmap0 = spaces[0].mesh.geometry.dofmap
+    dofmap0 = spaces[0].mesh.geometry.dofmap()
     for result, space in zip(results[1:], spaces[1:]):
         # Get the data relating to two results
-        dofmap1 = space.mesh.geometry.dofmap
+        dofmap1 = space.mesh.geometry.dofmap()
 
         # For each cell
         for cell in range(2):
@@ -477,10 +477,10 @@ def test_plus_minus_matrix(cell_type, pm1, pm2, dtype):
 
     # Check that the above matrices all have the same values, but
     # permuted due to differently ordered dofs
-    dofmap0 = spaces[0].mesh.geometry.dofmap
+    dofmap0 = spaces[0].mesh.geometry.dofmap()
     for result, space in zip(results[1:], spaces[1:]):
         # Get the data relating to two results
-        dofmap1 = space.mesh.geometry.dofmap
+        dofmap1 = space.mesh.geometry.dofmap()
         dof_order = []
 
         # For each cell

python/test/unit/fem/test_expression.py

@@ -314,7 +314,7 @@ def e_exact(x):
     # # FIXME: Below is only for testing purposes,
     # # never to be used in user code!
     # # TODO: Replace when interpolation into Quadrature element works.
-    coord_dofs = mesh.geometry.dofmap
+    coord_dofs = mesh.geometry.dofmap()
     x_g = mesh.geometry.x
     tdim = mesh.topology.dim
     Q_dofs = Q.dofmap.list

python/test/unit/fem/test_petsc_custom_assembler.py

@@ -141,7 +141,7 @@ def test_custom_mesh_loop_petsc_rank2(set_vals, backend):
 
     # Unpack mesh and dofmap data
     num_owned_cells = mesh.topology.index_map(mesh.topology.dim).size_local
-    x_dofs = mesh.geometry.dofmap
+    x_dofs = mesh.geometry.dofmap()
     x = mesh.geometry.x
     dofmap = V.dofmap.list.astype(np.dtype(PETSc.IntType))
 

python/test/unit/mesh/test_mesh.py

@@ -98,7 +98,7 @@ def submesh_geometry_test(mesh, submesh, entity_map, geom_map, entity_dim):
         mesh.topology.create_entity_permutations()
         e_to_g = entities_to_geometry(mesh, entity_dim, np.array(submesh_to_mesh), True)
         for submesh_entity in range(len(submesh_to_mesh)):
-            submesh_x_dofs = submesh.geometry.dofmap[submesh_entity]
+            submesh_x_dofs = submesh.geometry.dofmap()[submesh_entity]
             # e_to_g[i] gets the mesh x_dofs of entities[i], which should
             # correspond to the x_dofs of cell i in the submesh
             mesh_x_dofs = e_to_g[submesh_entity]

python/test/unit/mesh/test_mixed_topology.py

@@ -89,8 +89,8 @@ def test_mixed_topology_mesh():
     )
     print(geom.x)
     print(geom.index_map().size_local)
-    print(geom.dofmaps(0))
-    print(geom.dofmaps(1))
+    print(geom.dofmap(0))
+    print(geom.dofmap(1))
     set_log_level(LogLevel.WARNING)
 
 
@@ -249,8 +249,8 @@ def test_parallel_mixed_mesh():
         topology, [tri._cpp_object, quad._cpp_object], nodes, xdofs, x.flatten(), 2
     )
 
-    assert len(geom.dofmaps(0)) == 2
-    assert len(geom.dofmaps(1)) == 1
+    assert len(geom.dofmap(0)) == 2
+    assert len(geom.dofmap(1)) == 1
 
     mesh = Mesh_float64(MPI.COMM_WORLD, topology, geom)
     tri = mesh.topology.connectivity((2, 0), (0, 0))