diff --git a/lenskit/util/csmatrix.pyi b/lenskit/util/csmatrix.pyi index 7fbabcc80..33c4b702e 100644 --- a/lenskit/util/csmatrix.pyi +++ b/lenskit/util/csmatrix.pyi @@ -1,5 +1,6 @@ import numpy as np import numpy.typing as npt +from scipy.sparse import csr_array class CSMatrix: nrows: int @@ -19,3 +20,5 @@ class CSMatrix: vs: npt.NDArray[np.float64], ): ... def row_ep(self, row: int) -> tuple[int, int]: ... + @staticmethod + def from_scipy(matrix: csr_array) -> CSMatrix: ... diff --git a/lenskit/util/csmatrix.pyx b/lenskit/util/csmatrix.pyx index 3e236cdf4..a0ac84ad4 100644 --- a/lenskit/util/csmatrix.pyx +++ b/lenskit/util/csmatrix.pyx @@ -14,6 +14,12 @@ cdef class CSMatrix: self.values = vs self.nnz = self.rowptr[nr] + @staticmethod + def from_scipy(m): + nr, nc = m.shape + + return CSMatrix(nr, nc, m.indptr, m.indices, m.data) + cpdef (int,int) row_ep(self, row): if row < 0 or row >= self.nrows: raise IndexError(f"invalid row {row} for {self.nrows}x{self.ncols} matrix") diff --git a/tests/test_csmatrix.py b/tests/test_csmatrix.py index 1a25a124e..3de548382 100644 --- a/tests/test_csmatrix.py +++ b/tests/test_csmatrix.py @@ -42,12 +42,21 @@ def test_init_matrix(m: sps.csr_array): assert m2.nnz == m.nnz +@given(sparse_matrices()) +def test_from_scipy(m: sps.csr_array): + print(m.shape, m.nnz, m.indptr.dtype, m.indices.dtype) + m2 = CSMatrix.from_scipy(m) + + assert m2.nrows == m.shape[0] + assert m2.ncols == m.shape[1] + assert m2.nnz == m.nnz + + @given(sparse_matrices()) def test_csm_row_ep(m: sps.csr_array): - nr, nc = m.shape - m2 = CSMatrix(nr, nc, m.indptr, m.indices, m.data) + m2 = CSMatrix.from_scipy(m) - for i in range(nr): + for i in range(m2.nrows): sp, ep = m2.row_ep(i) assert sp == m2.rowptr[i] assert ep == m2.rowptr[i + 1]