Merge pull request #712 from unnonouno/sparse-sum

Implement sum for cupy.sparse

cupy/sparse/base.py

@@ -321,7 +321,52 @@ def set_shape(self, shape):
 
     # TODO(unno): Implement setdiag
 
-    # TODO(unno): Implement sum
+    def sum(self, axis=None, dtype=None, out=None):
+        """Sums the matrix elements over a given axis.
+
+        Args:
+            axis (int or ``None``): Axis along which the sum is comuted.
+                If it is ``None``, it computes the sum of all the elements.
+                Select from ``{None, 0, 1, -2, -1}``.
+            dtype: The type of returned matrix. If it is not specified, type
+                of the array is used.
+            out (cupy.ndarray): Output matrix.
+
+        Returns:
+            cupy.ndarray: Summed array.
+
+        .. seealso::
+           :func:`scipy.sparse.spmatrix.sum`
+
+        """
+        util.validateaxis(axis)
+
+        # This implementation uses multiplication, though it is not efficient
+        # for some matrix types. These should override this function.
+
+        m, n = self.shape
+
+        if axis is None:
+            return self.dot(cupy.ones(n, dtype=self.dtype)).sum(
+                dtype=dtype, out=out)
+
+        if axis < 0:
+            axis += 2
+
+        if axis == 0:
+            ret = self.T.dot(cupy.ones(m, dtype=self.dtype)).reshape(1, n)
+        else:  # axis == 1
+            ret = self.dot(cupy.ones(n, dtype=self.dtype)).reshape(m, 1)
+
+        if out is not None:
+            if out.shape != ret.shape:
+                raise ValueError('dimensions do not match')
+            out[:] = ret
+            return out
+        elif dtype is not None:
+            return ret.astype(dtype, copy=False)
+        else:
+            return ret
 
     def toarray(self, order=None, out=None):
         """Return a dense ndarray representation of this matrix."""

cupy/sparse/util.py

@@ -36,3 +36,21 @@ def isshape(x):
         return False
     m, n = x
     return isintlike(m) and isintlike(n)
+
+
+def validateaxis(axis):
+    if axis is not None:
+        axis_type = type(axis)
+
+        if axis_type == tuple:
+            raise TypeError(
+                'Tuples are not accepted for the \'axis\' '
+                'parameter. Please pass in one of the '
+                'following: {-2, -1, 0, 1, None}.')
+
+        if not cupy.issubdtype(cupy.dtype(axis_type), cupy.integer):
+            raise TypeError('axis must be an integer, not {name}'
+                            .format(name=axis_type.__name__))
+
+        if not (-2 <= axis <= 1):
+            raise ValueError('axis out of range')

tests/cupy_tests/sparse_tests/test_coo.py

@@ -725,6 +725,21 @@ def test_pow_neg(self, xp, sp):
         m = _make_square(xp, sp, self.dtype)
         m ** -1
 
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_tuple_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=(0, 1))
+
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_float_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=0.0)
+
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_too_large_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=3)
+
     @testing.numpy_cupy_allclose(sp_name='sp')
     def test_transpose(self, xp, sp):
         m = self.make(xp, sp, self.dtype)
@@ -743,6 +758,36 @@ def test_eliminate_zeros(self, xp, sp):
         return m.toarray()
 
 
+@testing.parameterize(*testing.product({
+    'dtype': [numpy.float32, numpy.float64],
+    'ret_dtype': [None, numpy.float32, numpy.float64],
+    'axis': [None, 0, 1, -1, -2],
+}))
+@unittest.skipUnless(scipy_available, 'requires scipy')
+class TestCooMatrixSum(unittest.TestCase):
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype)
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum_with_out(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        if self.axis is None:
+            shape = ()
+        else:
+            shape = list(m.shape)
+            shape[self.axis] = 1
+            shape = tuple(shape)
+        out = xp.empty(shape, dtype=self.ret_dtype)
+        if xp is numpy:
+            # TODO(unno): numpy.matrix is used for scipy.sparse though
+            # cupy.ndarray is used for cupy.sparse.
+            out = xp.asmatrix(out)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype, out=out)
+
+
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float32, numpy.float64],
 }))

tests/cupy_tests/sparse_tests/test_csc.py

@@ -707,6 +707,16 @@ def test_sort_indices(self, xp, sp):
         m.sort_indices()
         return m.toarray()
 
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_tuple_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=(0, 1))
+
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_too_large_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=3)
+
     @testing.numpy_cupy_allclose(sp_name='sp')
     def test_sum_duplicates(self, xp, sp):
         m = _make_duplicate(xp, sp, self.dtype)
@@ -733,6 +743,36 @@ def test_eliminate_zeros(self, xp, sp):
         return m.toarray()
 
 
+@testing.parameterize(*testing.product({
+    'dtype': [numpy.float32, numpy.float64],
+    'ret_dtype': [None, numpy.float32, numpy.float64],
+    'axis': [None, 0, 1, -1, -2],
+}))
+@unittest.skipUnless(scipy_available, 'requires scipy')
+class TestCscMatrixSum(unittest.TestCase):
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype)
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum_with_out(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        if self.axis is None:
+            shape = ()
+        else:
+            shape = list(m.shape)
+            shape[self.axis] = 1
+            shape = tuple(shape)
+        out = xp.empty(shape, dtype=self.ret_dtype)
+        if xp is numpy:
+            # TODO(unno): numpy.matrix is used for scipy.sparse though
+            # cupy.ndarray is used for cupy.sparse.
+            out = xp.asmatrix(out)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype, out=out)
+
+
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float32, numpy.float64],
 }))

tests/cupy_tests/sparse_tests/test_csr.py

@@ -789,6 +789,21 @@ def test_sort_indices(self, xp, sp):
         m.sort_indices()
         return m.toarray()
 
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_tuple_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=(0, 1))
+
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_str_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis='test')
+
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_too_large_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=3)
+
     @testing.numpy_cupy_allclose(sp_name='sp')
     def test_sum_duplicates(self, xp, sp):
         m = _make_duplicate(xp, sp, self.dtype)
@@ -815,6 +830,36 @@ def test_eliminate_zeros(self, xp, sp):
         return m.toarray()
 
 
+@testing.parameterize(*testing.product({
+    'dtype': [numpy.float32, numpy.float64],
+    'ret_dtype': [None, numpy.float32, numpy.float64],
+    'axis': [None, 0, 1, -1, -2],
+}))
+@unittest.skipUnless(scipy_available, 'requires scipy')
+class TestCsrMatrixSum(unittest.TestCase):
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype)
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum_with_out(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        if self.axis is None:
+            shape = ()
+        else:
+            shape = list(m.shape)
+            shape[self.axis] = 1
+            shape = tuple(shape)
+        out = xp.empty(shape, dtype=self.ret_dtype)
+        if xp is numpy:
+            # TODO(unno): numpy.matrix is used for scipy.sparse though
+            # cupy.ndarray is used for cupy.sparse.
+            out = xp.asmatrix(out)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype, out=out)
+
+
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float32, numpy.float64],
 }))

tests/cupy_tests/sparse_tests/test_dia.py

@@ -150,6 +150,21 @@ def test_A(self, xp, sp):
         m = self.make(xp, sp, self.dtype)
         return m.A
 
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_tuple_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=(0, 1))
+
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_float_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=0.0)
+
+    @testing.numpy_cupy_raises(sp_name='sp')
+    def test_sum_too_large_axis(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        m.sum(axis=3)
+
     @testing.numpy_cupy_allclose(sp_name='sp')
     def test_tocoo(self, xp, sp):
         m = self.make(xp, sp, self.dtype)
@@ -192,6 +207,36 @@ def test_transpose(self, xp, sp):
         return m.transpose().toarray()
 
 
+@testing.parameterize(*testing.product({
+    'dtype': [numpy.float32, numpy.float64],
+    'ret_dtype': [None, numpy.float32, numpy.float64],
+    'axis': [None, 0, 1, -1, -2],
+}))
+@unittest.skipUnless(scipy_available, 'requires scipy')
+class TestDiaMatrixSum(unittest.TestCase):
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype)
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_sum_with_out(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        if self.axis is None:
+            shape = ()
+        else:
+            shape = list(m.shape)
+            shape[self.axis] = 1
+            shape = tuple(shape)
+        out = xp.empty(shape, dtype=self.ret_dtype)
+        if xp is numpy:
+            # TODO(unno): numpy.matrix is used for scipy.sparse though
+            # cupy.ndarray is used for cupy.sparse.
+            out = xp.asmatrix(out)
+        return m.sum(axis=self.axis, dtype=self.ret_dtype, out=out)
+
+
 class TestIsspmatrixDia(unittest.TestCase):
 
     def test_dia(self):