From bf1b8a1c0f9b871cf76510876092a903b892f8ba Mon Sep 17 00:00:00 2001 From: schnellerhase <56360279+schnellerhase@users.noreply.github.com> Date: Wed, 13 Nov 2024 18:22:22 +0100 Subject: [PATCH] Add adjacency list python wrapper --- python/dolfinx/geometry.py | 21 ++++---- python/dolfinx/graph.py | 48 +++++++++++++------ python/dolfinx/io/gmshio.py | 7 +-- python/dolfinx/mesh.py | 7 ++- python/test/unit/fem/test_assemble_domains.py | 6 +-- python/test/unit/fem/test_assembler.py | 2 +- python/test/unit/fem/test_dofmap.py | 3 +- python/test/unit/io/test_adios2.py | 2 +- python/test/unit/mesh/test_dual_graph.py | 4 +- python/test/unit/mesh/test_mesh.py | 2 +- 10 files changed, 64 insertions(+), 38 deletions(-) diff --git a/python/dolfinx/geometry.py b/python/dolfinx/geometry.py index 5e126939783..371e2100ab9 100644 --- a/python/dolfinx/geometry.py +++ b/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 diff --git a/python/dolfinx/graph.py b/python/dolfinx/graph.py index dd16514a79a..517caa568a5 100644 --- a/python/dolfinx/graph.py +++ b/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) diff --git a/python/dolfinx/io/gmshio.py b/python/dolfinx/io/gmshio.py index 6cc89917d26..63f5630f3e2 100644 --- a/python/dolfinx/io/gmshio.py +++ b/python/dolfinx/io/gmshio.py @@ -19,6 +19,7 @@ from dolfinx import cpp as _cpp from dolfinx import default_real_type from dolfinx.cpp.graph import AdjacencyList_int32 +from dolfinx.graph import AdjacencyList, adjacencylist from dolfinx.io.utils import distribute_entity_data from dolfinx.mesh import CellType, Mesh, create_mesh, meshtags, meshtags_from_entities @@ -298,7 +299,7 @@ def model_to_mesh( mesh, mesh.topology.dim, cells, cell_values ) mesh.topology.create_connectivity(mesh.topology.dim, 0) - adj = _cpp.graph.AdjacencyList_int32(local_entities) + adj = adjacencylist(local_entities) ct = meshtags_from_entities( mesh, mesh.topology.dim, adj, local_values.astype(np.int32, copy=False) ) @@ -323,7 +324,7 @@ def model_to_mesh( mesh, tdim - 1, marked_facets, facet_values ) mesh.topology.create_connectivity(topology.dim - 1, tdim) - adj = _cpp.graph.AdjacencyList_int32(local_entities) + adj = adjacencylist(local_entities) ft = meshtags_from_entities(mesh, tdim - 1, adj, local_values.astype(np.int32, copy=False)) ft.name = "Facet tags" else: @@ -338,7 +339,7 @@ def read_from_msh( rank: int = 0, gdim: int = 3, partitioner: typing.Optional[ - typing.Callable[[_MPI.Comm, int, int, AdjacencyList_int32], AdjacencyList_int32] + typing.Callable[[_MPI.Comm, int, int, AdjacencyList], AdjacencyList_int32] ] = None, ) -> tuple[Mesh, _cpp.mesh.MeshTags_int32, _cpp.mesh.MeshTags_int32]: """Read a Gmsh .msh file and return a :class:`dolfinx.mesh.Mesh` and cell facet markers. diff --git a/python/dolfinx/mesh.py b/python/dolfinx/mesh.py index e58563a2bf1..1e69576a1c5 100644 --- a/python/dolfinx/mesh.py +++ b/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( diff --git a/python/test/unit/fem/test_assemble_domains.py b/python/test/unit/fem/test_assemble_domains.py index 839fbcc31c4..324346c8025 100644 --- a/python/test/unit/fem/test_assemble_domains.py +++ b/python/test/unit/fem/test_assemble_domains.py @@ -11,9 +11,9 @@ import pytest import ufl -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 +33,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) diff --git a/python/test/unit/fem/test_assembler.py b/python/test/unit/fem/test_assembler.py index 452d90eb077..f5507147c2f 100644 --- a/python/test/unit/fem/test_assembler.py +++ b/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 diff --git a/python/test/unit/fem/test_dofmap.py b/python/test/unit/fem/test_dofmap.py index 0115893407f..4275fc37a04 100644 --- a/python/test/unit/fem/test_dofmap.py +++ b/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) diff --git a/python/test/unit/io/test_adios2.py b/python/test/unit/io/test_adios2.py index 0b97b0f98a3..692d0aa126f 100644 --- a/python/test/unit/io/test_adios2.py +++ b/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) diff --git a/python/test/unit/mesh/test_dual_graph.py b/python/test/unit/mesh/test_dual_graph.py index e71c0ef2ff4..e1b43b63fdc 100644 --- a/python/test/unit/mesh/test_dual_graph.py +++ b/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 diff --git a/python/test/unit/mesh/test_mesh.py b/python/test/unit/mesh/test_mesh.py index e620e1c0569..d7907f868bb 100644 --- a/python/test/unit/mesh/test_mesh.py +++ b/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)