version 2 bartlett and fill_diagonal 28/05/2012

theano/tensor/extra_ops.py

-import theano
 import numpy as np
+import numpy
+
+import theano
 import basic
+from theano import gof, tensor, function, scalar
+from theano.sandbox.linalg.ops import diag
 
 
 class DiffOp(theano.Op):
@@ -348,3 +352,141 @@ def repeat(x, repeats, axis=None):
 
     """
     return RepeatOp(axis=axis)(x, repeats)
+
+
+class Bartlett(gof.Op):
+    """
+    An instance of this class returns the Bartlett spectral window in the
+    time-domain. The Bartlett window is very similar to a triangular window,
+    except that the end points are at zero. It is often used in signal
+    processing for tapering a signal, without generating too much ripple in
+    the frequency domain.
+
+    input : (integer scalar) Number of points in the output window. If zero or
+    less, an empty vector is returned.
+
+    output : (vector of doubles) The triangular window, with the maximum value
+    normalized to one (the value one appears only if the number of samples is
+    odd), with the first and last samples equal to zero.
+    """
+
+    def __eq__(self, other):
+        return type(self) == type(other)
+
+    def __hash__(self):
+        return hash(type(self))
+
+    def __str__(self):
+        return self.__class__.__name__
+
+    def make_node(self, M):
+        M = tensor.as_tensor_variable(M)
+        if M.ndim != 0:
+            raise TypeError('%s only works on scalar input'
+                            % self.__class__.__name__)
+        elif (not M.dtype.startswith('int')) and \
+              (not M.dtype.startswith('uint')):
+        # dtype is a theano attribute here
+            raise TypeError('%s only works on integer input'
+                            % self.__class__.__name__)
+        return gof.Apply(self, [M], [tensor.dvector()])
+
+    def perform(self, node, inputs, out_):
+        M = inputs[0]
+        if (not M.dtype.name.startswith('int')) and \
+              (not M.dtype.startswith('uint')):
+        # dtype is an instance of numpy dtype class here
+            raise TypeError('%s only works on integers'
+                            % self.__class__.__name__)
+        out, = out_
+        out[0] = numpy.bartlett(M)
+
+    def infer_shape(self, node, in_shapes):
+        temp = node.inputs[0]
+        M = tensor.switch(tensor.lt(temp, 0),
+            tensor.cast(0, temp.dtype), temp)
+        return [[M]]
+
+    def grad(self, inputs, output_grads):
+        return [None for i in inputs]
+
+
+bartlett = Bartlett()
+
+
+class FillDiagonal(gof.Op):
+    """
+    An instance of this class returns a copy of an array with all elements of
+    the main diagonal set to a specified scalar value.
+
+    inputs:
+
+    a : Rectangular array of at least two dimensions.
+    val : Scalar value to fill the diagonal whose type must be compatible with
+    that of array 'a' (i.e. 'val' cannot be viewed as an upcast of 'a').
+
+    output:
+
+    An array identical to 'a' except that its main diagonal is filled with
+    scalar 'val'. (For an array 'a' with a.ndim >= 2, the main diagonal is the
+    list of locations a[i, i, ..., i] (i.e. with indices all identical).)
+    """
+
+    def __eq__(self, other):
+        return type(self) == type(other)
+
+    def __hash_(self):
+        return hash(type(self))
+
+    def __str__(self):
+        return self.__class__.__name__
+
+    def infer_shape(self, node, in_shapes):
+        return [in_shapes[0]]
+
+    def make_node(self, a, val):
+        a = tensor.as_tensor_variable(a)
+        val = tensor.as_tensor_variable(val)
+        if a.ndim < 2:
+            raise TypeError('%s: first parameter must have at least'
+                            ' two dimensions' % self.__class__.__name__)
+        elif val.ndim != 0:
+            raise TypeError('%s: second parameter must be a scalar'
+                            % self.__class__.__name__)
+        val = tensor.cast(val, dtype=scalar.upcast(a.dtype, val.dtype))
+        if val.dtype != a.dtype:
+            raise TypeError('%s: type of second parameter must be compatible'
+                          ' with first\'s' % self.__class__.__name__)
+        return gof.Apply(self, [a, val], [a.type()])
+
+    def perform(self, node, inputs, output_storage):
+        if inputs[0].ndim < 2:
+            raise TypeError('%s: first parameter must have at least'
+                            ' two dimensions' % self.__class__.__name__)
+        elif inputs[1].ndim != 0:
+            raise TypeError('%s: second parameter must be a scalar'
+                            % self.__class__.__name__)
+        a = inputs[0].copy()
+        val = inputs[1]
+        numpy.fill_diagonal(a, val)
+        output_storage[0][0] = a
+
+    def grad(self, inp, cost_grad):
+        """
+        Note: The gradient is currently implemented for matrices
+        only.
+        """
+        a, val = inp
+        grad = cost_grad[0]
+        if (a.dtype == 'complex64' or a.dtype == 'complex128' or
+            val.dtype == 'complex64' or val.dtype == 'complex128'):
+            return [None, None]
+        elif a.ndim > 2:
+            raise NotImplementedError('%s: gradient is currently implemented'
+                            ' for matrices only' % self.__class__.__name__)
+        wr_a = fill_diagonal(grad, 0)  # valid for any number of dimensions
+        wr_val = diag(grad).sum()  # diag is only valid for matrices
+        return [wr_a, wr_val]
+
+
+fill_diagonal = FillDiagonal()

theano/tensor/tests/test_extra_ops.py

-import theano
 import numpy as np
-from theano import tensor as T
-from theano.tests import unittest_tools as utt
+import numpy
 
+import theano
+from theano.tests import unittest_tools as utt
 from theano.tensor.extra_ops import *
+from theano import tensor as T
+from theano import tensor, function, scalar
 
 
 class TestBinCountOp(utt.InferShapeTester):
@@ -185,3 +187,79 @@ class TestRepeatOp(utt.InferShapeTester):
         def repeat_(a):
             return RepeatOp()(a, 3)
         utt.verify_grad(repeat_, [a])
+
+
+class TestBartlett(utt.InferShapeTester):
+
+    def setUp(self):
+        super(TestBartlett, self).setUp()
+        self.op_class = Bartlett
+        self.op = bartlett
+
+    def test_perform(self):
+        x = tensor.lscalar()
+        f = function([x], self.op(x))
+        M = numpy.random.random_integers(3, 50, size=())
+        assert numpy.allclose(f(M), numpy.bartlett(M))
+        assert numpy.allclose(f(0), numpy.bartlett(0))
+        assert numpy.allclose(f(-1), numpy.bartlett(-1))
+        b = numpy.array([17], dtype='uint8')
+        assert numpy.allclose(f(b[0]), numpy.bartlett(b[0]))
+
+    def test_infer_shape(self):
+        x = tensor.lscalar()
+        self._compile_and_check([x], [self.op(x)],
+                                [numpy.random.random_integers(3, 50, size=())],
+                                self.op_class)
+
+if __name__ == "__main__":
+    t = TestBartlett('setUp')
+    t.setUp()
+    t.test_perform()
+    t.test_infer_shape()
+
+
+class TestFillDiagonal(utt.InferShapeTester):
+
+    rng = numpy.random.RandomState(43)
+
+    def setUp(self):
+        super(TestFillDiagonal, self).setUp()
+        self.op_class = FillDiagonal
+        self.op = fill_diagonal
+
+    def test_perform(self):
+        x = tensor.dmatrix()
+        y = tensor.dscalar()
+        f = function([x, y], fill_diagonal(x, y))
+        a = numpy.random.rand(8, 5)
+        val = numpy.random.rand()
+        out = f(a, val)
+        numpy.fill_diagonal(a, val)
+        # remember that numpy.fill_diagonal works in place
+        assert numpy.allclose(out, a)
+
+    def test_gradient(self):
+        utt.verify_grad(fill_diagonal, [numpy.random.rand(5, 8),
+                                        numpy.random.rand()],
+                        n_tests=1, rng=TestFillDiagonal.rng)
+        #  TODO: check why gradient wrto val does not match when a has
+        # more rows than cols: might be problem with testing procedure
+        utt.verify_grad(fill_diagonal, [numpy.random.rand(8, 5),
+                                        numpy.random.rand()],
+                        n_tests=1, rng=TestFillDiagonal.rng)
+
+    def test_infer_shape(self):
+        x = tensor.dmatrix()
+        y = tensor.dscalar()
+        self._compile_and_check([x, y], [self.op(x, y)],
+                                [numpy.random.rand(8, 5), numpy.random.rand()],
+                                self.op_class)
+
+if __name__ == "__main__":
+    t = TestFillDiagonal('setUp')
+    t.setUp()
+    t.test_perform()
+    t.test_gradient()
+    t.test_infer_shape()
+