Merge pull request #36 from scientificcomputing/dokken/vertex_to_dofmap

Add vertex_to_dofmap

Author: jorgensd

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)

README.md

@@ -1,18 +1,20 @@
 # Scientific Computing Tools for Finite Element Methods
+
 This package contains a collection of tools for scientific computing with a focus on finite element methods. The tools are written in Python and are intended to be used in conjunction with the [dolfinx](https://github.com/FEniCS/dolfinx).
 
 Many users that are transitioning from legacy FEniCS to FEniCSx may find the transition difficult due to the lack of some functionalities in FEniCSx.
 This package aims to provide some of the functionalities that are missing in FEniCSx.
 The package is still in its early stages and many functionalities are still missing.
 
 ## Features
+
 - Real-space implementation for usage in DOLFINx (>=v0.8.0)
 - Save quadrature functions as point clouds
 - Save any function that can tabulate dof coordinates as point clouds.
 - Point sources for usage in DOLFINx (>=v0.8.0)
   - Point sources in vector spaces are only supported on v0.9.0, post [DOLFINx PR 3429](https://github.com/FEniCS/dolfinx/pull/3429).
     For older versions, apply one point source in each sub space.
-
+- Maps between degrees of freedom and vertices: `vertex_to_dofmap` and `dof_to_vertex`
 
 ## Installation
 

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,24 @@
 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
 
+__all__ = [
+    "PointSource",
+    "assemble_scalar",
+    "xdmf",
+    "create_real_functionspace",
+    "assemble_scalar",
+    "PointSource",
+    "xdmf",
+    "vertex_to_dofmap",
+    "dof_to_vertexmap",
+]
+
 
 def create_real_functionspace(
     mesh: dolfinx.mesh.Mesh, value_shape: tuple[int, ...] = ()
@@ -37,4 +50,37 @@ 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
+
+    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)
+
+
+def dof_to_vertexmap(V: dolfinx.fem.FunctionSpace) -> npt.NDArray[np.int32]:
+    """
+    Create a map from the degrees of freedom to the vertices of the mesh.
+    As not every degree of freedom is associated with a vertex, every dof that is not
+    associated with a vertex returns -1
+
+    Args:
+        V: The function space
+
+    Returns:
+        An array mapping local dof i to a local vertex
+    """
+    num_dofs_local = V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts
+    dof_to_vertex_map = np.full(num_dofs_local, -1, dtype=np.int32)
+    v_to_d = _scifem.vertex_to_dofmap(V.mesh._cpp_object.topology, V.dofmap._cpp_object)
+    dof_to_vertex_map[v_to_d] = np.arange(len(v_to_d), dtype=np.int32)
+    return dof_to_vertex_map

tests/test_vertex_to_dofmap.py

@@ -0,0 +1,36 @@
+from mpi4py import MPI
+import pytest
+import dolfinx
+from scifem import vertex_to_dofmap, dof_to_vertexmap
+import numpy as np
+
+
+@pytest.mark.parametrize("degree", range(1, 4))
+@pytest.mark.parametrize("value_size", [(), (2,), (2, 3)])
+def test_vertex_to_dofmap_P(degree, value_size):
+    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 2)
+    V = dolfinx.fem.functionspace(mesh, ("Lagrange", degree, value_size))
+
+    v_to_d = vertex_to_dofmap(V)
+    mesh.topology.create_connectivity(0, mesh.topology.dim)
+    geom_indices = dolfinx.mesh.entities_to_geometry(
+        mesh, 0, np.arange(len(v_to_d), dtype=np.int32)
+    )
+
+    x_V = V.tabulate_dof_coordinates()
+    x_g = mesh.geometry.x
+    tol = 1e3 * np.finfo(x_g.dtype).eps
+    mesh = V.mesh
+
+    np.testing.assert_allclose(x_V[v_to_d], x_g[geom_indices.flatten()], atol=tol)
+
+    # Test inverse map
+    d_to_v = dof_to_vertexmap(V)
+    np.testing.assert_allclose(d_to_v[v_to_d], np.arange(len(v_to_d)))
+
+    # Mark all dofs connected with a vertex
+    dof_marker = np.full(len(d_to_v), True, dtype=np.bool_)
+    dof_marker[v_to_d] = False
+
+    other_dofs = np.flatnonzero(dof_marker)
+    np.testing.assert_allclose(d_to_v[other_dofs], -1)