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,29 +31,46 @@
     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:
-        data: The adjacency array. If the array is one-dimensional,
-            offsets should be supplied. If the array is two-dimensional
-            the number of edges per node is the second dimension.
-        offsets: The offsets array with the number of edges per node.
+    data: The adjacency array. If the array is one-dimensional,
+        offsets should be supplied. If the array is two-dimensional
+        the number of edges per node is the second dimension.
+    offsets: The offsets array with the number of edges per node.
 
     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_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)