Add vertex_to_dofmap + test

CMakeLists.txt

@@ -32,8 +32,6 @@ if(Basix_FOUND)
   message(STATUS "Found Basix at ${Basix_DIR}")
 endif()
 
-
-
 find_package(DOLFINX REQUIRED CONFIG)
 
 if(DOLFINX_FOUND)

src/scifem.cpp

@@ -10,11 +10,46 @@
 #include <dolfinx/mesh/Mesh.h>
 #include <memory>
 #include <nanobind/nanobind.h>
+#include <nanobind/ndarray.h>
 #include <nanobind/stl/shared_ptr.h>
 #include <nanobind/stl/tuple.h>
 #include <nanobind/stl/vector.h>
 
-using namespace dolfinx;
+namespace
+{
+// Copyright (c) 2020-2024 Chris Richardson, Matthew Scroggs and Garth N. Wells
+// FEniCS Project (Basix)
+// SPDX-License-Identifier:    MIT
+
+template <typename V>
+auto as_nbarray(V&& x, std::size_t ndim, const std::size_t* shape)
+{
+  using _V = std::decay_t<V>;
+  _V* ptr = new _V(std::move(x));
+  return nanobind::ndarray<typename _V::value_type, nanobind::numpy>(
+      ptr->data(), ndim, shape,
+      nanobind::capsule(ptr, [](void* p) noexcept { delete (_V*)p; }));
+}
+
+template <typename V>
+auto as_nbarray(V&& x, const std::initializer_list<std::size_t> shape)
+{
+  return as_nbarray(x, shape.size(), shape.begin());
+}
+
+template <typename V>
+auto as_nbarray(V&& x)
+{
+  return as_nbarray(std::move(x), {x.size()});
+}
+
+template <typename V, std::size_t U>
+auto as_nbarrayp(std::pair<V, std::array<std::size_t, U>>&& x)
+{
+  return as_nbarray(std::move(x.first), x.second.size(), x.second.data());
+}
+
+} // namespace
 
 namespace scifem
 {
@@ -26,7 +61,7 @@ create_real_functionspace(std::shared_ptr<const dolfinx::mesh::Mesh<T>> mesh,
 
   basix::FiniteElement e_v = basix::create_element<T>(
       basix::element::family::P,
-      mesh::cell_type_to_basix_type(mesh->topology()->cell_type()), 0,
+      dolfinx::mesh::cell_type_to_basix_type(mesh->topology()->cell_type()), 0,
       basix::element::lagrange_variant::unset,
       basix::element::dpc_variant::unset, true);
 
@@ -45,8 +80,8 @@ create_real_functionspace(std::shared_ptr<const dolfinx::mesh::Mesh<T>> mesh,
   int index_map_bs = value_size;
   int bs = value_size;
   // Element dof layout
-  fem::ElementDofLayout dof_layout(value_size, e_v.entity_dofs(),
-                                   e_v.entity_closure_dofs(), {}, {});
+  dolfinx::fem::ElementDofLayout dof_layout(value_size, e_v.entity_dofs(),
+                                            e_v.entity_closure_dofs(), {}, {});
   std::size_t num_cells_on_process
       = mesh->topology()->index_map(mesh->topology()->dim())->size_local()
         + mesh->topology()->index_map(mesh->topology()->dim())->num_ghosts();
@@ -63,6 +98,47 @@ create_real_functionspace(std::shared_ptr<const dolfinx::mesh::Mesh<T>> mesh,
 
   return dolfinx::fem::FunctionSpace<T>(mesh, d_el, real_dofmap, value_shape);
 }
+
+std::vector<std::int32_t>
+create_vertex_to_dofmap(std::shared_ptr<const dolfinx::mesh::Topology> topology,
+                        std::shared_ptr<const dolfinx::fem::DofMap> dofmap)
+{
+  // Get number of vertices
+  const std::shared_ptr<const dolfinx::common::IndexMap> v_map
+      = topology->index_map(0);
+  std::size_t num_vertices = v_map->size_local() + v_map->num_ghosts();
+
+  // Get cell to vertex connectivity
+  std::shared_ptr<const dolfinx::graph::AdjacencyList<std::int32_t>> c_to_v
+      = topology->connectivity(topology->dim(), 0);
+  assert(c_to_v);
+
+  // Get vertex dof layout
+  const dolfinx::fem::ElementDofLayout& layout = dofmap->element_dof_layout();
+  const std::vector<std::vector<std::vector<int>>>& entity_dofs
+      = layout.entity_dofs_all();
+  const std::vector<std::vector<int>> vertex_dofs = entity_dofs.front();
+
+  // Get number of cells on process
+  const std::shared_ptr<const dolfinx::common::IndexMap> c_map
+      = topology->index_map(topology->dim());
+  std::size_t num_cells = c_map->size_local() + c_map->num_ghosts();
+
+  std::vector<std::int32_t> vertex_to_dofmap(num_vertices);
+  for (std::size_t i = 0; i < num_cells; i++)
+  {
+    auto vertices = c_to_v->links(i);
+    auto dofs = dofmap->cell_dofs(i);
+    for (std::size_t j = 0; j < vertices.size(); ++j)
+    {
+      const std::vector<int>& vdof = vertex_dofs[j];
+      assert(vdof.size() == 1);
+      vertex_to_dofmap[vertices[j]] = dofs[vdof.front()];
+    }
+  }
+  return vertex_to_dofmap;
+}
+
 } // namespace scifem
 
 namespace scifem_wrapper
@@ -77,6 +153,16 @@ void declare_real_function_space(nanobind::module_& m, std::string type)
          std::vector<std::size_t> value_shape)
       { return scifem::create_real_functionspace<T>(mesh, value_shape); },
       "Create a real function space");
+  m.def(
+      "vertex_to_dofmap",
+      [](std::shared_ptr<const dolfinx::mesh::Topology> topology,
+         std::shared_ptr<const dolfinx::fem::DofMap> dofmap)
+      {
+        std::vector<std::int32_t> v_to_d
+            = scifem::create_vertex_to_dofmap(topology, dofmap);
+        return as_nbarray(v_to_d);
+      },
+      "Create a vertex to dofmap.");
 }
 
 } // namespace scifem_wrapper

src/scifem/__init__.py

@@ -1,11 +1,17 @@
 import dolfinx
 import basix
 import numpy as np
+import numpy.typing as npt
 from . import _scifem  # type: ignore
 from .point_source import PointSource
 from .assembly import assemble_scalar
 from . import xdmf
 
+from .mesh import create_meshtags
+
+__all__ = ["PointSource", "assemble_scalar", "create_meshtags", "xdmf","create_real_functionspace", "assemble_scalar", "PointSource", "xdmf",
+           "vertex_to_dofmap"]
+
 
 def create_real_functionspace(
     mesh: dolfinx.mesh.Mesh, value_shape: tuple[int, ...] = ()
@@ -37,4 +43,18 @@ def create_real_functionspace(
     return dolfinx.fem.FunctionSpace(mesh, ufl_e, cppV)
 
 
-__all__ = ["create_real_functionspace", "assemble_scalar", "PointSource", "xdmf"]
+def vertex_to_dofmap(V: dolfinx.fem.FunctionSpace)-> npt.NDArray[np.int32]:
+    """
+    Create a map from the vertices (local to the process) to the correspondning degrees
+    of freedom.
+
+    Args:
+        V: The function space
+    
+    Returns:
+        An array mapping local vertex i to local degree of freedom i
+
+    Note:
+        If using a blocked space this map is not unrolled for the DofMap block size.
+    """
+    return _scifem.vertex_to_dofmap(V.mesh._cpp_object.topology, V.dofmap._cpp_object)