Route choice API, link loading, path file generation

.github/workflows/unit_tests.yml

@@ -31,7 +31,7 @@ jobs:
     runs-on: ${{ matrix.os}}
     strategy:
       matrix:
-        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
+        python-version: ['3.9', '3.10', '3.11', '3.12']
         os: [windows-latest, ubuntu-latest]
 
       max-parallel: 20

.pre-commit-config.yaml

@@ -1,9 +1,14 @@
 repos:
--   repo: https://github.com/ambv/black
-    rev: 22.3.0
-    hooks:
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.3.0
+  hooks:
+    # Run the linter.
+    - id: ruff
+      args: [ --fix ]
+    # Run the formatter.
+    - id: ruff-format
+- repo: https://github.com/ambv/black
+  rev: 24.1.1
+  hooks:
     - id: black
--   repo: https://github.com/pycqa/flake8
-    rev: 4.0.1
-    hooks:
-    - id: flake8

aequilibrae/matrix/__init__.py

@@ -1,2 +1,3 @@
 from .aequilibrae_matrix import AequilibraeMatrix, matrix_export_types
 from .aequilibrae_data import AequilibraeData, data_export_types
+from .sparse_matrix import Sparse, COO

aequilibrae/matrix/sparse_matrix.pxd

@@ -0,0 +1,13 @@
+from libcpp.vector cimport vector
+
+cdef class Sparse:
+    pass
+
+cdef class COO(Sparse):
+    cdef:
+        vector[size_t] *row
+        vector[size_t] *col
+        vector[double] *data
+        readonly object shape
+
+    cdef void append(COO self, size_t i, size_t j, double v) noexcept nogil

aequilibrae/matrix/sparse_matrix.pyx

@@ -0,0 +1,122 @@
+from libcpp.vector cimport vector
+from libcpp cimport nullptr
+from cython.operator cimport dereference as d
+
+import scipy.sparse
+import numpy as np
+import openmatrix as omx
+
+cdef class Sparse:
+    """
+    A class to implement sparse matrix operations such as reading, writing, and indexing
+    """
+
+    def __cinit__(self):
+        """C level init. For C memory allocation and initialisation. Called exactly once per object."""
+        pass
+
+    def __init__(self):
+        """Python level init, may be called multiple times, for things that can't be done in __cinit__."""
+        pass
+
+    def __dealloc__(self):
+        """
+        C level deallocation. For freeing memory allocated by this object. *Must* have GIL, `self` may be in a
+        partially deallocated state already.
+        """
+        pass
+
+    def to_disk(self, path, name: str):
+        f = omx.open_file(path, "a")
+        try:
+            f[name] = self.to_scipy().tocsr().toarray()
+        finally:
+            f.close()
+
+    @classmethod
+    def from_disk(cls, path, names=None, aeq=False):
+        """
+        Read a OMX file and return a dictionary of matrix names to a scipy.sparse matrix, or
+        aequilibrae.matrix.sparse matrix.
+        """
+        f = omx.open_file(path, "r")
+        res = {}
+        try:
+            for matrix in (f.list_matrices() if names is None else names):
+                if aeq:
+                    res[matrix] = cls.from_matrix(f[matrix])
+                else:
+                    res[matrix] = scipy.sparse.csr_matrix(f[matrix])
+            return res
+        finally:
+            f.close()
+
+
+cdef class COO(Sparse):
+    """
+    A class to implement sparse matrix operations such as reading, writing, and indexing
+    """
+
+    def __cinit__(self):
+        """C level init. For C memory allocation and initialisation. Called exactly once per object."""
+
+        self.row = new vector[size_t]()
+        self.col = new vector[size_t]()
+        self.data = new vector[double]()
+
+    def __init__(self, shape=None):
+        """Python level init, may be called multiple times, for things that can't be done in __cinit__."""
+
+        self.shape = shape
+
+    def __dealloc__(self):
+        """
+        C level deallocation. For freeing memory allocated by this object. *Must* have GIL, `self` may be in a
+        partially deallocated state already.
+        """
+
+        del self.row
+        self.row = <vector[size_t] *>nullptr
+
+        del self.col
+        self.col = <vector[size_t] *>nullptr
+
+        del self.data
+        self.data = <vector[double] *>nullptr
+
+    def to_scipy(self, shape=None, dtype=np.float64):
+        """
+        Create scipy.sparse.coo_matrix from this COO matrix.
+        """
+        row = <size_t[:self.row.size()]>&d(self.row)[0]
+        col = <size_t[:self.col.size()]>&d(self.col)[0]
+        data = <double[:self.data.size()]>&d(self.data)[0]
+
+        if shape is None:
+            shape = self.shape
+
+        return scipy.sparse.coo_matrix((data, (row, col)), dtype=dtype, shape=shape)
+
+    @classmethod
+    def from_matrix(cls, m):
+        """
+        Create COO matrix from an dense or scipy-like matrix.
+        """
+        if not isinstance(m, scipy.sparse.coo_matrix):
+            m = scipy.sparse.coo_matrix(m)
+
+        self = <COO?>cls()
+
+        cdef size_t[:] row = m.row.astype(np.uint64), col = m.row.astype(np.uint64)
+        cdef double[:] data = m.data
+
+        self.row.insert(self.row.end(), &row[0], &row[-1] + 1)
+        self.col.insert(self.col.end(), &col[0], &col[-1] + 1)
+        self.data.insert(self.data.end(), &data[0], &data[-1] + 1)
+
+        return self
+
+    cdef void append(COO self, size_t i, size_t j, double v) noexcept nogil:
+        self.row.push_back(i)
+        self.col.push_back(j)
+        self.data.push_back(v)

aequilibrae/paths/__init__.py

@@ -8,6 +8,7 @@
 from aequilibrae.paths.traffic_assignment import TrafficAssignment, TransitAssignment
 from aequilibrae.paths.vdf import VDF
 from aequilibrae.paths.graph import Graph, TransitGraph
+from aequilibrae.paths.route_choice import RouteChoice
 
 from aequilibrae import global_logger
 

aequilibrae/paths/graph.py

@@ -4,14 +4,35 @@
 from datetime import datetime
 from os.path import join
 from typing import List, Tuple, Optional
+import dataclasses
 
 import numpy as np
 import pandas as pd
-from aequilibrae.paths.graph_building import build_compressed_graph
+from aequilibrae.paths.graph_building import build_compressed_graph, create_compressed_link_network_mapping
 
 from aequilibrae.context import get_logger
 
 
+@dataclasses.dataclass
+class NetworkGraphIndices:
+    network_ab_idx: np.array
+    network_ba_idx: np.array
+    graph_ab_idx: np.array
+    graph_ba_idx: np.array
+
+
+def _get_graph_to_network_mapping(lids, direcs):
+    num_uncompressed_links = int(np.unique(lids).shape[0])
+    indexing = np.zeros(int(lids.max()) + 1, np.uint64)
+    indexing[np.unique(lids)[:]] = np.arange(num_uncompressed_links)
+
+    graph_ab_idx = direcs > 0
+    graph_ba_idx = direcs < 0
+    network_ab_idx = indexing[lids[graph_ab_idx]]
+    network_ba_idx = indexing[lids[graph_ba_idx]]
+    return NetworkGraphIndices(network_ab_idx, network_ba_idx, graph_ab_idx, graph_ba_idx)
+
+
 class GraphBase(ABC):  # noqa: B024
     """
     Graph class.
@@ -95,6 +116,9 @@ def __init__(self, logger=None):
 
         self.dead_end_links = np.array([])
 
+        self.compressed_link_network_mapping_idx = None
+        self.compressed_link_network_mapping_data = None
+
         # Randomly generate a unique Graph ID randomly
         self._id = uuid.uuid4().hex
 
@@ -167,6 +191,11 @@ def prepare_graph(self, centroids: Optional[np.ndarray]) -> None:
             self.__build_compressed_graph()
             self.compact_num_links = self.compact_graph.shape[0]
 
+        # The cache property should be recalculated when the graph has been re-prepared
+        self.compressed_link_network_mapping_idx = None
+        self.compressed_link_network_mapping_data = None
+        self.network_compressed_node_mapping = None
+
     def __build_compressed_graph(self):
         build_compressed_graph(self)
 
@@ -505,6 +534,27 @@ def save_compressed_correspondence(self, path, mode_name, mode_id):
         node_path = join(path, f"nodes_to_indices_c{mode_name}_{mode_id}.feather")
         pd.DataFrame(self.nodes_to_indices, columns=["node_index"]).to_feather(node_path)
 
+    def create_compressed_link_network_mapping(self):
+        """
+        Create two arrays providing a mapping of compressed id to link id.
+
+        Uses sparse compression. Index ``idx`` by the by compressed id and compressed id + 1, the
+        network IDs are then in the range ``idx[id]:idx[id + 1]``.
+
+        .. code-block:: python
+
+            >>> idx, data = graph.compressed_link_network_mapping
+            >>> data[idx[id]:idx[id + 1]]  # ==> Slice of network ID's corresponding to the compressed ID
+
+        Links not in the compressed graph are not contained within the ``data`` array.
+
+        :Returns:
+            **idx** (:obj:`np.array`): index array for ``data``
+            **data** (:obj:`np.array`): array of link ids
+        """
+
+        return create_compressed_link_network_mapping(self)
+
 
 class Graph(GraphBase):
     def __init__(self, *args, **kwargs):

aequilibrae/paths/graph_building.pyx

@@ -325,6 +325,10 @@ def build_compressed_graph(graph):
             "link_id": np.arange(slink),
         }
     )
+
+    # Link compression can introduce new simple cycles into the graph
+    comp_lnk = comp_lnk[comp_lnk.a_node != comp_lnk.b_node]
+
     max_link_id = link_id_max * 10
     comp_lnk.link_id += max_link_id
 
@@ -377,3 +381,76 @@ def build_compressed_graph(graph):
     # If will refer all the links that have no correlation to an element beyond the last link
     # This element will always be zero during assignment
     graph.graph.__compressed_id__ = graph.graph.__compressed_id__.fillna(graph.compact_graph.id.max() + 1).astype(np.int64)
+
+
+@cython.embedsignature(True)
+@cython.boundscheck(False)
+@cython.initializedcheck(False)
+def create_compressed_link_network_mapping(graph):
+        # Cache the result, this isn't a huge computation but isn't worth doing twice
+        if (
+            graph.compressed_link_network_mapping_idx is not None
+            and graph.compressed_link_network_mapping_data is not None
+            and graph.network_compressed_node_mapping is not None
+        ):
+            return (
+                graph.compressed_link_network_mapping_idx,
+                graph.compressed_link_network_mapping_data,
+                graph.network_compressed_node_mapping,
+            )
+
+        cdef:
+            long long i, j, a_node, x, b_node, tmp, compressed_id
+            long long[:] b
+            long long[:] values
+            np.uint32_t[:] idx
+            np.uint32_t[:] data
+            np.int32_t[:] node_mapping
+
+        # This method requires that graph.graph is sorted on the a_node IDs, since that's done already we don't
+        # bother redoing sorting it.
+
+        # Some links are completely removed from the network, they are assigned ID `graph.compact_graph.id.max() + 1`,
+        # we skip them.
+        filtered = graph.graph[graph.graph.__compressed_id__ != graph.compact_graph.id.max() + 1]
+        gb = filtered.groupby(by="__compressed_id__", sort=True)
+        idx = np.zeros(graph.compact_num_links + 1, dtype=np.uint32)
+        data = np.zeros(len(filtered), dtype=np.uint32)
+
+        node_mapping = np.full(graph.num_nodes, -1, dtype=np.int32)
+
+        i = 0
+        for compressed_id, df in gb:
+            idx[compressed_id] = i
+            values = df.link_id.values
+            a = df.a_node.values
+            b = df.b_node.values
+
+            # In order to ensure that the link IDs come out in the correct order we must walk the links
+            # we do this assuming the `a` array is sorted.
+            j = 0
+            # Find the missing a_node, this is the starting of the chain. We cannot rely on the node ordering to do a simple lookup
+
+            a_node = x = a[np.isin(a, b, invert=True, assume_unique=True)][0]
+            while True:
+                tmp = a.searchsorted(x)
+                if tmp < len(a) and a[tmp] == x:
+                    x = b[tmp]
+                    data[i + j] = values[tmp]
+                else:
+                    break
+                j += 1
+
+            b_node = x
+            node_mapping[a_node] = graph.compact_graph["a_node"].iat[compressed_id]
+            node_mapping[b_node] = graph.compact_graph["b_node"].iat[compressed_id]
+
+            i += len(values)
+
+        idx[-1] = i
+
+        graph.compressed_link_network_mapping_idx = idx
+        graph.compressed_link_network_mapping_data = data
+        graph.network_compressed_node_mapping = node_mapping
+
+        return idx, data, node_mapping

aequilibrae/paths/results/assignment_results.py

@@ -1,10 +1,9 @@
-import dataclasses
 import multiprocessing as mp
 from abc import ABC, abstractmethod
 
 import numpy as np
 from aequilibrae.matrix import AequilibraeMatrix, AequilibraeData
-from aequilibrae.paths.graph import Graph, TransitGraph, GraphBase
+from aequilibrae.paths.graph import Graph, TransitGraph, GraphBase, _get_graph_to_network_mapping
 from aequilibrae.parameters import Parameters
 from aequilibrae import global_logger
 from pathlib import Path
@@ -22,14 +21,6 @@
 """
 
 
-@dataclasses.dataclass
-class NetworkGraphIndices:
-    network_ab_idx: np.array
-    network_ba_idx: np.array
-    graph_ab_idx: np.array
-    graph_ba_idx: np.array
-
-
 class AssignmentResultsBase(ABC):
     """Assignment results base class for traffic and transit assignments."""
 
@@ -249,15 +240,7 @@ def total_flows(self) -> None:
         sum_axis1(self.total_link_loads, self.link_loads, self.cores)
 
     def get_graph_to_network_mapping(self):
-        num_uncompressed_links = int(np.unique(self.lids).shape[0])
-        indexing = np.zeros(int(self.lids.max()) + 1, np.uint64)
-        indexing[np.unique(self.lids)[:]] = np.arange(num_uncompressed_links)
-
-        graph_ab_idx = self.direcs > 0
-        graph_ba_idx = self.direcs < 0
-        network_ab_idx = indexing[self.lids[graph_ab_idx]]
-        network_ba_idx = indexing[self.lids[graph_ba_idx]]
-        return NetworkGraphIndices(network_ab_idx, network_ba_idx, graph_ab_idx, graph_ba_idx)
+        return _get_graph_to_network_mapping(self.lids, self.direcs)
 
     def get_load_results(self) -> AequilibraeData:
         """