1from __future__ import absolute_import, print_function, division
2import unittest
3
4from nose.plugins.skip import SkipTest
5import numpy as np
6try:
7    import scipy.sparse as sp
8except ImportError:
9    pass  # The variable enable_sparse will be used to disable the test file.
10
11import theano
12from theano import config
13from theano import tensor
14from theano import sparse
15
16if not theano.sparse.enable_sparse:
17    raise SkipTest('Optional package sparse disabled')
18
19from theano.sparse.sandbox.sp2 import (
20    Poisson, poisson, Binomial, Multinomial, multinomial)
21
22from theano.tests import unittest_tools as utt
23from theano.sparse.tests.test_basic import as_sparse_format
24
25
26class PoissonTester(utt.InferShapeTester):
27    x = {}
28    a = {}
29
30    for format in sparse.sparse_formats:
31        variable = getattr(theano.sparse, format + '_matrix')
32
33        rand = np.array(np.random.randint(1, 4, size=(3, 4)) - 1,
34                           dtype=theano.config.floatX)
35
36        x[format] = variable()
37        a[format] = as_sparse_format(rand, format)
38
39    def setUp(self):
40        super(PoissonTester, self).setUp()
41        self.op_class = Poisson
42
43    def test_op(self):
44        for format in sparse.sparse_formats:
45            f = theano.function(
46                [self.x[format]],
47                poisson(self.x[format]))
48
49            tested = f(self.a[format])
50
51            assert tested.format == format
52            assert tested.dtype == self.a[format].dtype
53            assert np.allclose(np.floor(tested.data), tested.data)
54            assert tested.shape == self.a[format].shape
55
56    def test_infer_shape(self):
57        for format in sparse.sparse_formats:
58            self._compile_and_check([self.x[format]],
59                                    [poisson(self.x[format])],
60                                    [self.a[format]],
61                                    self.op_class)
62
63
64class BinomialTester(utt.InferShapeTester):
65    n = tensor.scalar(dtype='int64')
66    p = tensor.scalar()
67    shape = tensor.lvector()
68    _n = 5
69    _p = .25
70    _shape = np.asarray([3, 5], dtype='int64')
71
72    inputs = [n, p, shape]
73    _inputs = [_n, _p, _shape]
74
75    def setUp(self):
76        super(BinomialTester, self).setUp()
77        self.op_class = Binomial
78
79    def test_op(self):
80        for sp_format in sparse.sparse_formats:
81            for o_type in sparse.float_dtypes:
82                f = theano.function(
83                    self.inputs,
84                    Binomial(sp_format, o_type)(*self.inputs))
85
86                tested = f(*self._inputs)
87
88                assert tested.shape == tuple(self._shape)
89                assert tested.format == sp_format
90                assert tested.dtype == o_type
91                assert np.allclose(np.floor(tested.todense()),
92                                   tested.todense())
93
94    def test_infer_shape(self):
95        for sp_format in sparse.sparse_formats:
96            for o_type in sparse.float_dtypes:
97                self._compile_and_check(
98                    self.inputs,
99                    [Binomial(sp_format, o_type)(*self.inputs)],
100                    self._inputs,
101                    self.op_class)
102
103
104class MultinomialTester(utt.InferShapeTester):
105    p = sparse.csr_matrix()
106    _p = sp.csr_matrix(np.asarray([[0.0, 0.5, 0.0, 0.5],
107                                      [0.1, 0.2, 0.3, 0.4],
108                                      [0.0, 1.0, 0.0, 0.0],
109                                      [0.3, 0.3, 0.0, 0.4]],
110                                     dtype=config.floatX))
111
112    def setUp(self):
113        super(MultinomialTester, self).setUp()
114        self.op_class = Multinomial
115
116    def test_op(self):
117        n = tensor.lscalar()
118        f = theano.function([self.p, n], multinomial(n, self.p))
119
120        _n = 5
121        tested = f(self._p, _n)
122        assert tested.shape == self._p.shape
123        assert np.allclose(np.floor(tested.todense()), tested.todense())
124        assert tested[2, 1] == _n
125
126        n = tensor.lvector()
127        f = theano.function([self.p, n], multinomial(n, self.p))
128
129        _n = np.asarray([1, 2, 3, 4], dtype='int64')
130        tested = f(self._p, _n)
131        assert tested.shape == self._p.shape
132        assert np.allclose(np.floor(tested.todense()), tested.todense())
133        assert tested[2, 1] == _n[2]
134
135    def test_infer_shape(self):
136        self._compile_and_check([self.p],
137                                [multinomial(5, self.p)],
138                                [self._p],
139                                self.op_class,
140                                warn=False)
141
142
143if __name__ == '__main__':
144    unittest.main()
145