Add unravel_index and ravel_multi_index.

theano/tensor/basic.py

@@ -2588,6 +2588,175 @@ def nonzero_values(a):
     return a.flatten()[flatnonzero(a)]
 
 
+class UnravelIndex(gof.Op):
+    __props__ = ('order',)
+
+    def __init__(self, order='C'):
+        self.order = order
+
+    def make_node(self, indices, dims, ndim):
+        indices = as_tensor_variable(indices)
+        dims = as_tensor_variable(dims)
+
+        if indices.dtype not in int_dtypes:
+            raise TypeError("'%s' object cannot be interpreted as an index" % str(indices.dtype))
+        if dims.dtype not in int_dtypes:
+            raise TypeError("'%s' object cannot be interpreted as an index" % str(dims.dtype))
+        if dims.ndim != 1:
+            raise TypeError("dims must be a 1D array")
+        if not isinstance(ndim, int):
+            raise TypeError("ndim must be an integer")
+
+        return gof.Apply(
+            self, [indices, dims],
+            [TensorType(dtype=indices.dtype, broadcastable=(False,) * indices.ndim)()
+             for i in xrange(ndim)])
+
+    def infer_shape(self, node, input_shapes):
+        return [input_shapes[0]] * len(node.outputs)
+
+    def perform(self, node, inp, out):
+        indices, dims = inp
+        res = np.unravel_index(indices, dims)
+        assert len(res) == len(out)
+        for i in xrange(len(out)):
+            out[i][0] = res[i]
+
+
+def unravel_index(indices, dims, order='C', ndim=None):
+    """
+    Converts a flat index or array of flat indices into a tuple
+    of coordinate arrays.
+
+    This method is similar to the NumPy version, except for the
+    additional ``ndim`` parameter. This parameter is required if
+    the length of ``dims`` cannot be determined automatically, for
+    example because ``dims`` is a Theano vector.
+
+    For example:
+
+    >>> x = T.matrix()
+    >>> idx = T.ivector()
+    >>> unraveled_idx = T.unravel_index(idx, x.shape, ndim=x.ndim)
+
+    Parameters
+    ----------
+    indices : Theano or NumPy array
+        An integer array whose elements are indices into the flattened
+        version of an array of dimensions ``dims``.
+    dims : tuple of ints
+        The shape of the array to use for unraveling ``indices``.
+    order : {'C', 'F'}, optional
+        Determines whether the indices should be viewed as indexing in
+        row-major (C-style) or column-major (Fortran-style) order.
+    ndim : int, optional
+        Specifies the number of dimensions, i.e., the length of
+        ``dims``. This is required if the dimensions cannot be determined
+        from ``dims`` itself, for example, if ``dims`` is a Theano vector.
+
+    Returns
+    -------
+    unraveled_coords : tuple of ndarray
+        Each array in the tuple has the same shape as the ``indices``
+        array.
+
+    See Also
+    --------
+    ravel_multi_index
+
+    """
+    if ndim is None:
+        if isinstance(dims, (tuple, list)):
+            ndim = len(dims)
+        elif isinstance(dims, np.ndarray) and dims.ndim == 1:
+            ndim = dims.shape[0]
+        else:
+            raise TypeError('unravel_index was called with a dimension '
+                            'list with an unspecified length (dim = %s). '
+                            'Use the ndim parameter of unravel_index to '
+                            'set the number of dimensions. ' % str(dims))
+
+    res = UnravelIndex(order=order)(indices, dims, ndim)
+    if ndim == 1:
+        return (res,)
+    else:
+        return tuple(res)
+
+
+class RavelMultiIndex(gof.Op):
+    __props__ = ('mode', 'order')
+
+    def __init__(self, mode='raise', order='C'):
+        self.mode = mode
+        self.order = order
+
+    def make_node(self, *inp):
+        multi_index = [as_tensor_variable(i) for i in inp[:-1]]
+        dims = as_tensor_variable(inp[-1])
+
+        for i in multi_index:
+            if i.dtype not in int_dtypes:
+                raise TypeError("'%s' object cannot be interpreted as an index" % str(i.dtype))
+        if dims.dtype not in int_dtypes:
+            raise TypeError("'%s' object cannot be interpreted as an index" % str(dims.dtype))
+        if dims.ndim != 1:
+            raise TypeError("dims must be a 1D array")
+
+        return gof.Apply(
+            self, multi_index + [dims],
+            [TensorType(dtype=multi_index[0].dtype,
+                        broadcastable=(False,) * multi_index[0].ndim)()])
+
+    def infer_shape(self, node, input_shapes):
+        return [input_shapes[0]]
+
+    def perform(self, node, inp, out):
+        multi_index, dims = inp[:-1], inp[-1]
+        out[0][0] = np.ravel_multi_index(multi_index, dims,
+                                         mode=self.mode, order=self.order)
+
+
+def ravel_multi_index(multi_index, dims, mode='raise', order='C'):
+    """
+    Converts a tuple of index arrays into an array of flat
+    indices, applying boundary modes to the multi-index.
+
+    Parameters
+    ----------
+    multi_index : tuple of Theano or NumPy arrays
+        A tuple of integer arrays, one array for each dimension.
+    dims : tuple of ints
+        The shape of array into which the indices from ``multi_index`` apply.
+    mode : {'raise', 'wrap', 'clip'}, optional
+        Specifies how out-of-bounds indices are handled.  Can specify
+        either one mode or a tuple of modes, one mode per index.
+        * 'raise' -- raise an error (default)
+        * 'wrap' -- wrap around
+        * 'clip' -- clip to the range
+        In 'clip' mode, a negative index which would normally
+        wrap will clip to 0 instead.
+    order : {'C', 'F'}, optional
+        Determines whether the multi-index should be viewed as
+        indexing in row-major (C-style) or column-major
+        (Fortran-style) order.
+
+    Returns
+    -------
+    raveled_indices : Theano array
+        An array of indices into the flattened version of an array
+        of dimensions ``dims``.
+
+    See Also
+    --------
+    unravel_index
+
+    """
+    if not isinstance(multi_index, (tuple, list)):
+        raise TypeError('multi_index must be a tuple or a list.')
+    args = tuple(multi_index) + (dims,)
+    return RavelMultiIndex(mode=mode, order=order)(*args)
+
+
 class Tri(gof.Op):
 
     __props__ = ("dtype",)

theano/tensor/tests/test_basic.py

@@ -50,6 +50,7 @@ from theano.tensor import (
     dtensor3, SpecifyShape, Mean,
     itensor3, Tile, switch, ExtractDiag, AllocDiag,
     nonzero, flatnonzero, nonzero_values,
+    unravel_index, UnravelIndex, ravel_multi_index, RavelMultiIndex,
     stacklists, DimShuffle, hessian, ptp, power,
     swapaxes, choose, Choose, NoneConst, AllocEmpty,
     isclose, allclose, mgrid, ogrid, extract_constant,
@@ -2744,6 +2745,130 @@ class test_nonzero(unittest.TestCase):
         check(rand4d)
 
 
+class test_unravel_index(utt.InferShapeTester):
+    def test_unravel_index(self):
+        def check(shape, index_ndim, order):
+            indices = np.arange(np.product(shape))
+            # test with scalars and higher-dimensional indices
+            if index_ndim == 0:
+                indices = indices[-1]
+            elif index_ndim == 2:
+                indices = indices[:, np.newaxis]
+            indices_symb = theano.shared(indices)
+
+            # reference result
+            ref = np.unravel_index(indices, shape)
+
+            def fn(i, d, nd=None):
+                if nd is None:
+                    return function([], unravel_index(i, d, order=order))
+                else:
+                    return function([], unravel_index(i, d, order=order, ndim=nd))
+
+            # shape given as a tuple
+            f_array_tuple = fn(indices, shape)
+            f_symb_tuple = fn(indices_symb, shape)
+            np.testing.assert_equal(ref, f_array_tuple())
+            np.testing.assert_equal(ref, f_symb_tuple())
+
+            # shape given as an array
+            shape_array = np.array(shape)
+            f_array_array = fn(indices, shape_array)
+            f_symb_array = fn(indices_symb, shape_array)
+            np.testing.assert_equal(ref, f_array_array())
+            np.testing.assert_equal(ref, f_symb_array())
+
+            # shape given as a theano variable
+            shape_symb = theano.shared(shape_array)
+            f_array_symb = fn(indices, shape_symb, len(shape))
+            f_symb_symb = fn(indices_symb, shape_symb, len(shape))
+            np.testing.assert_equal(ref, f_array_symb())
+            np.testing.assert_equal(ref, f_symb_symb())
+
+            # shape testing
+            self._compile_and_check([],
+                                    unravel_index(indices, shape_symb, order=order, ndim=len(shape)),
+                                    [], UnravelIndex)
+
+        for order in ('C', 'F'):
+            for index_ndim in (0, 1, 2):
+                check((3,), index_ndim, order)
+                check((3, 4), index_ndim, order)
+                check((3, 4, 5), index_ndim, order)
+
+        # must specify ndim if length of dims is not fixed
+        self.assertRaises(TypeError, unravel_index, ivector(), ivector())
+
+        # must provide integers
+        self.assertRaises(TypeError, unravel_index, fvector(), (3, 4))
+        self.assertRaises(TypeError, unravel_index, (3, 4), (3.4, 3.2))
+        self.assertRaises(TypeError, unravel_index, (3, 4), (3, 3), ndim=5.4)
+
+        # dims must be a 1D sequence
+        self.assertRaises(TypeError, unravel_index, (3, 4), 3)
+        self.assertRaises(TypeError, unravel_index, (3, 4), ((3, 4),))
+
+
+class test_ravel_multi_index(utt.InferShapeTester):
+    def test_ravel_multi_index(self):
+        def check(shape, index_ndim, mode, order):
+            multi_index = np.unravel_index(np.arange(np.product(shape)), shape, order=order)
+            # create some invalid indices to test the mode
+            if mode in ('wrap', 'clip'):
+                multi_index = (multi_index[0] - 1,) + multi_index[1:]
+            # test with scalars and higher-dimensional indices
+            if index_ndim == 0:
+                multi_index = tuple(i[-1] for i in multi_index)
+            elif index_ndim == 2:
+                multi_index = tuple(i[:, np.newaxis] for i in multi_index)
+            multi_index_symb = [theano.shared(i) for i in multi_index]
+
+            # reference result
+            ref = np.ravel_multi_index(multi_index, shape, mode, order)
+
+            def fn(mi, s):
+                return function([], ravel_multi_index(mi, s, mode, order))
+
+            # shape given as a tuple
+            f_array_tuple = fn(multi_index, shape)
+            f_symb_tuple = fn(multi_index_symb, shape)
+            np.testing.assert_equal(ref, f_array_tuple())
+            np.testing.assert_equal(ref, f_symb_tuple())
+
+            # shape given as an array
+            shape_array = np.array(shape)
+            f_array_array = fn(multi_index, shape_array)
+            f_symb_array = fn(multi_index_symb, shape_array)
+            np.testing.assert_equal(ref, f_array_array())
+            np.testing.assert_equal(ref, f_symb_array())
+
+            # shape given as a theano variable
+            shape_symb = theano.shared(shape_array)
+            f_array_symb = fn(multi_index, shape_symb)
+            f_symb_symb = fn(multi_index_symb, shape_symb)
+            np.testing.assert_equal(ref, f_array_symb())
+            np.testing.assert_equal(ref, f_symb_symb())
+
+            # shape testing
+            self._compile_and_check([],
+                                    [ravel_multi_index(multi_index, shape_symb, mode, order)],
+                                    [], RavelMultiIndex)
+
+        for mode in ('raise', 'wrap', 'clip'):
+            for order in ('C', 'F'):
+                for index_ndim in (0, 1, 2):
+                    check((3,), index_ndim, mode, order)
+                    check((3, 4), index_ndim, mode, order)
+                    check((3, 4, 5), index_ndim, mode, order)
+
+        # must provide integers
+        self.assertRaises(TypeError, ravel_multi_index, (fvector(), ivector()), (3, 4))
+        self.assertRaises(TypeError, ravel_multi_index, ((3, 4), ivector()), (3.4, 3.2))
+
+        # dims must be a 1D sequence
+        self.assertRaises(TypeError, ravel_multi_index, ((3, 4),), ((3, 4),))
+
+
 def test_identity():
     def check(dtype):
         obj = rand_of_dtype((2,), dtype)