Add indexing functions by Fernando Perez.

4 files changed, 506 insertions, 23 deletions

diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
index b6eaae29f..737fc0a60 100644
--- a/numpy/lib/index_tricks.py
+++ b/numpy/lib/index_tricks.py
@@ -3,16 +3,19 @@ __all__ = ['unravel_index',
            'ogrid',
            'r_', 'c_', 's_',
            'index_exp', 'ix_',
-           'ndenumerate','ndindex']
+           'ndenumerate','ndindex',
+           'fill_diagonal','diag_indices','diag_indices_from']
 
 import sys
 import numpy.core.numeric as _nx
-from numpy.core.numeric import asarray, ScalarType, array
+from numpy.core.numeric import ( asarray, ScalarType, array, alltrue, cumprod,
+                                 arange )
 from numpy.core.numerictypes import find_common_type
 import math
 
 import function_base
 import numpy.core.defmatrix as matrix
+from function_base import diff
 makemat = matrix.matrix
 
 # contributed by Stefan van der Walt
@@ -665,3 +668,153 @@ index_exp = IndexExpression(maketuple=True)
 s_ = IndexExpression(maketuple=False)
 
 # End contribution from Konrad.
+
+# I'm not sure this is the best place in numpy for these functions, but since
+# they handle multidimensional arrays, it seemed better than twodim_base.
+
+def fill_diagonal(a,val):
+    """Fill the main diagonal of the given array of any dimensionality.
+
+    For an array with ndim > 2, the diagonal is the list of locations with
+    indices a[i,i,...,i], all identical.
+
+    This function modifies the input array in-place, it does not return a
+    value.
+
+    This functionality can be obtained via diag_indices(), but internally this
+    version uses a much faster implementation that never constructs the indices
+    and uses simple slicing.
+
+    Parameters
+    ----------
+    a : array, at least 2-dimensional.
+      Array whose diagonal is to be filled, it gets modified in-place.
+
+    val : scalar
+      Value to be written on the diagonal, its type must be compatible with
+      that of the array a.
+
+    Examples
+    --------
+    >>> a = zeros((3,3),int)
+    >>> fill_diagonal(a,5)
+    >>> a
+    array([[5, 0, 0],
+           [0, 5, 0],
+           [0, 0, 5]])
+
+    The same function can operate on a 4-d array:
+    >>> a = zeros((3,3,3,3),int)
+    >>> fill_diagonal(a,4)
+
+    We only show a few blocks for clarity:
+    >>> a[0,0]
+    array([[4, 0, 0],
+           [0, 0, 0],
+           [0, 0, 0]])
+    >>> a[1,1]
+    array([[0, 0, 0],
+           [0, 4, 0],
+           [0, 0, 0]])
+    >>> a[2,2]
+    array([[0, 0, 0],
+           [0, 0, 0],
+           [0, 0, 4]])
+
+    See also
+    --------
+    - diag_indices: indices to access diagonals given shape information.
+    - diag_indices_from: indices to access diagonals given an array.
+    """
+    if a.ndim < 2:
+        raise ValueError("array must be at least 2-d")
+    if a.ndim == 2:
+        # Explicit, fast formula for the common case.  For 2-d arrays, we
+        # accept rectangular ones.
+        step = a.shape[1] + 1
+    else:
+        # For more than d=2, the strided formula is only valid for arrays with
+        # all dimensions equal, so we check first.
+        if not alltrue(diff(a.shape)==0):
+            raise ValueError("All dimensions of input must be of equal length")
+        step = cumprod((1,)+a.shape[:-1]).sum()
+
+    # Write the value out into the diagonal.
+    a.flat[::step] = val
+
+
+def diag_indices(n,ndim=2):
+    """Return the indices to access the main diagonal of an array.
+
+    This returns a tuple of indices that can be used to access the main
+    diagonal of an array with ndim (>=2) dimensions and shape (n,n,...,n).  For
+    ndim=2 this is the usual diagonal, for ndim>2 this is the set of indices
+    to access A[i,i,...,i] for i=[0..n-1].
+
+    Parameters
+    ----------
+    n : int
+      The size, along each dimension, of the arrays for which the returned
+      indices can be used.
+
+    ndim : int, optional
+      The number of dimensions 
+
+    Examples
+    --------
+    Create a set of indices to access the diagonal of a (4,4) array:
+    >>> di = diag_indices(4)
+
+    >>> a = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])
+    >>> a
+    array([[ 1,  2,  3,  4],
+           [ 5,  6,  7,  8],
+           [ 9, 10, 11, 12],
+           [13, 14, 15, 16]])
+    >>> a[di] = 100
+    >>> a
+    array([[100,   2,   3,   4],
+           [  5, 100,   7,   8],
+           [  9,  10, 100,  12],
+           [ 13,  14,  15, 100]])
+
+    Now, we create indices to manipulate a 3-d array:
+    >>> d3 = diag_indices(2,3)
+
+    And use it to set the diagonal of a zeros array to 1:
+    >>> a = zeros((2,2,2),int)
+    >>> a[d3] = 1
+    >>> a
+    array([[[1, 0],
+            [0, 0]],
+
+           [[0, 0],
+            [0, 1]]])
+
+    See also
+    --------
+    - diag_indices_from: create the indices based on the shape of an existing
+    array. 
+    """
+    idx = arange(n)
+    return (idx,)*ndim
+
+
+def diag_indices_from(arr):
+    """Return the indices to access the main diagonal of an n-dimensional array.
+
+    See diag_indices() for full details.
+
+    Parameters
+    ----------
+    arr : array, at least 2-d
+    """
+
+    if not arr.ndim >= 2:
+        raise ValueError("input array must be at least 2-d")
+    # For more than d=2, the strided formula is only valid for arrays with
+    # all dimensions equal, so we check first.
+    if not alltrue(diff(a.shape)==0):
+        raise ValueError("All dimensions of input must be of equal length")
+
+    return diag_indices(a.shape[0],a.ndim)
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py
index 641737d43..50dc5ac15 100644
--- a/numpy/lib/tests/test_index_tricks.py
+++ b/numpy/lib/tests/test_index_tricks.py
@@ -1,5 +1,6 @@
 from numpy.testing import *
-from numpy import array, ones, r_, mgrid, unravel_index, ndenumerate
+from numpy import ( array, ones, r_, mgrid, unravel_index, zeros, where,
+                    fill_diagonal, diag_indices, diag_indices_from )
 
 class TestUnravelIndex(TestCase):
     def test_basic(self):
@@ -68,5 +69,48 @@ class TestNdenumerate(TestCase):
         assert_equal(list(ndenumerate(a)),
                      [((0,0), 1), ((0,1), 2), ((1,0), 3), ((1,1), 4)])
 
+
+def test_fill_diagonal():
+    a = zeros((3, 3),int)
+    fill_diagonal(a, 5)
+    yield (assert_array_equal, a,
+           array([[5, 0, 0],
+                  [0, 5, 0],
+                  [0, 0, 5]]))
+
+    # The same function can operate on a 4-d array:
+    a = zeros((3, 3, 3, 3), int)
+    fill_diagonal(a, 4)
+    i = array([0, 1, 2])
+    yield (assert_equal, where(a != 0), (i, i, i, i))
+
+
+def test_diag_indices():
+    di = diag_indices(4)
+    a = array([[1, 2, 3, 4],
+               [5, 6, 7, 8],
+               [9, 10, 11, 12],
+               [13, 14, 15, 16]])
+    a[di] = 100
+    yield (assert_array_equal, a,
+           array([[100,   2,   3,   4],
+                  [  5, 100,   7,   8],
+                  [  9,  10, 100,  12],
+                  [ 13,  14,  15, 100]]))
+
+    # Now, we create indices to manipulate a 3-d array:
+    d3 = diag_indices(2, 3)
+
+    # And use it to set the diagonal of a zeros array to 1:
+    a = zeros((2, 2, 2),int)
+    a[d3] = 1
+    yield (assert_array_equal, a,
+           array([[[1, 0],
+                   [0, 0]],
+
+                  [[0, 0],
+                   [0, 1]]]) )
+
+
 if __name__ == "__main__":
     run_module_suite()
diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py
index 32c4ca58e..a1e58d53b 100644
--- a/numpy/lib/tests/test_twodim_base.py
+++ b/numpy/lib/tests/test_twodim_base.py
@@ -3,8 +3,11 @@
 """
 
 from numpy.testing import *
-from numpy import arange, rot90, add, fliplr, flipud, zeros, ones, eye, \
-     array, diag, histogram2d, tri
+
+from numpy import ( arange, rot90, add, fliplr, flipud, zeros, ones, eye,
+                    array, diag, histogram2d, tri, mask_indices, triu_indices,
+                    triu_indices_from, tril_indices, tril_indices_from )
+
 import numpy as np
 
 def get_mat(n):
@@ -52,32 +55,32 @@ class TestEye(TestCase):
 
 class TestDiag(TestCase):
     def test_vector(self):
-        vals = (100*arange(5)).astype('l')
-        b = zeros((5,5))
+        vals = (100 * arange(5)).astype('l')
+        b = zeros((5, 5))
         for k in range(5):
-            b[k,k] = vals[k]
-        assert_equal(diag(vals),b)
-        b = zeros((7,7))
+            b[k, k] = vals[k]
+        assert_equal(diag(vals), b)
+        b = zeros((7, 7))
         c = b.copy()
         for k in range(5):
-            b[k,k+2] = vals[k]
-            c[k+2,k] = vals[k]
-        assert_equal(diag(vals,k=2), b)
-        assert_equal(diag(vals,k=-2), c)
+            b[k, k + 2] = vals[k]
+            c[k + 2, k] = vals[k]
+        assert_equal(diag(vals, k=2), b)
+        assert_equal(diag(vals, k=-2), c)
 
     def test_matrix(self):
-        vals = (100*get_mat(5)+1).astype('l')
+        vals = (100 * get_mat(5) + 1).astype('l')
         b = zeros((5,))
         for k in range(5):
             b[k] = vals[k,k]
-        assert_equal(diag(vals),b)
-        b = b*0
+        assert_equal(diag(vals), b)
+        b = b * 0
         for k in range(3):
-            b[k] = vals[k,k+2]
-        assert_equal(diag(vals,2),b[:3])
+            b[k] = vals[k, k + 2]
+        assert_equal(diag(vals, 2), b[:3])
         for k in range(3):
-            b[k] = vals[k+2,k]
-        assert_equal(diag(vals,-2),b[:3])
+            b[k] = vals[k + 2, k]
+        assert_equal(diag(vals, -2), b[:3])
 
 class TestFliplr(TestCase):
     def test_basic(self):
@@ -193,5 +196,76 @@ class TestTri(TestCase):
         assert_array_equal(tri(3,dtype=bool),out.astype(bool))
 
 
+def test_mask_indices():
+    # simple test without offset
+    iu = mask_indices(3, np.triu)
+    a = np.arange(9).reshape(3, 3)
+    yield (assert_array_equal, a[iu], array([0, 1, 2, 4, 5, 8]))
+    # Now with an offset
+    iu1 = mask_indices(3, np.triu, 1)
+    yield (assert_array_equal, a[iu1], array([1, 2, 5]))
+
+
+def test_tril_indices():
+    # indices without and with offset
+    il1 = tril_indices(4)
+    il2 = tril_indices(4, 2)
+
+    a = np.array([[1, 2, 3, 4],
+                  [5, 6, 7, 8],
+                  [9, 10, 11, 12],
+                  [13, 14, 15, 16]])
+
+    # indexing:
+    yield (assert_array_equal, a[il1],
+           array([ 1,  5,  6,  9, 10, 11, 13, 14, 15, 16]) )
+
+    # And for assigning values:
+    a[il1] = -1
+    yield (assert_array_equal, a,
+    array([[-1,  2,  3,  4],
+           [-1, -1,  7,  8],
+           [-1, -1, -1, 12],
+           [-1, -1, -1, -1]]) )
+
+    # These cover almost the whole array (two diagonals right of the main one):
+    a[il2] = -10
+    yield (assert_array_equal, a,
+    array([[-10, -10, -10,   4],
+           [-10, -10, -10, -10],
+           [-10, -10, -10, -10],
+           [-10, -10, -10, -10]]) )
+
+
+def test_triu_indices():
+    iu1 = triu_indices(4)
+    iu2 = triu_indices(4, 2)
+
+    a = np.array([[1, 2, 3, 4],
+                  [5, 6, 7, 8],
+                  [9, 10, 11, 12],
+                  [13, 14, 15, 16]])
+
+    # Both for indexing:
+    yield (assert_array_equal, a[iu1],
+           array([ 1,  5,  6,  9, 10, 11, 13, 14, 15, 16]) )
+
+    # And for assigning values:
+    a[iu1] = -1
+    yield (assert_array_equal, a,
+    array([[-1, -1, -1, -1],
+           [ 5, -1, -1, -1],
+           [ 9, 10, -1, -1],
+           [13, 14, 15, -1]])  )
+
+    # These cover almost the whole array (two diagonals right of the main one):
+    a[iu2] = -10
+    yield ( assert_array_equal, a,
+    array([[ -1,  -1, -10, -10],
+           [  5,  -1,  -1, -10],
+           [  9,  10,  -1,  -1],
+           [ 13,  14,  15,  -1]]) )
+
+
 if __name__ == "__main__":
     run_module_suite()
diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py
index f0abf3122..533408887 100644
--- a/numpy/lib/twodim_base.py
+++ b/numpy/lib/twodim_base.py
@@ -3,10 +3,12 @@
 """
 
 __all__ = ['diag','diagflat','eye','fliplr','flipud','rot90','tri','triu',
-           'tril','vander','histogram2d']
+           'tril','vander','histogram2d','mask_indices',
+           'tril_indices','tril_indices_from','triu_indices','triu_indices_from',
+           ]
 
 from numpy.core.numeric import asanyarray, equal, subtract, arange, \
-     zeros, greater_equal, multiply, ones, asarray
+     zeros, greater_equal, multiply, ones, asarray, alltrue, where
 
 def fliplr(m):
     """
@@ -559,3 +561,213 @@ def histogram2d(x,y, bins=10, range=None, normed=False, weights=None):
         bins = [xedges, yedges]
     hist, edges = histogramdd([x,y], bins, range, normed, weights)
     return hist, edges[0], edges[1]
+
+    
+def mask_indices(n,mask_func,k=0):
+    """Return the indices to access (n,n) arrays, given a masking function.
+
+    Assume mask_func() is a function that, for a square array a of size (n,n)
+    with a possible offset argument k, when called as mask_func(a,k) returns a
+    new array with zeros in certain locations (functions like triu() or tril()
+    do precisely this).  Then this function returns the indices where the
+    non-zero values would be located.
+
+    Parameters
+    ----------
+    n : int
+      The returned indices will be valid to access arrays of shape (n,n).
+
+    mask_func : callable
+      A function whose api is similar to that of numpy.tri{u,l}.  That is,
+      mask_func(x,k) returns a boolean array, shaped like x.  k is an optional
+      argument to the function.
+
+    k : scalar
+      An optional argument which is passed through to mask_func().  Functions
+      like tri{u,l} take a second argument that is interpreted as an offset.
+
+    Returns
+    -------
+    indices : an n-tuple of index arrays.
+      The indices corresponding to the locations where mask_func(ones((n,n)),k)
+      is True.
+
+    Examples
+    --------
+    These are the indices that would allow you to access the upper triangular
+    part of any 3x3 array:
+    >>> iu = mask_indices(3,np.triu)
+
+    For example, if `a` is a 3x3 array:
+    >>> a = np.arange(9).reshape(3,3)
+    >>> a
+    array([[0, 1, 2],
+           [3, 4, 5],
+           [6, 7, 8]])
+
+    Then:
+    >>> a[iu]
+    array([0, 1, 2, 4, 5, 8])
+
+    An offset can be passed also to the masking function.  This gets us the
+    indices starting on the first diagonal right of the main one:
+    >>> iu1 = mask_indices(3,np.triu,1)
+
+    with which we now extract only three elements:
+    >>> a[iu1]
+    array([1, 2, 5])
+    """ 
+    m = ones((n,n),int)
+    a = mask_func(m,k)
+    return where(a != 0)
+
+
+def tril_indices(n,k=0):
+    """Return the indices for the lower-triangle of an (n,n) array.
+
+    Parameters
+    ----------
+    n : int
+      Sets the size of the arrays for which the returned indices will be valid.
+
+    k : int, optional
+      Diagonal offset (see tril() for details).
+
+    Examples
+    --------
+    Commpute two different sets of indices to access 4x4 arrays, one for the
+    lower triangular part starting at the main diagonal, and one starting two
+    diagonals further right:
+    
+    >>> il1 = tril_indices(4)
+    >>> il2 = tril_indices(4,2)
+
+    Here is how they can be used with a sample array:
+    >>> a = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])
+    >>> a
+    array([[ 1,  2,  3,  4],
+           [ 5,  6,  7,  8],
+           [ 9, 10, 11, 12],
+           [13, 14, 15, 16]])
+
+    Both for indexing:
+    >>> a[il1]
+    array([ 1,  5,  6,  9, 10, 11, 13, 14, 15, 16])
+
+    And for assigning values:
+    >>> a[il1] = -1
+    >>> a
+    array([[-1,  2,  3,  4],
+           [-1, -1,  7,  8],
+           [-1, -1, -1, 12],
+           [-1, -1, -1, -1]])
+
+    These cover almost the whole array (two diagonals right of the main one):
+    >>> a[il2] = -10 
+    >>> a
+    array([[-10, -10, -10,   4],
+           [-10, -10, -10, -10],
+           [-10, -10, -10, -10],
+           [-10, -10, -10, -10]])
+
+    See also
+    --------
+    - triu_indices : similar function, for upper-triangular.
+    - mask_indices : generic function accepting an arbitrary mask function.
+    """
+    return mask_indices(n,tril,k)
+
+
+def tril_indices_from(arr,k=0):
+    """Return the indices for the lower-triangle of an (n,n) array.
+
+    See tril_indices() for full details.
+    
+    Parameters
+    ----------
+    n : int
+      Sets the size of the arrays for which the returned indices will be valid.
+
+    k : int, optional
+      Diagonal offset (see tril() for details).
+
+    """
+    if not arr.ndim==2 and arr.shape[0] == arr.shape[1]:
+        raise ValueError("input array must be 2-d and square")
+    return tril_indices(arr.shape[0],k)
+
+    
+def triu_indices(n,k=0):
+    """Return the indices for the upper-triangle of an (n,n) array.
+
+    Parameters
+    ----------
+    n : int
+      Sets the size of the arrays for which the returned indices will be valid.
+
+    k : int, optional
+      Diagonal offset (see triu() for details).
+
+    Examples
+    --------
+    Commpute two different sets of indices to access 4x4 arrays, one for the
+    lower triangular part starting at the main diagonal, and one starting two
+    diagonals further right:
+
+    >>> iu1 = triu_indices(4)
+    >>> iu2 = triu_indices(4,2)
+
+    Here is how they can be used with a sample array:
+    >>> a = np.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])
+    >>> a
+    array([[ 1,  2,  3,  4],
+           [ 5,  6,  7,  8],
+           [ 9, 10, 11, 12],
+           [13, 14, 15, 16]])
+
+    Both for indexing:
+    >>> a[il1]
+    array([ 1,  5,  6,  9, 10, 11, 13, 14, 15, 16])
+
+    And for assigning values:       
+    >>> a[iu] = -1
+    >>> a
+    array([[-1, -1, -1, -1],
+           [ 5, -1, -1, -1],
+           [ 9, 10, -1, -1],
+           [13, 14, 15, -1]])
+
+    These cover almost the whole array (two diagonals right of the main one):
+    >>> a[iu2] = -10
+    >>> a
+    array([[ -1,  -1, -10, -10],
+           [  5,  -1,  -1, -10],
+           [  9,  10,  -1,  -1],
+           [ 13,  14,  15,  -1]])
+
+    See also
+    --------
+    - tril_indices : similar function, for lower-triangular.
+    - mask_indices : generic function accepting an arbitrary mask function.
+    """
+    return mask_indices(n,triu,k)
+
+
+def triu_indices_from(arr,k=0):
+    """Return the indices for the lower-triangle of an (n,n) array.
+
+    See triu_indices() for full details.
+    
+    Parameters
+    ----------
+    n : int
+      Sets the size of the arrays for which the returned indices will be valid.
+
+    k : int, optional
+      Diagonal offset (see triu() for details).
+
+    """
+    if not arr.ndim==2 and arr.shape[0] == arr.shape[1]:
+        raise ValueError("input array must be 2-d and square")
+    return triu_indices(arr.shape[0],k)
+