Merge pull request #4017 from toslunar/sparsenorm

Add `sparse.linalg.norm`

cupyx/scipy/sparse/linalg/__init__.py

@@ -2,4 +2,5 @@
 # https://docs.scipy.org/doc/scipy/reference/sparse.linalg.html
 
 # "NOQA" to suppress flake8 warning
+from cupyx.scipy.sparse.linalg._norm import norm  # NOQA
 from cupyx.scipy.sparse.linalg._solve import lsqr  # NOQA

cupyx/scipy/sparse/linalg/_norm.py

@@ -0,0 +1,113 @@
+import numpy
+
+import cupy
+import cupyx.scipy.sparse
+
+__all__ = ['norm']
+
+
+def _sparse_frobenius_norm(x):
+    if cupy.issubdtype(x.dtype, cupy.complexfloating):
+        sqnorm = abs(x).power(2).sum()
+    else:
+        sqnorm = x.power(2).sum()
+    return cupy.sqrt(sqnorm)
+
+
+def norm(x, ord=None, axis=None):
+    """Norm of a cupy.scipy.spmatrix
+
+    This function is able to return one of seven different sparse matrix norms,
+    depending on the value of the ``ord`` parameter.
+
+    Args:
+        x (sparse matrix) : Input sparse matrix.
+        ord (non-zero int, inf, -inf, 'fro', optional) : Order of the norm (see
+            table under ``Notes``). inf means numpy's `inf` object.
+        axis : (int, 2-tuple of ints, None, optional): If `axis` is an
+            integer, it specifies the axis of `x` along which to
+            compute the vector norms.  If `axis` is a 2-tuple, it specifies the
+            axes that hold 2-D matrices, and the matrix norms of these matrices
+            are computed.  If `axis` is None then either a vector norm
+            (when `x` is 1-D) or a matrix norm (when `x` is 2-D) is returned.
+    Returns:
+        ndarray : 0-D or 1-D array or norm(s).
+
+    .. seealso:: :func:`scipy.sparse.linalg.norm`
+    """
+
+    if not cupyx.scipy.sparse.issparse(x):
+        raise TypeError(("input is not sparse. use cupy.linalg.norm"))
+
+    # Check the default case first and handle it immediately.
+    if axis is None and ord in (None, 'fro', 'f'):
+        return _sparse_frobenius_norm(x)
+
+    # Some norms require functions that are not implemented for all types.
+    x = x.tocsr()
+
+    if axis is None:
+        axis = (0, 1)
+    elif not isinstance(axis, tuple):
+        msg = "'axis' must be None, an integer or a tuple of integers"
+        try:
+            int_axis = int(axis)
+        except TypeError:
+            raise TypeError(msg)
+        if axis != int_axis:
+            raise TypeError(msg)
+        axis = (int_axis,)
+
+    nd = 2
+    if len(axis) == 2:
+        row_axis, col_axis = axis
+        if not (-nd <= row_axis < nd and -nd <= col_axis < nd):
+            raise ValueError('Invalid axis %r for an array with shape %r' %
+                             (axis, x.shape))
+        if row_axis % nd == col_axis % nd:
+            raise ValueError('Duplicate axes given.')
+        if ord == 2:
+            raise NotImplementedError
+            # return _multi_svd_norm(x, row_axis, col_axis, amax)
+        elif ord == -2:
+            raise NotImplementedError
+            # return _multi_svd_norm(x, row_axis, col_axis, amin)
+        elif ord == 1:
+            return abs(x).sum(axis=row_axis).max()
+        elif ord == numpy.Inf:
+            return abs(x).sum(axis=col_axis).max()
+        elif ord == -1:
+            return abs(x).sum(axis=row_axis).min()
+        elif ord == -numpy.Inf:
+            return abs(x).sum(axis=col_axis).min()
+        elif ord in (None, 'f', 'fro'):
+            # The axis order does not matter for this norm.
+            return _sparse_frobenius_norm(x)
+        else:
+            raise ValueError("Invalid norm order for matrices.")
+    elif len(axis) == 1:
+        a, = axis
+        if not (-nd <= a < nd):
+            raise ValueError('Invalid axis %r for an array with shape %r' %
+                             (axis, x.shape))
+        if ord == numpy.Inf:
+            return abs(x).max(axis=a).A.ravel()
+        elif ord == -numpy.Inf:
+            return abs(x).min(axis=a).A.ravel()
+        elif ord == 0:
+            # Zero norm
+            return (x != 0).astype(numpy.float32).sum(axis=a).ravel().astype(
+                numpy.int_)
+        elif ord == 1:
+            # special case for speedup
+            return abs(x).sum(axis=a).ravel()
+        elif ord in (2, None):
+            return cupy.sqrt(abs(x).power(2).sum(axis=a)).ravel()
+        else:
+            try:
+                ord + 1
+            except TypeError:
+                raise ValueError('Invalid norm order for vectors.')
+            return cupy.power(abs(x).power(ord).sum(axis=a), 1 / ord).ravel()
+    else:
+        raise ValueError("Improper number of dimensions to norm.")

docs/source/reference/sparse.rst

@@ -90,3 +90,4 @@ Linear Algebra
    :nosignatures:
 
    cupyx.scipy.sparse.linalg.lsqr
+   cupyx.scipy.sparse.linalg.norm

tests/cupyx_tests/scipy_tests/sparse_tests/test_linalg.py

@@ -57,3 +57,55 @@ def test_ndarray(self, xp, sp):
         b = xp.array(self.b, dtype=self.dtype)
         x = sp.linalg.lsqr(A, b)
         return x[0]
+
+
+@testing.parameterize(*testing.product({
+    'ord': [None, -numpy.Inf, -2, -1, 0, 1, 2, 3, numpy.Inf, 'fro'],
+    'dtype': [
+        numpy.float32,
+        numpy.float64,
+        numpy.complex64,
+        numpy.complex128
+    ],
+    'axis': [None, (0, 1), (1, -2)],
+}))
+@unittest.skipUnless(scipy_available, 'requires scipy')
+@testing.gpu
+class TestMatrixNorm(unittest.TestCase):
+
+    @testing.numpy_cupy_allclose(rtol=1e-3, atol=1e-4, sp_name='sp',
+                                 accept_error=(ValueError,
+                                               NotImplementedError))
+    def test_matrix_norm(self, xp, sp):
+        a = xp.arange(9, dtype=self.dtype) - 4
+        b = a.reshape((3, 3))
+        b = sp.csr_matrix(b, dtype=self.dtype)
+        return sp.linalg.norm(b, ord=self.ord, axis=self.axis)
+
+
+@testing.parameterize(*testing.product({
+    'ord': [None, -numpy.Inf, -2, -1, 0, 1, 2, numpy.Inf, 'fro'],
+    'dtype': [
+        numpy.float32,
+        numpy.float64,
+        numpy.complex64,
+        numpy.complex128
+    ],
+    'transpose': [True, False],
+    'axis': [0, (1,), (-2,), -1],
+})
+)
+@unittest.skipUnless(scipy_available, 'requires scipy')
+@testing.gpu
+class TestVectorNorm(unittest.TestCase):
+    @testing.numpy_cupy_allclose(rtol=1e-3, atol=1e-4, sp_name='sp',
+                                 accept_error=(ValueError,))
+    def test_vector_norm(self, xp, sp):
+        a = xp.arange(9, dtype=self.dtype) - 4
+        b = a.reshape((3, 3))
+        b = sp.csr_matrix(b, dtype=self.dtype)
+        if self.transpose:
+            b = b.T
+        return sp.linalg.norm(b, ord=self.ord, axis=self.axis)
+
+# TODO : TestVsNumpyNorm