Add adjacency list python wrapper

python/dolfinx/geometry.py

@@ -13,11 +13,10 @@
 import numpy.typing as npt
 
 if typing.TYPE_CHECKING:
-    from dolfinx.cpp.graph import AdjacencyList_int32
     from dolfinx.mesh import Mesh
 
-
 from dolfinx import cpp as _cpp
+from dolfinx.graph import AdjacencyList
 
 __all__ = [
     "BoundingBoxTree",
@@ -155,9 +154,7 @@ def compute_collisions_trees(
     return _cpp.geometry.compute_collisions_trees(tree0._cpp_object, tree1._cpp_object)
 
 
-def compute_collisions_points(
-    tree: BoundingBoxTree, x: npt.NDArray[np.floating]
-) -> _cpp.graph.AdjacencyList_int32:
+def compute_collisions_points(tree: BoundingBoxTree, x: npt.NDArray[np.floating]) -> AdjacencyList:
     """Compute collisions between points and leaf bounding boxes.
 
     Bounding boxes can overlap, therefore points can collide with more
@@ -172,7 +169,7 @@ def compute_collisions_points(
        point.
 
     """
-    return _cpp.geometry.compute_collisions_points(tree._cpp_object, x)
+    return AdjacencyList(_cpp.geometry.compute_collisions_points(tree._cpp_object, x))
 
 
 def compute_closest_entity(
@@ -216,8 +213,8 @@ def create_midpoint_tree(mesh: Mesh, dim: int, entities: npt.NDArray[np.int32])
 
 
 def compute_colliding_cells(
-    mesh: Mesh, candidates: AdjacencyList_int32, x: npt.NDArray[np.floating]
-):
+    mesh: Mesh, candidates: AdjacencyList, x: npt.NDArray[np.floating]
+) -> AdjacencyList:
     """From a mesh, find which cells collide with a set of points.
 
     Args:
@@ -231,10 +228,14 @@ def compute_colliding_cells(
         collide with the ith point.
 
     """
-    return _cpp.geometry.compute_colliding_cells(mesh._cpp_object, candidates, x)
+    return AdjacencyList(
+        _cpp.geometry.compute_colliding_cells(mesh._cpp_object, candidates._cpp_object, x)
+    )
 
 
-def squared_distance(mesh: Mesh, dim: int, entities: list[int], points: npt.NDArray[np.floating]):
+def squared_distance(
+    mesh: Mesh, dim: int, entities: list[int], points: npt.NDArray[np.floating]
+) -> npt.NDArray[np.floating]:
     """Compute the squared distance between a point and a mesh entity.
 
     The distance is computed between the ith input points and the ith

python/dolfinx/graph.py

@@ -1,4 +1,4 @@
-# Copyright (C) 2021 Garth N. Wells
+# Copyright (C) 2021-2024 Garth N. Wells and Paul T. Kühner
 #
 # This file is part of DOLFINx (https://www.fenicsproject.org)
 #
@@ -7,7 +7,10 @@
 
 from __future__ import annotations
 
+from typing import Optional, Union
+
 import numpy as np
+import numpy.typing as npt
 
 from dolfinx import cpp as _cpp
 from dolfinx.cpp.graph import partitioner
@@ -28,10 +31,34 @@
     pass
 
 
-__all__ = ["adjacencylist", "partitioner"]
+__all__ = ["AdjacencyList", "adjacencylist", "partitioner"]
+
+
+class AdjacencyList:
+    _cpp_object: Union[_cpp.la.AdjacencyList_int32, _cpp.la.AdjacencyList_int64]
+
+    def __init__(self, cpp_object: Union[_cpp.la.AdjacencyList_int32, _cpp.la.AdjacencyList_int64]):
+        self._cpp_object = cpp_object
+
+    def links(self, node: Union[np.int32, np.int64]) -> npt.NDArray[Union[np.int32, np.int64]]:
+        return self._cpp_object.links(node)
 
+    @property
+    def array(self) -> npt.NDArray[Union[np.int32, np.int64]]:
+        return self._cpp_object.array
 
-def adjacencylist(data: np.ndarray, offsets=None):
+    @property
+    def offsets(self) -> npt.NDArray[np.int32]:
+        return self._cpp_object.offsets
+
+    @property
+    def num_nodes(self) -> np.int32:
+        return self._cpp_object.num_nodes
+
+
+def adjacencylist(
+    data: npt.NDArray[Union[np.int32, np.int64]], offsets: Optional[npt.NDArray[np.int32]] = None
+) -> AdjacencyList:
     """Create an AdjacencyList for int32 or int64 datasets.
 
     Args:
@@ -42,15 +69,8 @@ def adjacencylist(data: np.ndarray, offsets=None):
 
     Returns:
         An adjacency list.
-
     """
-    if offsets is None:
-        try:
-            return _cpp.graph.AdjacencyList_int32(data)
-        except TypeError:
-            return _cpp.graph.AdjacencyList_int64(data)
-    else:
-        try:
-            return _cpp.graph.AdjacencyList_int32(data, offsets)
-        except TypeError:
-            return _cpp.graph.AdjacencyList_int64(data, offsets)
+    is_32bit = np.isdtype(data.dtype, np.int32)
+    cpp_t = _cpp.graph.AdjacencyList_int32 if is_32bit else _cpp.graph.AdjacencyList_int64
+    cpp_object = cpp_t(data, offsets) if offsets is not None else cpp_t(data)
+    return AdjacencyList(cpp_object)

python/dolfinx/mesh.py

@@ -33,6 +33,7 @@
 from dolfinx.cpp.refinement import RefinementOption
 from dolfinx.fem import CoordinateElement as _CoordinateElement
 from dolfinx.fem import coordinate_element as _coordinate_element
+from dolfinx.graph import AdjacencyList
 
 __all__ = [
     "meshtags_from_entities",
@@ -735,7 +736,7 @@ def meshtags(
 
 
 def meshtags_from_entities(
-    msh: Mesh, dim: int, entities: _cpp.graph.AdjacencyList_int32, values: npt.NDArray[typing.Any]
+    msh: Mesh, dim: int, entities: AdjacencyList, values: npt.NDArray[typing.Any]
 ):
     """Create a :class:dolfinx.mesh.MeshTags` object that associates
     data with a subset of mesh entities, where the entities are defined
@@ -762,7 +763,9 @@ def meshtags_from_entities(
     elif isinstance(values, float):
         values = np.full(entities.num_nodes, values, dtype=np.double)
     values = np.asarray(values)
-    return MeshTags(_cpp.mesh.create_meshtags(msh.topology._cpp_object, dim, entities, values))
+    return MeshTags(
+        _cpp.mesh.create_meshtags(msh.topology._cpp_object, dim, entities._cpp_object, values)
+    )
 
 
 def create_interval(

python/test/unit/fem/test_assemble_domains.py

@@ -14,6 +14,7 @@
 from dolfinx import cpp as _cpp
 from dolfinx import default_scalar_type, fem, la
 from dolfinx.fem import Constant, Function, assemble_scalar, dirichletbc, form, functionspace
+from dolfinx.graph import adjacencylist
 from dolfinx.mesh import (
     GhostMode,
     Mesh,
@@ -33,9 +34,7 @@ def mesh():
 def create_cell_meshtags_from_entities(mesh: Mesh, dim: int, cells: np.ndarray, values: np.ndarray):
     mesh.topology.create_connectivity(mesh.topology.dim, 0)
     cell_to_vertices = mesh.topology.connectivity(mesh.topology.dim, 0)
-    entities = _cpp.graph.AdjacencyList_int32(
-        np.array([cell_to_vertices.links(cell) for cell in cells])
-    )
+    entities = adjacencylist(np.array([cell_to_vertices.links(cell) for cell in cells]))
     return meshtags_from_entities(mesh, dim, entities, values)
 
 

python/test/unit/fem/test_assembler.py

@@ -1059,7 +1059,7 @@ def test_assemble_empty_rank_mesh(self):
         def partitioner(comm, nparts, local_graph, num_ghost_nodes):
             """Leave cells on the curent rank"""
             dest = np.full(len(cells), comm.rank, dtype=np.int32)
-            return graph.adjacencylist(dest)
+            return graph.adjacencylist(dest)._cpp_object
 
         if comm.rank == 0:
             # Put cells on rank 0

python/test/unit/fem/test_dofmap.py

@@ -436,7 +436,8 @@ def test_empty_rank_collapse():
     def self_partitioner(comm: MPI.Intracomm, n, m, topo):
         dests = np.full(len(topo[0]) // 2, comm.rank, dtype=np.int32)
         offsets = np.arange(len(topo[0]) // 2 + 1, dtype=np.int32)
-        return dolfinx.graph.adjacencylist(dests, offsets)
+        # TODO: can we improve on this interface? I.e. warp to do cpp type conversion automatically
+        return dolfinx.graph.adjacencylist(dests, offsets)._cpp_object
 
     mesh = create_mesh(MPI.COMM_WORLD, cells, nodes, c_el, partitioner=self_partitioner)
 

python/test/unit/io/test_adios2.py

@@ -184,7 +184,7 @@ def test_empty_rank_mesh(self, tempdir):
         def partitioner(comm, nparts, local_graph, num_ghost_nodes):
             """Leave cells on the current rank"""
             dest = np.full(len(cells), comm.rank, dtype=np.int32)
-            return adjacencylist(dest)
+            return adjacencylist(dest)._cpp_object
 
         if comm.rank == 0:
             cells = np.array([[0, 1, 2], [0, 2, 3]], dtype=np.int64)

python/test/unit/mesh/test_dual_graph.py

@@ -25,7 +25,9 @@ def test_dgrsph_1d():
         x = 0
     # Circular chain of interval cells
     cells = [[n0, n0 + 1], [n0 + 1, n0 + 2], [n0 + 2, x]]
-    w = mesh.build_dual_graph(MPI.COMM_WORLD, mesh.CellType.interval, to_adj(cells, np.int64))
+    w = mesh.build_dual_graph(
+        MPI.COMM_WORLD, mesh.CellType.interval, to_adj(cells, np.int64)._cpp_object
+    )
     assert w.num_nodes == 3
     for i in range(w.num_nodes):
         assert len(w.links(i)) == 2

python/test/unit/mesh/test_mesh.py

@@ -555,7 +555,7 @@ def test_empty_rank_mesh(dtype):
     def partitioner(comm, nparts, local_graph, num_ghost_nodes):
         """Leave cells on the curent rank"""
         dest = np.full(len(cells), comm.rank, dtype=np.int32)
-        return graph.adjacencylist(dest)
+        return graph.adjacencylist(dest)._cpp_object
 
     if comm.rank == 0:
         cells = np.array([[0, 1, 2], [0, 2, 3]], dtype=np.int64)