Use adjacencylists and pure numpy for arrays.

src/networks_fenicsx/assembly.py

@@ -42,22 +42,26 @@ def compute_integration_data(
     """
 
     # Pack all bifurcation in and out fluxes per colored edge in graph
-    influx_color_to_bifurcations: dict[int, list[int]] = {
-        int(color): [] for color in range(network_mesh._num_edge_colors)
+    influx_color_to_bifurcations: dict[int, npt.NDArray[np.int32]] = {
+        int(color): np.empty(0, dtype=np.int32) for color in range(network_mesh.num_edge_colors)
     }
-    outflux_color_to_bifurcations: dict[int, list[int]] = {
-        int(color): [] for color in range(network_mesh._num_edge_colors)
+    outflux_color_to_bifurcations: dict[int, npt.NDArray[np.int32]] = {
+        int(color): np.empty(0, dtype=np.int32) for color in range(network_mesh.num_edge_colors)
     }
-    for bifurcation in network_mesh.bifurcation_values:
-        for color in network_mesh.in_edges(bifurcation):
-            influx_color_to_bifurcations[color].append(int(bifurcation))
-        for color in network_mesh.out_edges(bifurcation):
-            outflux_color_to_bifurcations[color].append(int(bifurcation))
+    for i, bifurcation in enumerate(network_mesh.bifurcation_values):
+        for color in network_mesh.in_edges(i):
+            influx_color_to_bifurcations[color] = np.append(
+                influx_color_to_bifurcations[color], bifurcation
+            )
+        for color in network_mesh.out_edges(i):
+            outflux_color_to_bifurcations[color] = np.append(
+                outflux_color_to_bifurcations[color], bifurcation
+            )
     # Accumulate integration data for all in-edges on the same submesh.
     in_flux_entities: dict[int, npt.NDArray[np.int32]] = {}
     out_flux_entities: dict[int, npt.NDArray[np.int32]] = {}
 
-    for color in range(network_mesh._num_edge_colors):
+    for color in range(network_mesh.num_edge_colors):
         sm = network_mesh.submeshes[color]
         smfm = network_mesh.submesh_facet_markers[color]
         sm.topology.create_connectivity(sm.topology.dim - 1, sm.topology.dim)
@@ -183,7 +187,7 @@ def compute_forms(
            f: source term
            p_bc: neumann bc for pressure
         """
-        num_qs = self._network_mesh._num_edge_colors
+        num_flux_spaces = self._network_mesh.num_edge_colors
 
         test_functions = [ufl.TestFunction(fs) for fs in self.function_spaces]
         trial_functions = [ufl.TrialFunction(fs) for fs in self.function_spaces]
@@ -218,7 +222,6 @@ def compute_forms(
         network_mesh = self._network_mesh
 
         # Assemble edge contributions to a and L
-        num_qs = len(self._network_mesh.submeshes)
         P1_e = fem.functionspace(network_mesh.mesh, ("Lagrange", 1))
         p_bc = fem.Function(P1_e)
         if isinstance(p_bc_ex, ufl.core.expr.Expr):
@@ -244,15 +247,15 @@ def compute_forms(
             ds_edge = ufl.Measure("ds", domain=submesh, subdomain_data=facet_marker)
 
             a[i][i] += R * qs[i] * vs[i] * dx_edge
-            a[num_qs][i] += phi * ufl.dot(ufl.grad(qs[i]), tangent) * dx_edge
-            a[i][num_qs] = -p * ufl.dot(ufl.grad(vs[i]), tangent) * dx_edge
+            a[num_flux_spaces][i] += phi * ufl.dot(ufl.grad(qs[i]), tangent) * dx_edge
+            a[i][num_flux_spaces] = -p * ufl.dot(ufl.grad(vs[i]), tangent) * dx_edge
 
             # Add all boundary contributions
             L[i] = p_bc * vs[i] * ds_edge(network_mesh.in_marker) - p_bc * vs[i] * ds_edge(
                 network_mesh.out_marker
             )
 
-        L[num_qs] += f * phi * dx_global
+        L[num_flux_spaces] += f * phi * dx_global
 
         # Multiplier mesh and flux share common parent mesh.
         # We create unique integration entities for each in and out branch

src/networks_fenicsx/mesh.py

@@ -64,8 +64,9 @@ class NetworkMesh:
     # Graph properties
     _geom_dim: int
     _num_edge_colors: int
-    _bifurcation_in_color: dict[int, list[int]]
-    _bifurcation_out_color: dict[int, list[int]]
+    _bifurcation_values: npt.NDArray[np.int32]
+    _bifurcation_in_color: _graph.AdjacencyList
+    _bifurcation_out_color: _graph.AdjacencyList
 
     # Mesh properties
     _msh: mesh.Mesh | None
@@ -75,7 +76,6 @@ class NetworkMesh:
     _edge_meshes: list[mesh.Mesh]
     _edge_entity_maps: list[mesh.EntityMap]
     _tangent: fem.Function
-    _bifurcation_values: npt.NDArray[np.int32]
     _boundary_values: npt.NDArray[np.int32]
     _lm_mesh: mesh.Mesh | None
     _lm_map: mesh.EntityMap | None
@@ -167,8 +167,10 @@ def _build_mesh(
         bifurcation_values = None
         boundary_values = None
         number_of_nodes = None
-        bifurcation_in_color: dict[int, list[int]] | None = None
-        bifurcation_out_color: dict[int, list[int]] | None = None
+        bifurcation_in_color: npt.NDArray[np.int32] | None = None
+        bifurcation_in_offsets: npt.NDArray[np.int32] | None = None
+        bifurcation_out_color: npt.NDArray[np.int32] | None = None
+        bifurcation_out_offsets: npt.NDArray[np.int32] | None = None
         if comm.rank == graph_rank:
             assert isinstance(graph, nx.DiGraph), f"Directional graph not present of {graph_rank}"
             geom_dim = len(graph.nodes[1]["pos"])
@@ -184,29 +186,41 @@ def _build_mesh(
             boundary_values = np.flatnonzero(nodes_with_degree == 1)
             max_connections = np.max(nodes_with_degree)
 
-            bifurcation_in_color = {}
-            bifurcation_out_color = {}
+            bifurcation_in_color = np.empty(0, dtype=np.int32)
+            bifurcation_in_offsets = np.zeros(1, dtype=np.int32)
+            bifurcation_out_color = np.empty(0, dtype=np.int32)
+            bifurcation_out_offsets = np.zeros(1, dtype=np.int32)
             for bifurcation in bifurcation_values:
                 in_edges = graph.in_edges(bifurcation)
-                bifurcation_in_color[int(bifurcation)] = []
-                for edge in in_edges:
-                    bifurcation_in_color[int(bifurcation)].append(edge_coloring[edge])
+                bifurcation_in_offsets = np.append(
+                    bifurcation_in_offsets, len(bifurcation_in_color) + len(in_edges)
+                )
+                bifurcation_in_color = np.append(
+                    bifurcation_in_color, [edge_coloring[edge] for edge in in_edges]
+                )
                 out_edges = graph.out_edges(bifurcation)
-                bifurcation_out_color[int(bifurcation)] = []
-                for edge in out_edges:
-                    bifurcation_out_color[int(bifurcation)].append(edge_coloring[edge])
+                bifurcation_out_offsets = np.append(
+                    bifurcation_out_offsets, len(bifurcation_out_color) + len(out_edges)
+                )
+                bifurcation_out_color = np.append(
+                    bifurcation_out_color, [edge_coloring[edge] for edge in out_edges]
+                )
 
             # Map boundary_values to inlet and outlet data from graph
-            boundary_in_nodes = []
-            boundary_out_nodes = []
+            boundary_in_nodes = np.empty(0, dtype=np.int32)
+            boundary_out_nodes = np.empty(0, dtype=np.int32)
             for boundary in boundary_values:
                 in_edges = graph.in_edges(boundary)
                 out_edges = graph.out_edges(boundary)
                 assert len(in_edges) + len(out_edges) == 1, "Boundary node with multiple edges"
                 if len(in_edges) == 1:
-                    boundary_in_nodes.append(int(boundary))
+                    boundary_in_nodes = (
+                        np.append(boundary_in_nodes, np.int32(boundary)).flatten().astype(np.int32)
+                    )
                 else:
-                    boundary_out_nodes.append(int(boundary))
+                    boundary_out_nodes = (
+                        np.append(boundary_out_nodes, np.int32(boundary)).flatten().astype(np.int32)
+                    )
 
         comm.bcast(num_edge_colors, root=graph_rank)
         num_edge_colors, number_of_nodes, max_connections, geom_dim = comm.bcast(
@@ -217,8 +231,19 @@ def _build_mesh(
             (bifurcation_values, boundary_values), root=graph_rank
         )
         comm.barrier()
-        bifurcation_in_color, bifurcation_out_color = comm.bcast(
-            (bifurcation_in_color, bifurcation_out_color), root=graph_rank
+        (
+            bifurcation_in_color,
+            bifurcation_in_offsets,
+            bifurcation_out_color,
+            bifurcation_out_offsets,
+        ) = comm.bcast(
+            (
+                bifurcation_in_color,
+                bifurcation_in_offsets,
+                bifurcation_out_color,
+                bifurcation_out_offsets,
+            ),
+            root=graph_rank,
         )
         comm.barrier()
 
@@ -228,8 +253,12 @@ def _build_mesh(
         self._boundary_values = boundary_values
 
         # Create lookup of in and out colors for each bifurcation
-        self._bifurcation_in_color = bifurcation_in_color
-        self._bifurcation_out_color = bifurcation_out_color
+        self._bifurcation_in_color = _graph.adjacencylist(
+            bifurcation_in_color, bifurcation_in_offsets
+        )
+        self._bifurcation_out_color = _graph.adjacencylist(
+            bifurcation_out_color, bifurcation_out_offsets
+        )
 
         # Generate mesh
         tangents: npt.NDArray[np.float64] = np.empty((0, self._geom_dim), dtype=np.float64)
@@ -423,12 +452,6 @@ def _build_network_submeshes(self):
                 mesh.meshtags(edge_mesh, 0, marked_vertices, marked_values)
             )
 
-    def in_edges(self, bifurcation: int) -> list[int]:
-        return self._bifurcation_in_color[bifurcation]
-
-    def out_edges(self, bifurcation: int) -> list[int]:
-        return self._bifurcation_out_color[bifurcation]
-
     @property
     def submesh_facet_markers(self) -> list[mesh.MeshTags]:
         if self._submesh_facet_markers is None:
@@ -482,6 +505,23 @@ def bifurcation_values(self) -> npt.NDArray[np.int32]:
     def boundary_values(self) -> npt.NDArray[np.int32]:
         return self._boundary_values
 
+    def in_edges(self, bifurcation_idx: int) -> npt.NDArray[np.int32 | np.int64]:
+        """Return the list of in-edge colors for a given bifurcation node.
+        Index is is the index of the bifurcation in {py:meth}`self.bifurcation_values`."""
+        assert bifurcation_idx < len(self.bifurcation_values)
+        return self._bifurcation_in_color.links(np.int32(bifurcation_idx))
+
+    def out_edges(self, bifurcation_idx: int) -> npt.NDArray[np.int32 | np.int64]:
+        """Return the list of out-edge colors for a given bifurcation node.
+        Index is is the index of the bifurcation in {py:meth}`self.bifurcation_values`."""
+        assert bifurcation_idx < len(self.bifurcation_values)
+        return self._bifurcation_out_color.links(np.int32(bifurcation_idx))
+
+    @property
+    def num_edge_colors(self) -> int:
+        """Return the number of edge colors in the network mesh."""
+        return self._num_edge_colors
+
     @property
     def in_marker(self) -> int:
         return self._in_marker

tests/test_edge_info.py

@@ -0,0 +1,55 @@
+import networkx as nx
+import numpy as np
+import pytest
+
+from networks_fenicsx import NetworkMesh
+
+
+@pytest.mark.parametrize("N", [10, 50])
+def test_edge_info(N: int):
+    # Create manual graph where we have five bifurcations.
+    # One inlet (0) -> (1) -> (7).
+    G = nx.DiGraph()
+    G.add_node(0, pos=np.zeros(3))
+    G.add_node(1, pos=np.array([0.0, 0.0, 1.0]))
+    G.add_node(2, pos=np.array([0.2, 0.2, 2.0]))
+    G.add_node(3, pos=np.array([-0.2, 0.3, 2.0]))
+    G.add_node(4, pos=np.array([0.0, 0.1, 2.1]))
+    G.add_node(5, pos=np.array([0.1, -0.1, 3.0]))
+    G.add_node(6, pos=np.array([-0.3, 0.4, 4.0]))
+    G.add_node(7, pos=1.1 * G.nodes[1]["pos"])
+    G.add_edge(0, 1)
+    G.add_edge(1, 7)
+    # Split into three bifurcations (7)->(2), (7)->(3), (7)->(4)
+    G.add_edge(7, 2)
+    # First branch goes right to gathering
+    G.add_edge(2, 5)
+    # Second branch goes through an extra point(3), before meeting at point (4)
+    G.add_edge(7, 3)
+    G.add_edge(3, 4)
+    G.add_edge(4, 5)
+    # Last branch goes directly to gathering
+    G.add_edge(7, 4)
+    G.add_edge(5, 6)
+
+    network_mesh = NetworkMesh(G, N=N)
+    assert len(network_mesh.bifurcation_values) == 6
+    # Bifurcation values are sorted in increasing order
+    np.testing.assert_allclose([1, 2, 3, 4, 5, 7], network_mesh.bifurcation_values)
+    assert len(network_mesh.in_edges(0)) == 1
+    assert len(network_mesh.out_edges(0)) == 1
+
+    assert len(network_mesh.in_edges(1)) == 1
+    assert len(network_mesh.out_edges(1)) == 1
+
+    assert len(network_mesh.in_edges(2)) == 1
+    assert len(network_mesh.out_edges(2)) == 1
+
+    assert len(network_mesh.in_edges(3)) == 2
+    assert len(network_mesh.out_edges(3)) == 1
+
+    assert len(network_mesh.in_edges(4)) == 2
+    assert len(network_mesh.out_edges(4)) == 1
+
+    assert len(network_mesh.in_edges(5)) == 1
+    assert len(network_mesh.out_edges(5)) == 3