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):
@@ -56,34 +55,52 @@ def bincount(X, max_val=None, weights=None, minlength=None):
 
 def countnans(X, weights=None, axis=None, dtype=None, keepdims=False):
     """
-    Count the undefined elements in arr along given axis.
+    Count the undefined elements in an array along given axis.
 
     Parameters
     ----------
     X : array_like
     weights : array_like
         Weights to weight the nans with, before or after counting (depending
         on the weights shape).
+    axis : Optional[int]
 
     Returns
     -------
-    counts
+    Union[np.ndarray, float]
+
     """
     if not sp.issparse(X):
         X = np.asanyarray(X)
         isnan = np.isnan(X)
         if weights is not None and weights.shape == X.shape:
             isnan = isnan * weights
+
+        # In order to keep return types consistent with sparse vectors, we will
+        # handle `axis=1` given a regular 1d numpy array equivallently as
+        # `axis=0`. If we didn't do this, this would raise error, whereas the
+        # sparse counterpart would return the appropriate value.
+        axis = axis and min(axis, isnan.ndim - 1)
+        
         counts = isnan.sum(axis=axis, dtype=dtype, keepdims=keepdims)
         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)
+
+        # We want a scalar value if `axis=None` or if the sparse matrix is
+        # actually a vector (e.g. [[1 2 3]]), but has `ndim=2` due to scipy
+        # implementation
+        if axis is None or X.shape[0] == 1:
+            counts = counts.sum()
 
-        counts = _count_nans_per_row_sparse(X, weights)
     return counts
 
 

Orange/tests/test_statistics.py

@@ -31,10 +31,6 @@ def test_bincount(self):
         self.assertEqual(n_nans, 0)
         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])
-
     def test_contingency(self):
         x = np.array([0, 1, 0, 2, np.nan])
         y = np.array([0, 0, 1, np.nan, 0])
@@ -241,3 +237,71 @@ def test_var(self):
                     var(csr_matrix(data), axis=axis),
                     np.var(data, axis=axis)
                 )
+
+
+class TestCountnans(unittest.TestCase):
+    def test_countnans_on_sparse_array(self):
+        x = csr_matrix([0, 1, 0, 2, 2, np.nan, 1, np.nan, 0, 1])
+        self.assertEqual(countnans(x), 2)
+
+    def test_countnans_on_dense_array(self):
+        x = np.array([0, 1, 0, 2, 2, np.nan, 1, np.nan, 0, 1])
+        self.assertEqual(countnans(x), 2)
+
+    def test_shape_matches_dense_and_sparse_given_array_and_axis_None(self):
+        dense = np.array([0, 1, 0, 2, 2, np.nan, 1, np.nan, 0, 1])
+        sparse = csr_matrix(dense)
+        np.testing.assert_equal(countnans(dense), countnans(sparse))
+        self.assertEqual(countnans(dense), 2)
+
+    def test_shape_matches_dense_and_sparse_given_array_and_axis_0(self):
+        dense = np.array([0, 1, 0, 2, 2, np.nan, 1, np.nan, 0, 1])
+        sparse = csr_matrix(dense)
+        np.testing.assert_equal(countnans(dense, axis=0), countnans(sparse, axis=0))
+        self.assertEqual(countnans(dense, axis=0), 2)
+
+    def test_shape_matches_dense_and_sparse_given_array_and_axis_1(self):
+        dense = np.array([0, 1, 0, 2, 2, np.nan, 1, np.nan, 0, 1])
+        sparse = csr_matrix(dense)
+        np.testing.assert_equal(countnans(dense, axis=1), countnans(sparse, axis=1))
+        self.assertEqual(countnans(dense, axis=1), 2)
+
+    def test_countnans(self):
+        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])