Row with known nnz

Author: cwindolf

src/dartsort/cluster/gaussian_mixture.py

@@ -541,19 +541,21 @@ def log_likelihoods(
                 unit_neighb_info.append((j, neighbs, ns_unit))
             else:
                 assert previous_logliks is not None
-                if hasattr(previous_logliks, 'row_nnz'):
+                if hasattr(previous_logliks, "row_nnz"):
                     rnnz = previous_logliks.row_nnz[j]
-                    row = csc_sparse_getrow(previous_logliks, j, rnnz).tocoo(copy=False)
+                    six, data = csc_sparse_getrow(previous_logliks, j, rnnz)
                 else:
                     row = previous_logliks[[j]].tocoo(copy=True)
-                six = row.coords[1]
-                ns_unit = row.nnz
+                    six = row.coords[1]
+                    ns_unit = row.nnz
+                    data = row.data
+                ns_unit = len(six)
                 if "covered_neighbs" in unit.annotations:
                     covered_neighbs = unit.annotations["covered_neighbs"]
                 else:
                     covered_neighbs = full_core_neighborhoods.neighborhood_ids[six]
                     covered_neighbs = torch.unique(covered_neighbs)
-                unit_neighb_info.append((j, six, row.data, ns_unit))
+                unit_neighb_info.append((j, six, data, ns_unit))
             core_overlaps[covered_neighbs] += 1
             nnz += ns_unit
 

src/dartsort/util/sparse_util.py

@@ -184,27 +184,23 @@ def _csc_sparse_mask_rows(indices, indptr, data, oldrow_to_newrow, keep_mask):
 
 def csc_sparse_getrow(csc, row, rowcount):
     rowix_dtype = csc.indices.dtype
-    indptr_out = np.empty(rowcount + 1, dtype=rowix_dtype)
-    data_out = np.empty(rowcount, dtype=rowix_dtype)
-    indices_out = np.full(rowcount, row, dtype=rowix_dtype)
+    columns_out = np.empty(rowcount, dtype=rowix_dtype)
+    data_out = np.empty(rowcount, dtype=csc.data.dtype)
     _csc_sparse_getrow(
-        csc.indices, csc.indptr, csc.data, indptr_out, data_out, rowix_dtype.type(row)
+        csc.indices, csc.indptr, csc.data, columns_out, data_out, rowix_dtype.type(row)
     )
 
-    return csc_array(
-        (data_out, indices_out, indptr_out),
-        shape=(len(kept_row_inds), csc.shape[1]),
-    )
+    return columns_out, data_out
 
 
 sigs = [
-    "void(i8[::1], i8[::1], f4[::1], i8[::1], i8[::1], i8)",
-    "void(i4[::1], i4[::1], f4[::1], i4[::1], i4[::1], i4)",
+    "void(i8[::1], i8[::1], f4[::1], i8[::1], f4[::1], i8)",
+    "void(i4[::1], i4[::1], f4[::1], i4[::1], f4[::1], i4)",
 ]
 
 
 @numba.njit(sigs, error_model="numpy", nogil=True)
-def _csc_sparse_getrow(indices, indptr, data, indptr_out, data_out, the_row):
+def _csc_sparse_getrow(indices, indptr, data, columns_out, data_out, the_row):
     write_ix = 0
 
     column = 0
@@ -217,17 +213,11 @@ def _csc_sparse_getrow(indices, indptr, data, indptr_out, data_out, the_row):
 
         # write data for this sample
         data_out[write_ix] = data[read_ix]
-        write_ix += 1
-
         while read_ix >= column_end:
-            indptr_out[column + 1] = write_ix - 1
             column += 1
             column_end = indptr[column + 1]
-
-    while column < len(indptr) - 1:
-        indptr_out[column + 1] = write_ix
-        column += 1
-        column_end = indptr[column + 1]
+        columns_out[write_ix] = column
+        write_ix += 1
 
 
 # @numba.njit(sigs, error_model="numpy", nogil=True)

tests/test_sparse.py

@@ -80,6 +80,27 @@ def test_csc_mask():
         assert np.array_equal(x1.data, x0.data)
 
 
+def test_csc_getrow():
+    rg = np.random.default_rng(10)
+    ij = rg.integers(low=((0, 0),), high=(shape,), size=(nnz, 2))
+    ij = np.unique(ij, axis=0)
+    assert (np.diff(ij[:, 0]) >= 0).all()
+    assert not (np.diff(ij[:, 1]) >= 0).all()
+    vals = rg.normal(size=len(ij)).astype(np.float32)
+
+    x = coo_array((vals, ij.T), shape).tocsc()
+
+    for row in range(x.shape[0]):
+        x0 = x[[row]]
+        columns, data = sparse_util.csc_sparse_getrow(x, row, x0.nnz)
+
+        assert len(columns) == len(data) == x0.nnz
+        x0coo = x0.tocoo()
+        assert np.array_equal(columns, x0coo.coords[1])
+        assert np.array_equal(data, x0coo.data)
+
+
 if __name__ == "__main__":
+    test_csc_getrow()
     test_csc_insert()
     test_csc_mask()