Merge pull request #321 from unnonouno/dia-array-conversion

Array conversion from dia_matrix

cupy/sparse/dia.py

@@ -6,6 +6,7 @@
 
 import cupy
 from cupy import core
+from cupy.sparse import csc
 from cupy.sparse import data
 
 
@@ -114,6 +115,62 @@ def getnnz(self, axis=None):
             'a + b', 'nnz = a', '0', 'dia_nnz')(self.offsets, m, n)
         return int(nnz)
 
+    def toarray(self, order=None, out=None):
+        """Returns a dense matrix representing the same value."""
+        return self.tocsc().toarray(order=order, out=out)
+
+    def tocsc(self, copy=False):
+        """Converts the matrix to Compressed Sparse Column format.
+
+        Args:
+            copy (bool): If ``False``, it shares data arrays as much as
+                possible. Actually this option is ignored because all
+                arrays in a matrix cannot be shared in dia to csc conversion.
+
+        Returns:
+            cupy.sparse.csc_matrix: Converted matrix.
+
+        """
+        if self.nnz == 0:
+            return csc.csc_matrix(self.shape, dtype=self.dtype)
+
+        num_rows, num_cols = self.shape
+        num_offsets, offset_len = self.data.shape
+
+        row, mask = core.ElementwiseKernel(
+            'int32 offset_len, int32 offsets, int32 num_rows, '
+            'int32 num_cols, T data',
+            'int32 row, bool mask',
+            '''
+            int offset_inds = i % offset_len;
+            row = offset_inds - offsets;
+            mask = (row >= 0 && row < num_rows && offset_inds < num_cols
+                    && data != 0);
+            ''',
+            'dia_tocsc')(offset_len, self.offsets[:, None], num_rows,
+                         num_cols, self.data)
+        indptr = cupy.zeros(num_cols + 1, dtype='i')
+        indptr[1: offset_len + 1] = cupy.cumsum(mask.sum(axis=0))
+        indptr[offset_len + 1:] = indptr[offset_len]
+        indices = row.T[mask.T].astype('i', copy=False)
+        data = self.data.T[mask.T]
+        return csc.csc_matrix(
+            (data, indices, indptr), shape=self.shape, dtype=self.dtype)
+
+    def tocsr(self, copy=False):
+        """Converts the matrix to Compressed Sparse Row format.
+
+        Args:
+            copy (bool): If ``False``, it shares data arrays as much as
+                possible. Actually this option is ignored because all
+                arrays in a matrix cannot be shared in dia to csr conversion.
+
+        Returns:
+            cupy.sparse.csc_matrix: Converted matrix.
+
+        """
+        return self.tocsc().tocsr()
+
 
 def isspmatrix_dia(x):
     """Checks if a given matrix is of DIA format.

tests/cupy_tests/sparse_tests/test_dia.py

@@ -52,6 +52,23 @@ def test_ndim(self):
     def test_nnz(self):
         self.assertEqual(self.m.nnz, 5)
 
+    @unittest.skipUnless(scipy_available, 'requires scipy')
+    def test_str(self):
+        self.assertEqual(str(self.m), '''  (1, 1)\t1.0
+  (2, 2)\t2.0
+  (1, 0)\t3.0
+  (2, 1)\t4.0''')
+
+    def test_toarray(self):
+        m = self.m.toarray()
+        expect = [
+            [0, 0, 0, 0],
+            [3, 1, 0, 0],
+            [0, 4, 2, 0]
+        ]
+        self.assertTrue(m.flags.c_contiguous)
+        cupy.testing.assert_allclose(m, expect)
+
 
 @testing.parameterize(*testing.product({
     'dtype': [numpy.float32, numpy.float64],
@@ -111,3 +128,23 @@ def test_nnz_axis(self, xp, sp):
     def test_nnz_axis_not_none(self, xp, sp):
         m = _make(xp, sp, self.dtype)
         m.getnnz(axis=0)
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_toarray(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.toarray()
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_tocsc(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.tocsc().toarray()
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_tocsr(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.tocsr().toarray()
+
+    @testing.numpy_cupy_allclose(sp_name='sp')
+    def test_transpose(self, xp, sp):
+        m = _make(xp, sp, self.dtype)
+        return m.transpose().toarray()