Add hstack and vstack for sparse

theano/sparse/basic.py

@@ -21,7 +21,6 @@ import theano.tests.unittest_tools as utt
 
 sparse_formats = ['csc', 'csr']
 
-
 #TODO: move this decorator to the compile submodule
 def register_specialize(lopt, *tags, **kwargs):
     compile.optdb['specialize'].register((kwargs and kwargs.pop('name')) or
@@ -493,6 +492,14 @@ csr_dmatrix = SparseType(format='csr', dtype='float64')
 csc_fmatrix = SparseType(format='csc', dtype='float32')
 csr_fmatrix = SparseType(format='csr', dtype='float32')
 
+all_dtypes = SparseType.dtype_set
+complex_dtypes = [t for t in all_dtypes if t[:7] == 'complex']
+float_dtypes = [t for t in all_dtypes if t[:5] == 'float']
+int_dtypes = [t for t in all_dtypes if t[:3] == 'int']
+uint_dtypes = [t for t in all_dtypes if t[:4] == 'uint']
+
+continuous_dtypes = complex_dtypes + float_dtypes
+discrete_dtypes = int_dtypes + uint_dtypes
 
 # CONSTRUCTION
 class CSMProperties(gof.Op):

theano/sparse/sandbox/sp2.py

@@ -2,7 +2,7 @@ import theano
 import numpy
 import scipy.sparse
 
-from theano import gof, tensor, scalar
+from theano import gof, tensor, scalar, sparse
 from theano.tensor import blas
 
 from theano.sparse.basic import (
@@ -93,6 +93,160 @@ def dcast(x):
     return Cast('float64')(x)
 
 
+class HStack(gof.op.Op):
+    """Stack sparse matrices horizontally (column wise).
+
+    This wrap the method hstack from scipy.
+
+    :Parameters:
+    - `blocks`: Sequence of sparse array of compatible shape
+    - `format`: String representing the output format
+    - `dtype`: Output dtype
+
+    :return: the concatenation of the sparse arrays column wise.
+
+    The number of line of the sparse matrix must agree.
+    """
+
+    def __init__(self, format=None, dtype=None):
+        if format is None:
+            self.format = 'csc'
+        else:
+            self.format = format
+        if dtype is None:
+            self.dtype = theano.config.floatX
+        else:
+            self.dtype = dtype
+
+    def __eq__(self, other):
+        return (type(self) == type(other) and
+                self.format == other.format and
+                self.dtype == other.dtype)
+
+    def __hash__(self):
+        return hash(type(self)) ^ hash(self.format) ^ hash(self.dtype)
+
+    def make_node(self, *mat):
+        if not mat:
+            raise ValueError('Cannot join an empty list of sparses.')
+        var = [as_sparse_variable(x) for x in mat]
+        return gof.Apply(
+            self, var,
+            [SparseType(dtype=self.dtype, format=self.format).make_variable()])
+
+    def perform(self, node, block, (out, )):
+        for b in block:
+            assert _is_sparse(b)
+        out[0] = scipy.sparse.hstack(block, format=self.format,
+                                     dtype=self.dtype)
+
+    def grad(self, inputs, (gz, )):
+        is_continuous = [(inputs[i].dtype in tensor.continuous_dtypes)
+                        for i in range(len(inputs))]
+
+        if all(is_continuous):
+            if _is_sparse_variable(gz):
+                gz = sparse.DenseFromSparse()(gz)
+
+            split = tensor.Split(len(inputs))(gz, 1,
+                                              tensor.stack(
+                                                  *[x.shape[1]
+                                                    for x in inputs]))
+            if not isinstance(split, list):
+                split = [split]
+            return [sparse.SparseFromDense(self.format)(s) for s in split]
+        else:
+            return [None] * len(inputs)
+
+    def infer_shape(self, node, ins_shapes):
+        def _get(l):
+            return l[1]
+        d = sum(map(_get, ins_shapes))
+        return [(ins_shapes[0][0], d)]
+
+    def __str__(self):
+        return "%s(%s,%s)" % (self.__class__.__name__, self.format, self.dtype)
+
+
+def hstack(blocks, format=None, dtype=None):
+    """Stack sparse matrices horizontally (column wise).
+
+    This wrap the method hstack from scipy.
+
+    :Parameters:
+    - `blocks`: Sequence of sparse array of compatible shape
+    - `format`: String representing the output format
+    - `dtype`: Output dtype
+
+    :return: the concatenation of the sparse array column wise.
+
+    The number of line of the sparse matrix must agree.
+    """
+    return HStack(format=format, dtype=dtype)(*blocks)
+
+
+class VStack(HStack):
+    """Stack sparse matrices vertically (row wise).
+
+    This wrap the method vstack from scipy.
+
+    :Parameters:
+    - `blocks`: Sequence of sparse array of compatible shape
+    - `format`: String representing the output format
+    - `dtype`: Output dtype
+
+    :return: the concatenation of the sparse arrays row wise.
+
+    The number of column of the sparse matrix must agree.
+    """
+    def perform(self, node, block, (out, )):
+        for b in block:
+            assert _is_sparse(b)
+        out[0] = scipy.sparse.vstack(block, format=self.format,
+                                     dtype=self.dtype)
+
+    def grad(self, inputs, (gz, )):
+        is_continuous = [(inputs[i].dtype in tensor.continuous_dtypes)
+                        for i in range(len(inputs))]
+
+        if all(is_continuous):
+            if _is_sparse_variable(gz):
+                gz = sparse.DenseFromSparse()(gz)
+
+            split = tensor.Split(len(inputs))(gz, 0,
+                                              tensor.stack(
+                                                  *[x.shape[0]
+                                                    for x in inputs]))
+            if not isinstance(split, list):
+                split = [split]
+            return [sparse.SparseFromDense(self.format)(s) for s in split]
+        else:
+            return [None] * len(inputs)
+
+    def infer_shape(self, node, ins_shapes):
+        def _get(l):
+            return l[0]
+        d = sum(map(_get, ins_shapes))
+        return [(d, ins_shapes[0][1])]
+
+
+def hstack(blocks, format=None, dtype=None):
+    """Stack sparse matrices vertically (row wise).
+
+    This wrap the method vstack from scipy.
+
+    :Parameters:
+    - `blocks`: Sequence of sparse array of compatible shape
+    - `format`: String representing the output format
+    - `dtype`: Output dtype
+
+    :return: the concatenation of the sparse array row wise.
+
+    The number of column of the sparse matrix must agree.
+    """
+    return VStack(format=format, dtype=dtype)(*blocks)
+
+
 class AddSSData(gof.op.Op):
     '''Add two sparse matrices assuming they have the same sparsity
     pattern. '''

theano/sparse/tests/test_sp2.py

@@ -118,6 +118,71 @@ class TestCast(utt.InferShapeTester):
             verify_grad_sparse(S2.Cast('float64'), [a])
 
 
+class HVStackTester(utt.InferShapeTester):
+    nb = 3  # Number of sparse matrix to stack
+    x = {}
+    mat = {}
+
+    for format in sparse.sparse_formats:
+        variable = getattr(theano.sparse, format + '_matrix')
+        spa = getattr(sp, format + '_matrix')
+
+        x[format] = [variable() for t in range(nb)]
+        mat[format] = [spa(np.random.random_integers(5, size=(3, 4)) - 1,
+                           dtype=theano.config.floatX)
+                       for t in range(nb)]
+
+    def test_op(self):
+        for format in sparse.sparse_formats:
+            for out_f in sparse.sparse_formats:
+                for dtype in sparse.all_dtypes:
+                    blocks = self.mat[format]
+
+                    f = theano.function(
+                        self.x[format],
+                        self.op_class(
+                            format=out_f, dtype=dtype)(*self.x[format]),
+                        allow_input_downcast=True)
+
+                    tested = f(*blocks)
+                    expected = self.expected_f(blocks, format=out_f, dtype=dtype)
+
+                    assert np.allclose(tested.toarray(), expected.toarray())
+                    assert tested.format == expected.format
+                    assert tested.dtype == expected.dtype
+
+    def test_infer_shape(self):
+        for format in sparse.sparse_formats:
+            self._compile_and_check(self.x[format],
+                                    [self.op_class()(*self.x[format])],
+                                    self.mat[format],
+                                    self.op_class)
+
+    def test_grad(self):
+        for format in sparse.sparse_formats:
+            for out_f in sparse.sparse_formats:
+                for dtype in ['float64']:  # sparse.float_dtypes:
+                    verify_grad_sparse(
+                        self.op_class(format=out_f, dtype=dtype),
+                        self.mat[format],
+                        structured=False)
+
+
+def _hv_switch(op, expected_function):
+    class XStackTester(HVStackTester):
+        op_class = op
+
+        def expected_f(self, a, format=None, dtype=None):
+            return expected_function(a, format, dtype)
+
+        def setUp(self):
+            super(XStackTester, self).setUp()
+
+    return XStackTester
+
+HStackTester = _hv_switch(S2.HStack, sp.hstack)
+VStackTester = _hv_switch(S2.VStack, sp.vstack)
+
 class test_structured_add_s_v(unittest.TestCase):
     def setUp(self):
         utt.seed_rng()