Statistics.countnans: Fix sparse implementation and add axis support

Author: pavlin-policar

Orange/statistics/util.py

@@ -12,13 +12,12 @@
 
 def _count_nans_per_row_sparse(X, weights):
     """ Count the number of nans (undefined) values per row. """
-    items_per_row = 1 if X.ndim == 1 else X.shape[1]
-    counts = np.ones(X.shape[0]) * items_per_row
-    nnz_per_row = np.bincount(X.indices, minlength=len(counts))
-    counts -= nnz_per_row
+    counts = np.fromiter((np.isnan(row.data).sum() for row in X), dtype=np.float)
+
     if weights is not None:
         counts *= weights
-    return np.sum(counts)
+
+    return counts
 
 
 def bincount(X, max_val=None, weights=None, minlength=None):
@@ -54,7 +53,7 @@ def bincount(X, max_val=None, weights=None, minlength=None):
     return bc, nans
 
 
-def countnans(X, weights=None, axis=None, dtype=None, keepdims=False):
+def countnans(X, weights=None, axis=None):
     """
     Count the undefined elements in arr along given axis.
 
@@ -64,10 +63,12 @@ def countnans(X, weights=None, axis=None, dtype=None, keepdims=False):
     weights : array_like
         Weights to weight the nans with, before or after counting (depending
         on the weights shape).
+    axis : Optional[int]
 
     Returns
     -------
     counts
+    
     """
     if not sp.issparse(X):
         X = np.asanyarray(X)
@@ -78,12 +79,18 @@ def countnans(X, weights=None, axis=None, dtype=None, keepdims=False):
         if weights is not None and weights.shape != X.shape:
             counts = counts * weights
     else:
-        if any(attr is not None for attr in [axis, dtype]) or \
-                        keepdims is not False:
-            raise ValueError('Arguments axis, dtype and keepdims'
-                             'are not yet supported on sparse data!')
+        assert axis in [None, 0, 1], 'Only axis 0 and 1 are currently supported'
+        arr = X if axis == 1 else X.T
+
+        if weights is not None:
+            weights = weights if axis == 1 else weights.T
+
+        arr = arr.tocsr()
+        counts = _count_nans_per_row_sparse(arr, weights)
+
+        if axis is None:
+            counts = counts.sum()
 
-        counts = _count_nans_per_row_sparse(X, weights)
     return counts
 
 

Orange/tests/test_statistics.py

@@ -32,8 +32,44 @@ def test_bincount(self):
         np.testing.assert_equal(hist, [1, 1, 0, 1])
 
     def test_countnans(self):
-        np.testing.assert_equal(countnans([[1, np.nan],
-                                           [2, np.nan]], axis=0), [0, 2])
+        x = [[1, np.nan],
+             [2, np.nan]]
+        np.testing.assert_equal(
+            countnans(x), 2, 'Countnans fails on dense data')
+        np.testing.assert_equal(
+            countnans(csr_matrix(x)), 2, 'Countnans fails on sparse data.')
+
+    def test_countnans_columns(self):
+        x = [[1, np.nan],
+             [2, np.nan]]
+        np.testing.assert_equal(
+            countnans(x, axis=0), [0, 2],
+            'Countnans fails on dense data with `axis=0`')
+        np.testing.assert_equal(
+            countnans(csr_matrix(x), axis=0), [0, 2],
+            'Countnans fails on sparse data with `axis=0`')
+
+    def test_countnans_rows(self):
+        x = [[1, np.nan],
+             [2, np.nan]]
+        np.testing.assert_equal(
+            countnans(x, axis=1), [1, 1],
+            'Countnans fails on dense data with `axis=1`')
+        np.testing.assert_equal(
+            countnans(csr_matrix(x), axis=1), [1, 1],
+            'Countnans fails on sparse data with `axis=1`')
+
+    def test_countnans_weights(self):
+        x = [[1, np.nan],
+             [2, np.nan]]
+        w = np.array([[1, 1],
+                      [2, 2]])
+        np.testing.assert_equal(countnans(x, weights=w, axis=0), [0, 3])
+        np.testing.assert_equal(countnans(x, weights=w, axis=1), [1, 2])
+
+        w = np.array([1, 2])
+        np.testing.assert_equal(countnans(x, weights=w, axis=0), [0, 4])
+        np.testing.assert_equal(countnans(x, weights=w, axis=1), [1, 2])
 
     def test_contingency(self):
         x = np.array([0, 1, 0, 2, np.nan])