Repository: scipy/scipy
Pull request: ENH: sparse: add axis parameter to `count_nonzero` method
The sparse matrix/array method `count_nonzero()` mimics `np.count_nonzero` but currently does not support the axis parameter.  This PR adds support for the axis parameter in COO, CSR, CSC, and LIL formats. The other formats add the parameter but raise if not the default value `axis=None`. This mimics the support format pattern for an axis parameter in the private `_getnnz` method.

This PR sums duplicate entries as part of the computation, and it does not eliminate explicit zeros. For CSC, CSR formats, faster methods work if no explicit zeros are stored, so we check for explicit zeros before deciding which method to use. 

The incentive for this feature (see #19405) is a way to provide `axis` level nonzero counts recently lost during the move from `getnnz` to  a property `nnz`. That dropped feature counted stored values while this counts nonzero values.

Tests and docs are added.

Closes: #19405


Problem statement:
ENH: sparse.csr_array: keep axis functionality in `getnnz()`
### Is your feature request related to a problem? Please describe.

I require this functionality and don't see an efficient way to hack it.
It is planned to disappear in version 1-13...

### Describe the solution you'd like.

have it work like before

### Describe alternatives you've considered.

_No response_

### Additional context (e.g. screenshots, GIFs)

_No response_

Hints from issue discussion:
One way to get the results of the axis parameter from `A.getnnz` is:
```python
# for axis=0
np.array([vec.nnz for vec in A.T])
# for axis=1
np.array([vec.nnz for vec in A])
```
Here's some code to test this:
```python
import numpy as np
import scipy as sp
A = sp.sparse.csr_array([[4, 5, 0], [7, 0, 0]])
A.getnnz(axis=0)  # -> np.array([2, 1, 0])
A.getnnz(axis=1)  # -> np.array([2, 1], dtype=np.int32)

# for axis=0
np.array([vec.nnz for vec in A.T])  # -> np.array([2, 1, 0])
# for axis=1
np.array([vec.nnz for vec in A])  # -> np.array([2, 1])
```
Hi [dschuldt](https://github.com/dschult)! I agree that this provides the same functionality, but it will be much less efficient. The following is a comparison of the runtime of the two:
![759ffeb8-75e5-47ed-a854-1c0099d2320a](https://github.com/scipy/scipy/assets/79711833/699e1e26-bc03-4918-b474-3b61ca4c4bfc)
```
A = csr_array(np.ones((1000,1000)))
npstat = []
scstat = []
for i in range(100):
    start=time()
    A.getnnz(axis=0)  # -> np.array([2, 1, 0])
    A.getnnz(axis=1)  # -> np.array([2, 1], dtype=np.int32)
    scstat.append(time()-start)
    start=time()
    np.array([vec.nnz for vec in A.T])  # -> np.array([2, 1, 0])
    np.array([vec.nnz for vec in A])  # -> np.array([2, 1])
    npstat.append(time()-start)

plt.hist(npstat,color='b', alpha=0.5,label='comprrehended nnz')
plt.hist(scstat,color='r', alpha=0.5,label='getnnz(axis=0)')
plt.xlabel('runtime')
plt.ylabel('samples')
plt.title('runtime of 100 samples')
plt.legend()```

Good point!   
If you are willing to give up having it work for all sparse array formats, the following works for `csr` format and is just as fast or very slightly faster in my trials than the original function.  A similar (swapped) version would work for `csc`.

```python
A = sp.sparse.csr_array(np.ones((500,1000)))
assert (A.getnnz(axis=0) == np.bincount(A.indices, minlength=A.shape[1])).all()
assert (A.getnnz(axis=1) == np.diff(A.indptr)).all()

B = sp.sparse.csr_array([[4, 5, 0], [7, 0, 0]])
# for axis=0
A0 = np.bincount(B.indices, minlength=B.shape[1])  # -> np.array([2, 1, 0])
# for axis=1
A1 = np.diff(B.indptr)       # -> np.array([2, 1])

# both are 4.56ms on my machine
%timeit A.getnnz(axis=0)
%timeit np.bincount(A.indices, minlength=A.shape[1])

# getnnz is 15.8µs,  np.diff is 9.9µs
%timeit A.getnnz(axis=1)
%timeit np.diff(A.indptr)
```

Of course, it is also not as readable in my opinion. I'm not saying it is better... just a possible replacement.

I'm pretty sure that a reasonable and commonly occuring use-case would be enough to make the developers include this feature somehow.  We're looking for how many people use the `axis` keyword and what the use-case is. What are you using `axis=` for?
Thank you! I will use the `np.diff(A.indptr)`!

Commit message: fix downcast for bincount in COO

Patch to review:
diff --git a/scipy/sparse/_base.py b/scipy/sparse/_base.py
index 510f7565065d..87a2dad8dba1 100644
--- a/scipy/sparse/_base.py
+++ b/scipy/sparse/_base.py
@@ -263,7 +263,7 @@ def _getmaxprint(self):
         """Maximum number of elements to display when printed."""
         return self.maxprint
 
-    def count_nonzero(self):
+    def count_nonzero(self, axis=None):
         """Number of non-zero entries, equivalent to
 
         np.count_nonzero(a.toarray())
@@ -271,18 +271,63 @@ def count_nonzero(self):
         Unlike the nnz property, which return the number of stored
         entries (the length of the data attribute), this method counts the
         actual number of non-zero entries in data.
+
+        Duplicate entries are summed before counting.
+
+        Parameters
+        ----------
+        axis : {-2, -1, 0, 1, None} optional
+            Count nonzeros for the whole array, or along a specified axis.
+
+            .. versionadded:: 1.15.0
+
+        Returns
+        -------
+        numpy array
+            A reduced array (no axis `axis`) holding the number of nonzero values
+            for each of the indices of the nonaxis dimensions.
+
+        Notes
+        -----
+        If you want to count nonzero and explicit zero stored values (e.g. nnz)
+        along an axis, two fast idioms are provided by `numpy` functions for the
+        common CSR, CSC, COO formats.
+
+        For the major axis in CSR (rows) and CSC (cols) use `np.diff`:
+
+            >>> import numpy as np
+            >>> import scipy as sp
+            >>> A = sp.sparse.csr_array([[4, 5, 0], [7, 0, 0]])
+            >>> major_axis_stored_values = np.diff(A.indptr)  # -> np.array([2, 1])
+
+        For the minor axis in CSR (cols) and CSC (rows) use `numpy.bincount` with
+        minlength ``A.shape[1]`` for CSR and ``A.shape[0]`` for CSC:
+
+            >>> csr_minor_stored_values = np.bincount(A.indices, minlength=A.shape[1])
+
+        For COO, use the minor axis approach for either `axis`:
+
+            >>> A = A.tocoo()
+            >>> coo_axis0_stored_values = np.bincount(A.coords[0], minlength=A.shape[1])
+            >>> coo_axis1_stored_values = np.bincount(A.coords[1], minlength=A.shape[0])
+
+        Examples
+        --------
+
+            >>> A = sp.sparse.csr_array([[4, 5, 0], [7, 0, 0]])
+            >>> A.count_nonzero(axis=0)
+            array([2, 1, 0])
         """
-        raise NotImplementedError("count_nonzero not implemented for %s." %
-                                  self.__class__.__name__)
+        clsname = self.__class__.__name__
+        raise NotImplementedError(f"count_nonzero not implemented for {clsname}.")
 
     def _getnnz(self, axis=None):
         """Number of stored values, including explicit zeros.
 
         Parameters
         ----------
-        axis : None, 0, or 1
-            Select between the number of values across the whole array, in
-            each column, or in each row.
+        axis : {-2, -1, 0, 1, None} optional
+            Report stored values for the whole array, or along a specified axis.
 
         See also
         --------
diff --git a/scipy/sparse/_bsr.py b/scipy/sparse/_bsr.py
index 6a8a1be7dabe..dc7fb42ec8f1 100644
--- a/scipy/sparse/_bsr.py
+++ b/scipy/sparse/_bsr.py
@@ -219,6 +219,15 @@ def _getnnz(self, axis=None):
 
     _getnnz.__doc__ = _spbase._getnnz.__doc__
 
+    def count_nonzero(self, axis=None):
+        if axis is not None:
+            raise NotImplementedError(
+                "count_nonzero over axis is not implemented for BSR format."
+            )
+        return np.count_nonzero(self._deduped_data())
+
+    count_nonzero.__doc__ = _spbase.count_nonzero.__doc__
+
     def __repr__(self):
         _, fmt = _formats[self.format]
         sparse_cls = 'array' if isinstance(self, sparray) else 'matrix'
diff --git a/scipy/sparse/_compressed.py b/scipy/sparse/_compressed.py
index a86c09fc5edf..054393c2ba59 100644
--- a/scipy/sparse/_compressed.py
+++ b/scipy/sparse/_compressed.py
@@ -2,6 +2,7 @@
 __all__ = []
 
 from warnings import warn
+import itertools
 import operator
 
 import numpy as np
@@ -124,14 +125,42 @@ def _getnnz(self, axis=None):
             axis, _ = self._swap((axis, 1 - axis))
             _, N = self._swap(self.shape)
             if axis == 0:
-                return np.bincount(downcast_intp_index(self.indices),
-                                   minlength=N)
+                return np.bincount(downcast_intp_index(self.indices), minlength=N)
             elif axis == 1:
                 return np.diff(self.indptr)
             raise ValueError('axis out of bounds')
 
     _getnnz.__doc__ = _spbase._getnnz.__doc__
 
+    def count_nonzero(self, axis=None):
+        self.sum_duplicates()
+        if axis is None:
+            return np.count_nonzero(self.data)
+
+        if self.ndim == 1:
+            if axis not in (0, -1):
+                raise ValueError('axis out of bounds')
+            return np.count_nonzero(self.data)
+
+        if axis < 0:
+            axis += 2
+        axis, _ = self._swap((axis, 1 - axis))
+        if axis == 0:
+            _, N = self._swap(self.shape)
+            mask = self.data != 0
+            idx = self.indices if mask.all() else self.indices[mask]
+            return np.bincount(downcast_intp_index(idx), minlength=N)
+        elif axis == 1:
+            if self.data.all():
+                return np.diff(self.indptr)
+            x, y = itertools.tee(self.indptr)
+            next(y)
+            return np.array([np.count_nonzero(self.data[i:j]) for i, j in zip(x, y)])
+        else:
+            raise ValueError('axis out of bounds')
+
+    count_nonzero.__doc__ = _spbase.count_nonzero.__doc__
+
     def check_format(self, full_check=True):
         """Check whether the array/matrix respects the CSR or CSC format.
 
diff --git a/scipy/sparse/_coo.py b/scipy/sparse/_coo.py
index 2f4580b3f004..d6a124dffa69 100644
--- a/scipy/sparse/_coo.py
+++ b/scipy/sparse/_coo.py
@@ -182,6 +182,21 @@ def _getnnz(self, axis=None):
 
     _getnnz.__doc__ = _spbase._getnnz.__doc__
 
+    def count_nonzero(self, axis=None):
+        self.sum_duplicates()
+        if axis is None:
+            return np.count_nonzero(self.data)
+
+        if axis < 0:
+            axis += self.ndim
+        if axis >= self.ndim:
+            raise ValueError('axis out of bounds')
+        mask = self.data != 0
+        coord = self.coords[1 - axis][mask]
+        return np.bincount(downcast_intp_index(coord), minlength=self.shape[1 - axis])
+
+    count_nonzero.__doc__ = _spbase.count_nonzero.__doc__
+
     def _check(self):
         """ Checks data structure for consistency """
         if self.ndim != len(self.coords):
diff --git a/scipy/sparse/_data.py b/scipy/sparse/_data.py
index b619eeb90bbb..fe1e34a1da74 100644
--- a/scipy/sparse/_data.py
+++ b/scipy/sparse/_data.py
@@ -97,11 +97,6 @@ def copy(self):
 
     copy.__doc__ = _spbase.copy.__doc__
 
-    def count_nonzero(self):
-        return np.count_nonzero(self._deduped_data())
-
-    count_nonzero.__doc__ = _spbase.count_nonzero.__doc__
-
     def power(self, n, dtype=None):
         """
         This function performs element-wise power.
diff --git a/scipy/sparse/_dia.py b/scipy/sparse/_dia.py
index ab6d5dcdccad..75fdcef45e37 100644
--- a/scipy/sparse/_dia.py
+++ b/scipy/sparse/_dia.py
@@ -115,10 +115,16 @@ def _data_mask(self):
         mask &= (offset_inds < num_cols)
         return mask
 
-    def count_nonzero(self):
+    def count_nonzero(self, axis=None):
+        if axis is not None:
+            raise NotImplementedError(
+                "count_nonzero over an axis is not implemented for DIA format"
+            )
         mask = self._data_mask()
         return np.count_nonzero(self.data[mask])
 
+    count_nonzero.__doc__ = _spbase.count_nonzero.__doc__
+
     def _getnnz(self, axis=None):
         if axis is not None:
             raise NotImplementedError("_getnnz over an axis is not implemented "
@@ -133,7 +139,6 @@ def _getnnz(self, axis=None):
         return int(nnz)
 
     _getnnz.__doc__ = _spbase._getnnz.__doc__
-    count_nonzero.__doc__ = _spbase.count_nonzero.__doc__
 
     def sum(self, axis=None, dtype=None, out=None):
         validateaxis(axis)
diff --git a/scipy/sparse/_dok.py b/scipy/sparse/_dok.py
index 08a039136ff3..16ab11c10689 100644
--- a/scipy/sparse/_dok.py
+++ b/scipy/sparse/_dok.py
@@ -69,7 +69,11 @@ def _getnnz(self, axis=None):
             )
         return len(self._dict)
 
-    def count_nonzero(self):
+    def count_nonzero(self, axis=None):
+        if axis is not None:
+            raise NotImplementedError(
+                "count_nonzero over an axis is not implemented for DOK format."
+            )
         return sum(x != 0 for x in self.values())
 
     _getnnz.__doc__ = _spbase._getnnz.__doc__
diff --git a/scipy/sparse/_lil.py b/scipy/sparse/_lil.py
index 2503aa628b58..b0399998d3ae 100644
--- a/scipy/sparse/_lil.py
+++ b/scipy/sparse/_lil.py
@@ -106,10 +106,27 @@ def _getnnz(self, axis=None):
         else:
             raise ValueError('axis out of bounds')
 
-    def count_nonzero(self):
-        return sum(np.count_nonzero(rowvals) for rowvals in self.data)
-
     _getnnz.__doc__ = _spbase._getnnz.__doc__
+
+    def count_nonzero(self, axis=None):
+        if axis is None:
+            return sum(np.count_nonzero(rowvals) for rowvals in self.data)
+
+        if axis < 0:
+            axis += 2
+        if axis == 0:
+            out = np.zeros(self.shape[1], dtype=np.intp)
+            for row, data in zip(self.rows, self.data):
+                mask = [c for c, d in zip(row, data) if d != 0]
+                out[mask] += 1
+            return out
+        elif axis == 1:
+            return np.array(
+                [np.count_nonzero(rowvals) for rowvals in self.data], dtype=np.intp,
+            )
+        else:
+            raise ValueError('axis out of bounds')
+
     count_nonzero.__doc__ = _spbase.count_nonzero.__doc__
 
     def getrowview(self, i):
diff --git a/scipy/sparse/tests/test_base.py b/scipy/sparse/tests/test_base.py
index 1aa76b8ff26e..1ee5c40aa1e4 100644
--- a/scipy/sparse/tests/test_base.py
+++ b/scipy/sparse/tests/test_base.py
@@ -633,11 +633,21 @@ def test_empty(self):
         assert_equal(self.spcreator((3, 3)).toarray(), zeros((3, 3)))
         assert_equal(self.spcreator((3, 3)).nnz, 0)
         assert_equal(self.spcreator((3, 3)).count_nonzero(), 0)
+        if self.datsp.format in ["coo", "csr", "csc", "lil"]:
+            assert_equal(self.spcreator((3, 3)).count_nonzero(axis=0), array([0, 0, 0]))
 
     def test_count_nonzero(self):
-        expected = np.count_nonzero(self.datsp.toarray())
-        assert_equal(self.datsp.count_nonzero(), expected)
-        assert_equal(self.datsp.T.count_nonzero(), expected)
+        axis_support = self.datsp.format in ["coo", "csr", "csc", "lil"]
+        axes = [None, 0, 1, -1, -2] if axis_support else [None]
+
+        for A in (self.datsp, self.datsp.T):
+            for ax in axes:
+                expected = np.count_nonzero(A.toarray(), axis=ax)
+                assert_equal(A.count_nonzero(axis=ax), expected)
+
+        if not axis_support:
+            with assert_raises(NotImplementedError, match="not implemented .* format"):
+                self.datsp.count_nonzero(axis=0)
 
     def test_invalid_shapes(self):
         assert_raises(ValueError, self.spcreator, (-1,3))


Code review comments:
diff_hunk: @@ -124,14 +125,42 @@ def _getnnz(self, axis=None):
             axis, _ = self._swap((axis, 1 - axis))
             _, N = self._swap(self.shape)
             if axis == 0:
-                return np.bincount(downcast_intp_index(self.indices),
-                                   minlength=N)
+                return np.bincount(downcast_intp_index(self.indices), minlength=N)
             elif axis == 1:
                 return np.diff(self.indptr)
             raise ValueError('axis out of bounds')
 
     _getnnz.__doc__ = _spbase._getnnz.__doc__
 
+    def count_nonzero(self, axis=None):
+        self.sum_duplicates()
+        if axis is None:
+            return np.count_nonzero(self.data)
+
+        if self.ndim == 1:
+            if axis not in (0, -1):
+                raise ValueError('axis out of bounds')
+            return np.count_nonzero(self.data)
+
+        if axis < 0:
+            axis += 2
+        axis, _ = self._swap((axis, 1 - axis))
+        if axis == 0:
+            _, N = self._swap(self.shape)
+            mask = self.data != 0
+            idx = self.indices if mask.all() else self.indices[mask]
+            return np.bincount(downcast_intp_index(idx), minlength=N)
+        elif axis == 1:
+            if self.data.all():
+                return np.diff(self.indptr)
+            x, y = itertools.tee(self.indptr)
path: scipy/sparse/_compressed.py
text: @user1:
https://docs.python.org/3/library/itertools.html#itertools.pairwise

@author:
The python docs say `itertools.pairwise is only for Python 3.10 and later (we use it a lot in networkx, though we built our own long ago).  Scipy supports 3.9 still. (will that change soon?)

We could create something in  sputils.py  and drop it after 3.9 is dropped. But that would be more trouble than this code (I think).  

@author:
Ahh!  I read the page wrong -- Scipy supports 3.10+ now. So `itertools.pairwise` will work!  Thanks!

diff_hunk: @@ -263,26 +263,71 @@ def _getmaxprint(self):
         """Maximum number of elements to display when printed."""
         return self.maxprint
 
-    def count_nonzero(self):
+    def count_nonzero(self, axis=None):
         """Number of non-zero entries, equivalent to
 
         np.count_nonzero(a.toarray())
path: scipy/sparse/_base.py
text: @user1:
, axis=axis

diff_hunk: @@ -182,6 +182,21 @@ def _getnnz(self, axis=None):
 
     _getnnz.__doc__ = _spbase._getnnz.__doc__
 
+    def count_nonzero(self, axis=None):
+        self.sum_duplicates()
+        if axis is None:
+            return np.count_nonzero(self.data)
+
+        if axis < 0:
+            axis += self.ndim
+        if axis >= self.ndim:
path: scipy/sparse/_coo.py
text: @user1:
or axis < 0