Add multiple dofmaps and elements to Geometry

cpp/dolfinx/fem/utils.cpp

@@ -117,7 +117,6 @@ fem::DofMap fem::create_dofmap(
   // DOF numbering on each cell
   if (unpermute_dofs)
   {
-    const int D = topology.dim();
     const int num_cells = topology.connectivity(D, 0)->num_nodes();
     topology.create_entity_permutations();
     const std::vector<std::uint32_t>& cell_info
@@ -133,6 +132,55 @@ fem::DofMap fem::create_dofmap(
   return DofMap(layout, index_map, bs, std::move(dofmaps.front()), bs);
 }
 //-----------------------------------------------------------------------------
+std::vector<fem::DofMap> fem::create_dofmaps(
+    MPI_Comm comm, const std::vector<ElementDofLayout>& layouts,
+    mesh::Topology& topology,
+    std::function<void(std::span<std::int32_t>, std::uint32_t)> unpermute_dofs,
+    std::function<std::vector<int>(const graph::AdjacencyList<std::int32_t>&)>
+        reorder_fn)
+{
+  std::int32_t D = topology.dim();
+  assert(layouts.size() == topology.entity_types(D).size());
+
+  // Create required mesh entities
+  for (std::int32_t d = 0; d < D; ++d)
+  {
+    if (layouts.front().num_entity_dofs(d) > 0)
+      topology.create_entities(d);
+  }
+
+  auto [_index_map, bs, dofmaps]
+      = build_dofmap_data(comm, topology, layouts, reorder_fn);
+  auto index_map = std::make_shared<common::IndexMap>(std::move(_index_map));
+
+  // If the element's DOF transformations are permutations, permute the
+  // DOF numbering on each cell
+  if (unpermute_dofs)
+  {
+    if (layouts.size() != 1)
+    {
+      throw std::runtime_error(
+          "DOF transformations not yet supported in mixed topology.");
+    }
+    std::int32_t num_cells = topology.connectivity(D, 0)->num_nodes();
+    topology.create_entity_permutations();
+    const std::vector<std::uint32_t>& cell_info
+        = topology.get_cell_permutation_info();
+    std::int32_t dim = layouts.front().num_dofs();
+    for (std::int32_t cell = 0; cell < num_cells; ++cell)
+    {
+      std::span<std::int32_t> dofs(dofmaps.front().data() + cell * dim, dim);
+      unpermute_dofs(dofs, cell_info[cell]);
+    }
+  }
+
+  std::vector<DofMap> dms;
+  for (std::size_t i = 0; i < dofmaps.size(); ++i)
+    dms.emplace_back(layouts[i], index_map, bs, std::move(dofmaps[i]), bs);
+
+  return dms;
+}
+//-----------------------------------------------------------------------------
 std::vector<std::string> fem::get_coefficient_names(const ufcx_form& ufcx_form)
 {
   return std::vector<std::string>(ufcx_form.coefficient_name_map,

cpp/dolfinx/fem/utils.h

@@ -259,6 +259,23 @@ DofMap create_dofmap(
     std::function<std::vector<int>(const graph::AdjacencyList<std::int32_t>&)>
         reorder_fn);
 
+/// @brief Create a set of dofmaps on a given topology
+/// @param[in] comm MPI communicator
+/// @param[in] layouts Dof layout on each element type
+/// @param[in] topology Mesh topology
+/// @param[in] unpermute_dofs Function to un-permute dofs. `nullptr`
+/// when transformation is not required.
+/// @param[in] reorder_fn Graph reordering function called on the dofmaps
+/// @return The list of new dof maps
+/// @note The number of layouts must match the number of cell types in the
+/// topology
+std::vector<DofMap> create_dofmaps(
+    MPI_Comm comm, const std::vector<ElementDofLayout>& layouts,
+    mesh::Topology& topology,
+    std::function<void(std::span<std::int32_t>, std::uint32_t)> unpermute_dofs,
+    std::function<std::vector<int>(const graph::AdjacencyList<std::int32_t>&)>
+        reorder_fn);
+
 /// Get the name of each coefficient in a UFC form
 /// @param[in] ufcx_form The UFC form
 /// @return The name of each coefficient

cpp/dolfinx/mesh/Geometry.h

@@ -343,7 +343,7 @@ create_geometry(
   std::vector<T> xg(3 * shape0, 0);
   for (std::size_t i = 0; i < shape0; ++i)
   {
-    std::copy_n(std::next(x.cbegin(), shape1 * l2l[i]), shape1,
+    std::copy_n(std::next(x.begin(), shape1 * l2l[i]), shape1,
                 std::next(xg.begin(), 3 * i));
   }
 

python/dolfinx/wrappers/fem.cpp

@@ -849,6 +849,30 @@ void declare_real_functions(nb::module_& m)
       nb::arg("element"),
       "Create DofMap object from a pointer to ufcx_dofmap.");
 
+  m.def(
+      "create_dofmaps",
+      [](const dolfinx_wrappers::MPICommWrapper comm,
+         std::vector<std::uintptr_t> ufcx_dofmaps,
+         dolfinx::mesh::Topology& topology)
+      {
+        std::vector<dolfinx::fem::ElementDofLayout> layouts;
+        int D = topology.dim();
+        assert(ufcx_dofmaps.size() == topology.entity_types(D).size());
+        for (std::size_t i = 0; i < ufcx_dofmaps.size(); ++i)
+        {
+          ufcx_dofmap* p = reinterpret_cast<ufcx_dofmap*>(ufcx_dofmaps[i]);
+          assert(p);
+          layouts.push_back(dolfinx::fem::create_element_dof_layout(
+              *p, topology.entity_types(D)[i]));
+        }
+
+        return dolfinx::fem::create_dofmaps(comm.get(), layouts, topology,
+                                            nullptr, nullptr);
+      },
+      nb::arg("comm"), nb::arg("dofmap"), nb::arg("topology"),
+      "Create DofMap objects on a mixed topology mesh from pointers to "
+      "ufcx_dofmaps.");
+
   m.def(
       "locate_dofs_topological",
       [](const std::vector<

python/dolfinx/wrappers/mesh.cpp

@@ -164,6 +164,17 @@ void declare_mesh(nb::module_& m, std::string type)
                 dofs.data_handle(), {dofs.extent(0), dofs.extent(1)});
           },
           nb::rv_policy::reference_internal)
+      .def(
+          "dofmaps",
+          [](dolfinx::mesh::Geometry<T>& self, 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"),
+          "Get the geometry dofmap associated with coordinate element i (mixed "
+          "topology)")
       .def("index_map", &dolfinx::mesh::Geometry<T>::index_map)
       .def_prop_ro(
           "x",
@@ -384,6 +395,20 @@ void declare_mesh(nb::module_& m, std::string type)
         return as_nbarray(std::move(idx), {entities.size(), num_vertices});
       },
       nb::arg("mesh"), nb::arg("dim"), nb::arg("entities"), nb::arg("orient"));
+
+  m.def("create_geometry",
+        [](const dolfinx::mesh::Topology& topology,
+           const std::vector<dolfinx::fem::CoordinateElement<T>>& elements,
+           nb::ndarray<const std::int64_t, nb::ndim<1>, nb::c_contig> nodes,
+           nb::ndarray<const std::int64_t, nb::ndim<1>, nb::c_contig> xdofs,
+           nb::ndarray<const T, nb::ndim<1>, nb::c_contig> x, int dim)
+        {
+          return dolfinx::mesh::create_geometry(
+              topology, elements,
+              std::span<const std::int64_t>(nodes.data(), nodes.size()),
+              std::span<const std::int64_t>(xdofs.data(), xdofs.size()),
+              std::span<const T>(x.data(), x.size()), dim);
+        });
 }
 
 void mesh(nb::module_& m)

python/test/unit/fem/test_mixed_mesh_dofmap.py

@@ -0,0 +1,145 @@
+from mpi4py import MPI
+
+import numpy as np
+
+import basix
+import dolfinx.cpp as _cpp
+from dolfinx import jit
+from dolfinx.cpp.mesh import Mesh_float64, create_geometry, create_topology
+from dolfinx.fem import coordinate_element
+from dolfinx.fem.dofmap import DofMap
+from dolfinx.log import LogLevel, set_log_level
+from dolfinx.mesh import CellType
+
+
+def create_element_dofmap(mesh, cell_types, degree):
+    cpp_elements = []
+    dofmaps = []
+    for cell_type in cell_types:
+        ufl_e = basix.ufl.element("P", cell_type, degree)
+        form_compiler_options = {"scalar_type": np.float64}
+        (ufcx_element, ufcx_dofmap), module, code = jit.ffcx_jit(
+            mesh.comm, ufl_e, form_compiler_options=form_compiler_options
+        )
+        ffi = module.ffi
+        cpp_elements += [
+            _cpp.fem.FiniteElement_float64(ffi.cast("uintptr_t", ffi.addressof(ufcx_element)))
+        ]
+        dofmaps += [ufcx_dofmap]
+
+    cpp_dofmaps = _cpp.fem.create_dofmaps(
+        mesh.comm, [ffi.cast("uintptr_t", ffi.addressof(dm)) for dm in dofmaps], mesh.topology
+    )
+
+    return (cpp_elements, cpp_dofmaps)
+
+
+def test_dofmap_mixed_topology():
+    rank = MPI.COMM_WORLD.Get_rank()
+
+    # Two triangles and one quadrilateral
+    tri = [0, 1, 4, 0, 3, 4]
+    quad = [1, 4, 2, 5]
+    # cells with global indexing
+    cells = [[t + 3 * rank for t in tri], [q + 3 * rank for q in quad]]
+    orig_index = [[3 * rank, 1 + 3 * rank], [2 + 3 * rank]]
+    # No ghosting
+    ghost_owners = [[], []]
+    # All vertices are on boundary
+    boundary_vertices = [3 * rank + i for i in range(6)]
+
+    topology = create_topology(
+        MPI.COMM_WORLD,
+        [CellType.triangle, CellType.quadrilateral],
+        cells,
+        orig_index,
+        ghost_owners,
+        boundary_vertices,
+    )
+    # Create dofmaps for Geometry
+    tri = coordinate_element(CellType.triangle, 1)
+    quad = coordinate_element(CellType.quadrilateral, 1)
+    nodes = np.arange(6, dtype=np.int64) + 3 * rank
+    xdofs = np.array([0, 1, 4, 0, 3, 4, 1, 4, 2, 5], dtype=np.int64) + 3 * rank
+    x = np.array(
+        [[0.0, 0.0], [1.0, 0.0], [2.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, 1.0]], dtype=np.float64
+    )
+    x[:, 1] += 1.0 * rank
+
+    set_log_level(LogLevel.INFO)
+    geom = create_geometry(
+        topology, [tri._cpp_object, quad._cpp_object], nodes, xdofs, x.flatten(), 2
+    )
+    mesh = Mesh_float64(MPI.COMM_WORLD, topology, geom)
+
+    assert mesh.geometry.x.shape == (6, 3)
+
+    # Second order dofmap on mixed mesh
+    elements, dofmaps = create_element_dofmap(
+        mesh, [basix.CellType.triangle, basix.CellType.quadrilateral], 2
+    )
+
+    assert len(elements) == 2
+    assert elements[0].basix_element.cell_type.name == "triangle"
+    assert elements[1].basix_element.cell_type.name == "quadrilateral"
+
+    assert len(dofmaps) == 2
+    q0 = DofMap(dofmaps[0])
+    q1 = DofMap(dofmaps[1])
+    assert q0.index_map.size_local == q1.index_map.size_local
+    # Triangles
+    print(q0.list)
+    assert q0.list.shape == (2, 6)
+    assert len(q0.dof_layout.entity_dofs(2, 0)) == 0
+    # Quadrilaterals
+    print(q1.list)
+    assert q1.list.shape == (1, 9)
+    assert len(q1.dof_layout.entity_dofs(2, 0)) == 1
+
+
+def test_dofmap_prism_mesh():
+    # Prism mesh
+    cells = [[0, 1, 2, 3, 4, 5]]
+    # cells with global indexing
+    orig_index = [[0, 1, 2, 3, 4, 5]]
+    # No ghosting
+    ghost_owners = [[]]
+    # All vertices are on boundary
+    boundary_vertices = [0, 1, 2, 3, 4, 5]
+
+    topology = create_topology(
+        MPI.COMM_SELF, [CellType.prism], cells, orig_index, ghost_owners, boundary_vertices
+    )
+    topology.create_entities(2)
+
+    # Create dofmaps for Geometry
+    prism = coordinate_element(CellType.prism, 1)
+    nodes = np.array([0, 1, 2, 3, 4, 5], dtype=np.int64)
+    xdofs = np.array([0, 1, 2, 3, 4, 5], dtype=np.int64)
+    x = np.array(
+        [
+            [0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [0.0, 1.0, 1.0],
+        ],
+        dtype=np.float64,
+    )
+
+    set_log_level(LogLevel.INFO)
+    geom = create_geometry(topology, [prism._cpp_object], nodes, xdofs, x.flatten(), 3)
+    mesh = Mesh_float64(MPI.COMM_WORLD, topology, geom)
+
+    elements, dofmaps = create_element_dofmap(mesh, [basix.CellType.prism], 2)
+    print()
+    assert len(elements) == 1
+    assert len(dofmaps) == 1
+    q = DofMap(dofmaps[0])
+    assert q.index_map.size_local == 18
+    print(q.list)
+    facet_dofs = []
+    for j in range(5):
+        facet_dofs += q.dof_layout.entity_dofs(2, j)
+    assert len(facet_dofs) == 3