Merge pull request #441 from Tal-Golan/row_col_indicator_g_sparse

A sparse version of  row_col_indicator_g

Author: JasperVanDenBosch

src/rsatoolbox/rdm/compare.py

@@ -13,7 +13,7 @@
 from rsatoolbox.util.matrix import pairwise_contrast
 from rsatoolbox.util.rdm_utils import _get_n_from_reduced_vectors
 from rsatoolbox.util.rdm_utils import _get_n_from_length
-from rsatoolbox.util.matrix import row_col_indicator_g
+from rsatoolbox.util.matrix import row_col_indicator_g, row_col_indicator_g_sparse
 from rsatoolbox.util.rdm_utils import batch_to_matrices
 
 
@@ -446,7 +446,11 @@ def _cov_weighting(vector, nan_idx, sigma_k=None):
     N, n_dist = vector.shape
     n_cond = _get_n_from_length(nan_idx.shape[0])
     vector_w = -0.5 * np.c_[vector, np.zeros((N, n_cond))]
-    rowI, colI = row_col_indicator_g(n_cond)
+    SPARSE_THRESHOLD = 100 # threshold for switching to sparse matrices
+    if n_cond >= SPARSE_THRESHOLD:
+        rowI, colI = row_col_indicator_g_sparse(n_cond) # use sparse indicator matrices
+    else:
+        rowI, colI = row_col_indicator_g(n_cond) # use dense indicator matrices
     sumI = rowI + colI
     if np.all(nan_idx):
         # column and row means

src/rsatoolbox/util/matrix.py

@@ -159,6 +159,46 @@ def row_col_indicator_g(n_cond):
     np.fill_diagonal(col_i[-n_cond:, :], 1)
     return (row_i, col_i)
 
+def row_col_indicator_g_sparse(n_cond, dtype=np.int8, index_dtype=np.int64):
+    """
+    Generates row and column indicator matrices as sparse matrices (CSR format)
+    for a vectorized second moment matrix.
+    The vectorized version has the off-diagonal elements first, then appends the diagonal.
+
+    Args:
+        n_cond (int): Number of conditions underlying the second moment.
+        dtype (numpy.dtype, optional): Data type for the non-zero elements
+                                       of the sparse matrices. Defaults to np.int8.
+        index_dtype (numpy.dtype, optional): Data type for the indices in the sparse matrix.
+                                              Defaults to np.int64.
+
+    Returns:
+        tuple: (row_indicator_sparse, col_indicator_sparse)
+            row_indicator_sparse (scipy.sparse.csr_matrix): Sparse indicator matrix for rows.
+            col_indicator_sparse (scipy.sparse.csr_matrix): Sparse indicator matrix for columns.
+    """
+    num_off_diag = n_cond * (n_cond - 1) // 2
+    n_elem = num_off_diag + n_cond
+    shape = (n_elem, n_cond)
+    if n_elem == 0:
+        return (csr_matrix(shape, dtype=dtype),
+                csr_matrix(shape, dtype=dtype))
+    sparse_data = np.ones(n_elem, dtype=dtype)
+
+    sparse_row_coords = np.arange(n_elem, dtype=index_dtype)
+    row_i_col_coords = np.empty(n_elem, dtype=index_dtype)
+    col_i_col_coords = np.empty(n_elem, dtype=index_dtype)
+    if num_off_diag > 0:
+        upper_p, upper_q = np.triu_indices(n_cond, k=1)
+        row_i_col_coords[:num_off_diag] = upper_p.astype(index_dtype)
+        col_i_col_coords[:num_off_diag] = upper_q.astype(index_dtype)
+    if n_cond > 0:
+        diag_indices_in_matrix = np.arange(n_cond, dtype=index_dtype)
+        row_i_col_coords[num_off_diag:] = diag_indices_in_matrix
+        col_i_col_coords[num_off_diag:] = diag_indices_in_matrix
+    row_i_sparse = csr_matrix((sparse_data, (sparse_row_coords, row_i_col_coords)), shape=shape)
+    col_i_sparse = csr_matrix((sparse_data, (sparse_row_coords, col_i_col_coords)), shape=shape)
+    return (row_i_sparse, col_i_sparse)
 
 def run() -> spmatrix:
     a = csr_matrix(np.eye(3))

tests/test_util_matrix.py

@@ -45,6 +45,13 @@ def test_centering(self):
         self.assertEqual(n_row, 10)
         self.assertEqual(n_col, 10)
 
+    def test_row_col_indicator_g(self):
+        """ Tests the equivalence of the sparse and dense versions of row_col_indicator_g """
+        for n_cond in [0, 1, 2, 3, 4, 5, 10, 20, 50, 100]:
+            rowI_sparse, colI_sparse = rsu.matrix.row_col_indicator_g_sparse(n_cond)
+            rowI_dense, colI_dense = rsu.matrix.row_col_indicator_g(n_cond)
+            self.assertTrue(np.array_equal(rowI_sparse.toarray(), rowI_dense))
+            self.assertTrue(np.array_equal(colI_sparse.toarray(), colI_dense))
 
 if __name__ == '__main__':
     unittest.main()