Add sparse matrix tests and from disk method

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.pyx

@@ -33,6 +33,24 @@ cdef class Sparse:
         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):
     """
@@ -67,6 +85,9 @@ cdef class COO(Sparse):
         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]
@@ -76,6 +97,25 @@ cdef class COO(Sparse):
 
         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)

tests/aequilibrae/matrix/test_sparse_matrix.py

@@ -0,0 +1,38 @@
+from tempfile import gettempdir
+from aequilibrae.matrix import COO
+from unittest import TestCase
+from uuid import uuid4
+import scipy.sparse
+import numpy as np
+import pathlib
+
+
+class TestSparseMatrix(TestCase):
+    def setUp(self) -> None:
+        self.data = np.full((100, 100), 5.0)
+        self.dir = pathlib.Path(gettempdir()) / uuid4().hex
+        self.dir.mkdir()
+
+    def tearDown(self) -> None:
+        pass
+
+    def test_round_trip(self):
+        p = self.dir / "test.omx"
+
+        coo = COO.from_matrix(
+            self.data,
+        )
+        coo.to_disk(p, "m1")
+        coo.to_disk(p, "m2")
+
+        sp = coo.to_scipy()
+
+        coo1 = COO.from_disk(p)
+        coo2 = COO.from_disk(p, aeq=True)
+
+        for m in ["m1", "m2"]:
+            self.assertIsInstance(coo1[m], scipy.sparse.csr_matrix)
+            self.assertIsInstance(coo2[m], COO)
+
+            np.testing.assert_allclose(sp.A, coo1[m].A)
+            np.testing.assert_allclose(sp.A, coo2[m].to_scipy().A)