Merge pull request #313 from unnonouno/dia

Implement dia_matrix

cupy/sparse/__init__.py

@@ -7,12 +7,13 @@
 from cupy.sparse.csc import isspmatrix_csc  # NOQA
 from cupy.sparse.csr import csr_matrix  # NOQA
 from cupy.sparse.csr import isspmatrix_csr  # NOQA
+from cupy.sparse.dia import dia_matrix  # NOQA
+from cupy.sparse.dia import isspmatrix_dia  # NOQA
 
 from cupy.sparse.construct import eye  # NOQA
 from cupy.sparse.construct import identity  # NOQA
 
 # TODO(unno): implement bsr_matrix
-# TODO(unno): implement dia_matrix
 # TODO(unno): implement dok_matrix
 # TODO(unno): implement lil_matrix
 
@@ -37,4 +38,3 @@
 # TODO(unno): implement isspmatrix_bsr(x)
 # TODO(unno): implement isspmatrix_lil(x)
 # TODO(unno): implement isspmatrix_dok(x)
-# TODO(unno): implement isspmatrix_dia(x)

cupy/sparse/dia.py

@@ -0,0 +1,103 @@
+try:
+    import scipy.sparse
+    _scipy_available = True
+except ImportError:
+    _scipy_available = False
+
+import cupy
+from cupy.sparse import data
+
+
+class dia_matrix(data._data_matrix):
+
+    """Sparse matrix with DIAgonal storage.
+
+    Now it has only one initializer format below:
+
+    ``dia_matrix((data, offsets))``
+
+    Args:
+        arg1: Arguments for the initializer.
+        shape (tuple): Shape of a matrix. Its length must be two.
+        dtype: Data type. It must be an argument of :class:`numpy.dtype`.
+        copy (bool): If ``True``, copies of given arrays are always used.
+
+    .. see::
+       :class:`scipy.sparse.dia_matrix`
+
+    """
+
+    format = 'dia'
+
+    def __init__(self, arg1, shape=None, dtype=None, copy=False):
+        if isinstance(arg1, tuple):
+            data, offsets = arg1
+            if shape is None:
+                raise ValueError('expected a shape argument')
+
+        else:
+            raise ValueError(
+                'unrecognized form for dia_matrix constructor')
+
+        data = cupy.array(data, dtype=dtype, copy=copy)
+        data = cupy.atleast_2d(data)
+        offsets = cupy.array(offsets, dtype='i', copy=copy)
+        offsets = cupy.atleast_1d(offsets)
+
+        if offsets.ndim != 1:
+            raise ValueError('offsets array must have rank 1')
+
+        if data.ndim != 2:
+            raise ValueError('data array must have rank 2')
+
+        if data.shape[0] != len(offsets):
+            raise ValueError(
+                'number of diagonals (%d) does not match the number of '
+                'offsets (%d)'
+                % (data.shape[0], len(offsets)))
+
+        sorted_offsets = cupy.sort(offsets)
+        if (sorted_offsets[:-1] == sorted_offsets[1:]).any():
+            raise ValueError('offset array contains duplicate values')
+
+        self.data = data
+        self.offsets = offsets
+        self._shape = shape
+
+    def _with_data(self, data):
+        return dia_matrix((data, self.offsets), shape=self.shape)
+
+    def get(self, stream=None):
+        """Returns a copy of the array on host memory.
+
+        Args:
+            stream (cupy.cuda.Stream): CUDA stream object. If it is given, the
+                copy runs asynchronously. Otherwise, the copy is synchronous.
+
+        Returns:
+            scipy.sparse.dia_matrix: Copy of the array on host memory.
+
+        """
+        if not _scipy_available:
+            raise RuntimeError('scipy is not available')
+        data = self.data.get(stream)
+        offsets = self.offsets.get(stream)
+        return scipy.sparse.dia_matrix((data, offsets), shape=self._shape)
+
+    def get_shape(self):
+        """Returns the shape of the matrix.
+
+        Returns:
+            tuple: Shape of the matrix.
+        """
+        return self._shape
+
+
+def isspmatrix_dia(x):
+    """Checks if a given matrix is of DIA format.
+
+    Returns:
+        bool: Returns if ``x`` is :class:`cupy.sparse.dia_matrix`.
+
+    """
+    return isinstance(x, dia_matrix)

tests/cupy_tests/sparse_tests/test_dia.py

@@ -0,0 +1,93 @@
+import unittest
+
+
+import numpy
+try:
+    import scipy.sparse  # NOQA
+    scipy_available = True
+except ImportError:
+    scipy_available = False
+
+import cupy
+import cupy.sparse
+from cupy import testing
+
+
+def _make(xp, sp, dtype):
+    data = xp.array([[0, 1, 2], [3, 4, 5]], dtype)
+    offsets = xp.array([0, -1], 'i')
+    # 0, 0, 0, 0
+    # 3, 1, 0, 0
+    # 0, 4, 2, 0
+    return sp.dia_matrix((data, offsets), shape=(3, 4))
+
+
+@testing.parameterize(*testing.product({
+    'dtype': [numpy.float32, numpy.float64],
+}))
+class TestDiaMatrix(unittest.TestCase):
+
+    def setUp(self):
+        self.m = _make(cupy, cupy.sparse, self.dtype)
+
+    def test_dtype(self):
+        self.assertEqual(self.m.dtype, self.dtype)
+
+    def test_data(self):
+        self.assertEqual(self.m.data.dtype, self.dtype)
+        testing.assert_array_equal(
+            self.m.data, cupy.array([[0, 1, 2], [3, 4, 5]], self.dtype))
+
+    def test_offsets(self):
+        self.assertEqual(self.m.offsets.dtype, numpy.int32)
+        testing.assert_array_equal(
+            self.m.offsets, cupy.array([0, -1], self.dtype))
+
+    def test_shape(self):
+        self.assertEqual(self.m.shape, (3, 4))
+
+    def test_ndim(self):
+        self.assertEqual(self.m.ndim, 2)
+
+
+@testing.parameterize(*testing.product({
+    'dtype': [numpy.float32, numpy.float64],
+}))
+@unittest.skipUnless(scipy_available, 'requires scipy')
+class TestDiaMatrixInit(unittest.TestCase):
+
+    def setUp(self):
+        self.shape = (3, 4)
+
+    def data(self, xp):
+        return xp.array([[1, 2, 3], [4, 5, 6]], self.dtype)
+
+    def offsets(self, xp):
+        return xp.array([0, -1], 'i')
+
+    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
+    def test_shape_none(self, xp, sp):
+        sp.dia_matrix(
+            (self.data(xp), self.offsets(xp)), shape=None)
+
+    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
+    def test_large_rank_offset(self, xp, sp):
+        sp.dia_matrix(
+            (self.data(xp), self.offsets(xp)[None]))
+
+    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
+    def test_large_rank_data(self, xp, sp):
+        sp.dia_matrix(
+            (self.data(xp)[None], self.offsets(xp)))
+
+    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
+    def test_data_offsets_different_size(self, xp, sp):
+        offsets = xp.array([0, -1, 1], 'i')
+        sp.dia_matrix(
+            (self.data(xp), offsets))
+
+    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
+    def test_duplicated_offsets(self, xp, sp):
+        offsets = xp.array([1, 1], 'i')
+        sp.dia_matrix(
+            (self.data(xp), offsets))