CCW#37: Reviewed Remove0 sparse op.

theano/sparse/sandbox/sp.py

@@ -222,32 +222,27 @@ class Remove0(Op):
     """
     Remove explicit zeros from a sparse matrix, and resort indices
     """
+
+    def __init__(self, inplace=False, *args, **kwargs):
+        Op.__init__(self, *args, **kwargs)
+        self.inplace = inplace
+        if self.inplace:
+            self.destroy_map = {0: [0]}
+
+    def __eq__(self,other):
+        return type(self) == type(other) and self.inplace == other.inplace
+    
     def make_node(self, x):
         return gof.Apply(self, [x], [x.type()])
 
     def perform(self,node, (x,), (z,)):
-        if x.format != 'csc':
-            raise TypeError('Remove0 only works on csc matrices')
-
-        M, N = x.shape
-
-        data = x.data
-        indices = x.indices
-        indptr = x.indptr
-
-        #TODO: try using ndarrays and then prune() on the result
-        new_data = []
-        new_indices = []
-        new_indptr = [0]
-
-        for j in xrange(0, N):
-            for i_idx in xrange(indptr[j], indptr[j+1]):
-                if data[i_idx] != 0:
-                    new_data.append(data[i_idx])
-                    new_indices.append(indices[i_idx])
-            new_indptr.append(len(new_indices))
-
-        z[0] = sparse.csc_matrix((new_data, new_indices, new_indptr), (M,N))
+        if self.inplace:
+            c = x
+        else:
+            c = x.copy()
+        c.eliminate_zeros()
+        z[0] = c
+        return
 
     def grad(self, (x,), (gz,)):
         return [gz]
@@ -281,7 +276,7 @@ class EnsureSortedIndices(Op):
 
     def infer_shape(self, node, i0_shapes):
         return i0_shapes
-            
+
     def __str__(self):
         if self.inplace:
             return self.__class__.__name__ + "{inplace}"

theano/sparse/sandbox/test_sp.py

@@ -16,7 +16,6 @@ from theano import function, tensor
 import theano
 from theano.sparse.sandbox import sp
 from theano.tests import unittest_tools as utt
-from theano.sparse.tests.test_basic import random_lil
 
 
 class TestSP(unittest.TestCase):
@@ -426,6 +425,36 @@ class TestSP(unittest.TestCase):
                 #utt.verify_grad(SpSum(axis=None), [x_val])
                 print 'ok'
 
+def test_remove0():
+    print
+    print 'test_remove0()'
+    import scipy
+    import numpy as N
+    configs=[
+        # structure type, numpy matching class, optimized
+        ('csc',scipy.sparse.csc_matrix),
+        ('csr',scipy.sparse.csr_matrix),
+        ]
+
+    for format,matrix_class in configs:
+        
+        print 'config: format=\'%(format)s\', matrix_class=%(matrix_class)s'%locals()
+
+        # real
+        origin = N.arange(9).reshape((3,3)).astype(theano.config.floatX)
+        with0= matrix_class(origin).astype(theano.config.floatX)
+        with0[0,1]=with0[1,0]=with0[2,2]=0
+        
+        # symbolic
+        x = theano.sparse.SparseType(format=format,dtype=theano.config.floatX)()
+        f=theano.function([x],sp.Remove0()(x))
+        
+        # checking
+        with0.eliminate_zeros()
+        # makes sense to change its name
+        target=with0
+        assert (f(with0).todense()==target.todense()).all()
+
 
 def test_diagonal():
     for K in 1, 5:
@@ -456,13 +485,13 @@ def test_ensure_sorted_indices():
             # csr
             input_tensor = theano.sparse.csr_dmatrix()
             sample = scipy.sparse.csr_matrix(random_lil((x,y),'float64',sparsity))
-            
+
         sort_op = sp.ensure_sorted_indices(input_tensor)
         f = theano.function([input_tensor], sort_op)
         sorted_scipy = sample.sorted_indices()
         sorted_theano = f(sample)
         assert numpy.all(sorted_theano.todense() == sorted_scipy.todense())
-    
+
 def test_diagonal_grad():
     def d(x):
         return sp.sp_sum(sp.square_diagonal(x), sparse_grad=True)
@@ -532,6 +561,8 @@ def test_col_scale():
         print >> sys.stderr, "WARNING: skipping gradient test because verify_grad doesn't support sparse arguments"
 
 if __name__ == '__main__':
+    test_remove0()
+    exit()
     if 1:
         testcase =  TestSP
         suite = unittest.TestLoader()