Fix cov and add tests

Author: asistradition

scself/tests/test_correlation.py

@@ -0,0 +1,48 @@
+import unittest
+
+import numpy as np
+import scipy.sparse as sps
+
+from scself.utils.correlation import cov
+
+X = np.random.default_rng(100).random((500, 100))
+
+
+class TestCOV(unittest.TestCase):
+
+    def test_dense(self):
+
+        x_cov = cov(X)
+        np.testing.assert_almost_equal(
+            x_cov,
+            np.cov(X.T),
+            decimal=4
+        )
+
+    def test_dense_axis1(self):
+
+        x_cov = cov(X, axis=1)
+        np.testing.assert_almost_equal(
+            x_cov,
+            np.cov(X),
+            decimal=4
+        )
+
+    def test_sparse(self):
+
+        x_cov = cov(sps.csr_array(X))
+        np.testing.assert_almost_equal(
+            x_cov,
+            np.cov(X.T),
+            decimal=4
+        )
+
+
+    def test_sparse_axis1(self):
+
+        x_cov = cov(sps.csr_array(X), axis=1)
+        np.testing.assert_almost_equal(
+            x_cov,
+            np.cov(X),
+            decimal=4
+        )

scself/utils/correlation.py

@@ -25,34 +25,46 @@ def fit_transform(self, X):
 def cov(X, axis=0):
 
     if sps.issparse(X):
-        return cov_sparse(X)
+        return cov_sparse(X, axis=axis)
 
     # Center and get num rows
     avg, w_sum = np.average(X, axis=axis, weights=None, returned=True)
     w_sum = w_sum[0]
     X = X - (avg[None, :] if axis == 0 else avg[:, None])
 
     # Gram matrix
-    X = np.dot(X.T, X)
-    X *= np.true_divide(1, w_sum)
+    if axis == 0:
+        X = np.dot(X.T, X)
+    else:
+        X = np.dot(X, X.T)
+
+    X *= np.true_divide(1, w_sum - 1)
 
     return X
 
 def cov_sparse(X, axis=0):
 
-    avg = X.mean(axis)
+    axsum = X.sum(axis)
+    w_sum = X.shape[axis]
 
     # for spmatrix & sparray
     try:
-        avg = avg.A1
+        axsum = axsum.A1
     except AttributeError:
-        avg = avg.ravel()
+        axsum = axsum.ravel()
 
-    w_sum = X.shape[axis]
-    X = dot(X.T, X, dense=True)
-    X *= np.true_divide(1, w_sum)
+    axsum = axsum.reshape(-1, 1).dot(axsum.reshape(1, -1)) 
+    axsum /= w_sum
 
-    return X
+    if axis == 0:
+        X_cov = dot(X.T, X, dense=True)
+    else:
+        X_cov = dot(X, X.T, dense=True)
+
+    X_cov -= axsum
+    X_cov /= (w_sum - 1)
+
+    return X_cov
 
 def corrcoef(X, axis=0):