~jtaylor/ubuntu/precise/python-numpy/multiarch-fix-818867

« back to all changes in this revision

Viewing changes to numpy/lib/index_tricks.py

  • Committer: Bazaar Package Importer
  • Author(s): Sandro Tosi
  • Date: 2010-10-07 10:19:13 UTC
  • mfrom: (7.1.5 sid)
  • Revision ID: james.westby@ubuntu.com-20101007101913-8b1kmt8ho4upcl9s
Tags: 1:1.4.1-5
* debian/patches/10_use_local_python.org_object.inv_sphinx.diff
  - fixed small typo in description
* debian/patches/changeset_r8364.diff
  - fix memory corruption (double free); thanks to Joseph Barillari for the
    report and to Michael Gilbert for pushing resolution; Closes: #581058

Show diffs side-by-side

added added

removed removed

Lines of Context:
3
3
           'ogrid',
4
4
           'r_', 'c_', 's_',
5
5
           'index_exp', 'ix_',
6
 
           'ndenumerate','ndindex']
 
6
           'ndenumerate','ndindex',
 
7
           'fill_diagonal','diag_indices','diag_indices_from']
7
8
 
8
9
import sys
9
10
import numpy.core.numeric as _nx
10
 
from numpy.core.numeric import asarray, ScalarType, array
 
11
from numpy.core.numeric import ( asarray, ScalarType, array, alltrue, cumprod,
 
12
                                 arange )
11
13
from numpy.core.numerictypes import find_common_type
12
14
import math
13
15
 
14
16
import function_base
15
 
import numpy.core.defmatrix as matrix
 
17
import numpy.matrixlib as matrix
 
18
from function_base import diff
16
19
makemat = matrix.matrix
17
20
 
18
21
# contributed by Stefan van der Walt
19
22
def unravel_index(x,dims):
20
23
    """
21
 
    Convert a flat index into an index tuple for an array of given shape.
 
24
    Convert a flat index to an index tuple for an array of given shape.
22
25
 
23
26
    Parameters
24
27
    ----------
25
28
    x : int
26
29
        Flattened index.
27
 
    dims : shape tuple
28
 
        Input shape.
 
30
    dims : tuple of ints
 
31
        Input shape, the shape of an array into which indexing is
 
32
        required.
 
33
 
 
34
    Returns
 
35
    -------
 
36
    idx : tuple of ints
 
37
        Tuple of the same shape as `dims`, containing the unraveled index.
29
38
 
30
39
    Notes
31
40
    -----
34
43
 
35
44
    Examples
36
45
    --------
37
 
    >>> arr = np.arange(20).reshape(5,4)
 
46
    >>> arr = np.arange(20).reshape(5, 4)
38
47
    >>> arr
39
48
    array([[ 0,  1,  2,  3],
40
49
           [ 4,  5,  6,  7],
72
81
    return tuple(x/dim_prod % dims)
73
82
 
74
83
def ix_(*args):
75
 
    """ Construct an open mesh from multiple sequences.
76
 
 
77
 
    This function takes n 1-d sequences and returns n outputs with n
78
 
    dimensions each such that the shape is 1 in all but one dimension and
79
 
    the dimension with the non-unit shape value cycles through all n
80
 
    dimensions.
81
 
 
82
 
    Using ix_() one can quickly construct index arrays that will index
83
 
    the cross product.
84
 
 
85
 
    a[ix_([1,3,7],[2,5,8])]  returns the array
86
 
 
87
 
    a[1,2]  a[1,5]  a[1,8]
88
 
    a[3,2]  a[3,5]  a[3,8]
89
 
    a[7,2]  a[7,5]  a[7,8]
 
84
    """
 
85
    Construct an open mesh from multiple sequences.
 
86
 
 
87
    This function takes N 1-D sequences and returns N outputs with N
 
88
    dimensions each, such that the shape is 1 in all but one dimension
 
89
    and the dimension with the non-unit shape value cycles through all
 
90
    N dimensions.
 
91
 
 
92
    Using `ix_` one can quickly construct index arrays that will index
 
93
    the cross product. ``a[np.ix_([1,3],[2,5])]`` returns the array
 
94
    ``[[a[1,2] a[1,5]], [a[3,2] a[3,5]]]``.
 
95
 
 
96
    Parameters
 
97
    ----------
 
98
    args : 1-D sequences
 
99
 
 
100
    Returns
 
101
    -------
 
102
    out : tuple of ndarrays
 
103
        N arrays with N dimensions each, with N the number of input
 
104
        sequences. Together these arrays form an open mesh.
 
105
 
 
106
    See Also
 
107
    --------
 
108
    ogrid, mgrid, meshgrid
 
109
 
 
110
    Examples
 
111
    --------
 
112
    >>> a = np.arange(10).reshape(2, 5)
 
113
    >>> a
 
114
    array([[0, 1, 2, 3, 4],
 
115
           [5, 6, 7, 8, 9]])
 
116
    >>> ixgrid = np.ix_([0,1], [2,4])
 
117
    >>> ixgrid
 
118
    (array([[0],
 
119
           [1]]), array([[2, 4]]))
 
120
    >>> ixgrid[0].shape, ixgrid[1].shape
 
121
    ((2, 1), (1, 2))
 
122
    >>> a[ixgrid]
 
123
    array([[2, 4],
 
124
           [7, 9]])
 
125
 
90
126
    """
91
127
    out = []
92
128
    nd = len(args)
107
143
    """
108
144
    Construct a multi-dimensional "meshgrid".
109
145
 
110
 
    grid = nd_grid() creates an instance which will return a mesh-grid
 
146
    ``grid = nd_grid()`` creates an instance which will return a mesh-grid
111
147
    when indexed.  The dimension and number of the output arrays are equal
112
148
    to the number of indexing dimensions.  If the step length is not a
113
149
    complex number, then the stop is not inclusive.
117
153
    number of points to create between the start and stop values, where
118
154
    the stop value **is inclusive**.
119
155
 
120
 
    If instantiated with an argument of sparse=True, the mesh-grid is
 
156
    If instantiated with an argument of ``sparse=True``, the mesh-grid is
121
157
    open (or not fleshed out) so that only one-dimension of each returned
122
 
    argument is greater than 1
 
158
    argument is greater than 1.
 
159
 
 
160
    Parameters
 
161
    ----------
 
162
    sparse : bool, optional
 
163
        Whether the grid is sparse or not. Default is False.
 
164
 
 
165
    Notes
 
166
    -----
 
167
    Two instances of `nd_grid` are made available in the NumPy namespace,
 
168
    `mgrid` and `ogrid`::
 
169
 
 
170
        mgrid = nd_grid(sparse=False)
 
171
        ogrid = nd_grid(sparse=True)
 
172
 
 
173
    Users should use these pre-defined instances instead of using `nd_grid`
 
174
    directly.
123
175
 
124
176
    Examples
125
177
    --------
137
189
            [0, 1, 2, 3, 4]]])
138
190
    >>> mgrid[-1:1:5j]
139
191
    array([-1. , -0.5,  0. ,  0.5,  1. ])
 
192
 
140
193
    >>> ogrid = np.lib.index_tricks.nd_grid(sparse=True)
141
194
    >>> ogrid[0:5,0:5]
142
195
    [array([[0],
215
268
ogrid.__doc__ = None # set in numpy.add_newdocs
216
269
 
217
270
class AxisConcatenator(object):
218
 
    """Translates slice objects to concatenation along an axis.
 
271
    """
 
272
    Translates slice objects to concatenation along an axis.
 
273
 
 
274
    For detailed documentation on usage, see `r_`.
 
275
 
219
276
    """
220
277
    def _retval(self, res):
221
278
        if self.matrix:
338
395
# in help(r_)
339
396
 
340
397
class RClass(AxisConcatenator):
341
 
    """Translates slice objects to concatenation along the first axis.
342
 
 
343
 
    For example:
 
398
    """
 
399
    Translates slice objects to concatenation along the first axis.
 
400
 
 
401
    This is a simple way to build up arrays quickly. There are two use cases.
 
402
 
 
403
    1. If the index expression contains comma separated arrays, then stack
 
404
       them along their first axis.
 
405
    2. If the index expression contains slice notation or scalars then create
 
406
       a 1-D array with a range indicated by the slice notation.
 
407
 
 
408
    If slice notation is used, the syntax ``start:stop:step`` is equivalent
 
409
    to ``np.arange(start, stop, step)`` inside of the brackets. However, if
 
410
    ``step`` is an imaginary number (i.e. 100j) then its integer portion is
 
411
    interpreted as a number-of-points desired and the start and stop are
 
412
    inclusive. In other words ``start:stop:stepj`` is interpreted as
 
413
    ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.
 
414
    After expansion of slice notation, all comma separated sequences are
 
415
    concatenated together.
 
416
 
 
417
    Optional character strings placed as the first element of the index
 
418
    expression can be used to change the output. The strings 'r' or 'c' result
 
419
    in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)
 
420
    matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1
 
421
    (column) matrix is produced. If the result is 2-D then both provide the
 
422
    same matrix result.
 
423
 
 
424
    A string integer specifies which axis to stack multiple comma separated
 
425
    arrays along. A string of two comma-separated integers allows indication
 
426
    of the minimum number of dimensions to force each entry into as the
 
427
    second integer (the axis to concatenate along is still the first integer).
 
428
 
 
429
    A string with three comma-separated integers allows specification of the
 
430
    axis to concatenate along, the minimum number of dimensions to force the
 
431
    entries to, and which axis should contain the start of the arrays which
 
432
    are less than the specified number of dimensions. In other words the third
 
433
    integer allows you to specify where the 1's should be placed in the shape
 
434
    of the arrays that have their shapes upgraded. By default, they are placed
 
435
    in the front of the shape tuple. The third argument allows you to specify
 
436
    where the start of the array should be instead. Thus, a third argument of
 
437
    '0' would place the 1's at the end of the array shape. Negative integers
 
438
    specify where in the new shape tuple the last dimension of upgraded arrays
 
439
    should be placed, so the default is '-1'.
 
440
 
 
441
    Parameters
 
442
    ----------
 
443
    Not a function, so takes no parameters
 
444
 
 
445
 
 
446
    Returns
 
447
    -------
 
448
    A concatenated ndarray or matrix.
 
449
 
 
450
    See Also
 
451
    --------
 
452
    concatenate : Join a sequence of arrays together.
 
453
    c_ : Translates slice objects to concatenation along the second axis.
 
454
 
 
455
    Examples
 
456
    --------
344
457
    >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]
345
458
    array([1, 2, 3, 0, 0, 4, 5, 6])
 
459
    >>> np.r_[-1:1:6j, [0]*3, 5, 6]
 
460
    array([-1. , -0.6, -0.2,  0.2,  0.6,  1. ,  0. ,  0. ,  0. ,  5. ,  6. ])
 
461
 
 
462
    String integers specify the axis to concatenate along or the minimum
 
463
    number of dimensions to force entries into.
 
464
 
 
465
    >>> np.r_['-1', a, a] # concatenate along last axis
 
466
    array([[0, 1, 2, 0, 1, 2],
 
467
           [3, 4, 5, 3, 4, 5]])
 
468
    >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2
 
469
    array([[1, 2, 3],
 
470
           [4, 5, 6]])
 
471
 
 
472
    >>> np.r_['0,2,0', [1,2,3], [4,5,6]]
 
473
    array([[1],
 
474
           [2],
 
475
           [3],
 
476
           [4],
 
477
           [5],
 
478
           [6]])
 
479
    >>> np.r_['1,2,0', [1,2,3], [4,5,6]]
 
480
    array([[1, 4],
 
481
           [2, 5],
 
482
           [3, 6]])
 
483
 
 
484
    Using 'r' or 'c' as a first string argument creates a matrix.
 
485
 
 
486
    >>> np.r_['r',[1,2,3], [4,5,6]]
 
487
    matrix([[1, 2, 3, 4, 5, 6]])
346
488
 
347
489
    """
348
490
    def __init__(self):
351
493
r_ = RClass()
352
494
 
353
495
class CClass(AxisConcatenator):
354
 
    """Translates slice objects to concatenation along the second axis.
355
 
 
356
 
    For example:
 
496
    """
 
497
    Translates slice objects to concatenation along the second axis.
 
498
 
 
499
    This is short-hand for ``np.r_['-1,2,0', index expression]``, which is
 
500
    useful because of its common occurrence. In particular, arrays will be
 
501
    stacked along their last axis after being upgraded to at least 2-D with
 
502
    1's post-pended to the shape (column vectors made out of 1-D arrays).
 
503
 
 
504
    For detailed documentation, see `r_`.
 
505
 
 
506
    Examples
 
507
    --------
357
508
    >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]
358
 
    array([1, 2, 3, 0, 0, 4, 5, 6])
 
509
    array([[1, 2, 3, 0, 0, 4, 5, 6]])
 
510
 
359
511
    """
360
512
    def __init__(self):
361
513
        AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
373
525
    a : ndarray
374
526
      Input array.
375
527
 
 
528
    See Also
 
529
    --------
 
530
    ndindex, flatiter
 
531
 
376
532
    Examples
377
533
    --------
378
 
    >>> a = np.array([[1,2],[3,4]])
 
534
    >>> a = np.array([[1, 2], [3, 4]])
379
535
    >>> for index, x in np.ndenumerate(a):
380
536
    ...     print index, x
381
537
    (0, 0) 1
388
544
        self.iter = asarray(arr).flat
389
545
 
390
546
    def next(self):
 
547
        """
 
548
        Standard iterator method, returns the index tuple and array value.
 
549
 
 
550
        Returns
 
551
        -------
 
552
        coords : tuple of ints
 
553
            The indices of the current iteration.
 
554
        val : scalar
 
555
            The array element of the current iteration.
 
556
 
 
557
        """
391
558
        return self.iter.coords, self.iter.next()
392
559
 
393
560
    def __iter__(self):
399
566
    An N-dimensional iterator object to index arrays.
400
567
 
401
568
    Given the shape of an array, an `ndindex` instance iterates over
402
 
    the N-dimensional index of the array. At each iteration, the index of the
403
 
    last dimension is incremented by one.
 
569
    the N-dimensional index of the array. At each iteration a tuple
 
570
    of indices is returned, the last dimension is iterated over first.
404
571
 
405
572
    Parameters
406
573
    ----------
407
 
    `*args` : integers
408
 
      The size of each dimension in the counter.
 
574
    `*args` : ints
 
575
      The size of each dimension of the array.
 
576
 
 
577
    See Also
 
578
    --------
 
579
    ndenumerate, flatiter
409
580
 
410
581
    Examples
411
582
    --------
412
 
    >>> for index in np.ndindex(3,2,1):
 
583
    >>> for index in np.ndindex(3, 2, 1):
413
584
    ...     print index
414
585
    (0, 0, 0)
415
586
    (0, 1, 0)
442
613
            self._incrementone(axis-1)
443
614
 
444
615
    def ndincr(self):
 
616
        """
 
617
        Increment the multi-dimensional index by one.
 
618
 
 
619
        `ndincr` takes care of the "wrapping around" of the axes.
 
620
        It is called by `ndindex.next` and not normally used directly.
 
621
 
 
622
        """
445
623
        self._incrementone(self.nd-1)
446
624
 
447
625
    def next(self):
 
626
        """
 
627
        Standard iterator method, updates the index and returns the index tuple.
 
628
 
 
629
        Returns
 
630
        -------
 
631
        val : tuple of ints
 
632
            Returns a tuple containing the indices of the current iteration.
 
633
 
 
634
        """
448
635
        if (self.index >= self.total):
449
636
            raise StopIteration
450
637
        val = tuple(self.ind)
473
660
    """
474
661
    A nicer way to build up index tuples for arrays.
475
662
 
 
663
    .. note::
 
664
       Use one of the two predefined instances `index_exp` or `s_`
 
665
       rather than directly using `IndexExpression`.
 
666
 
476
667
    For any index combination, including slicing and axis insertion,
477
 
    'a[indices]' is the same as 'a[index_exp[indices]]' for any
478
 
    array 'a'. However, 'index_exp[indices]' can be used anywhere
 
668
    ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any
 
669
    array `a`. However, ``np.index_exp[indices]`` can be used anywhere
479
670
    in Python code and returns a tuple of slice objects that can be
480
671
    used in the construction of complex index expressions.
 
672
 
 
673
    Parameters
 
674
    ----------
 
675
    maketuple : bool
 
676
        If True, always returns a tuple.
 
677
 
 
678
    See Also
 
679
    --------
 
680
    index_exp : Predefined instance that always returns a tuple:
 
681
       `index_exp = IndexExpression(maketuple=True)`.
 
682
    s_ : Predefined instance without tuple conversion:
 
683
       `s_ = IndexExpression(maketuple=False)`.
 
684
 
 
685
    Notes
 
686
    -----
 
687
    You can do all this with `slice()` plus a few special objects,
 
688
    but there's a lot to remember and this version is simpler because
 
689
    it uses the standard array indexing syntax.
 
690
 
 
691
    Examples
 
692
    --------
 
693
    >>> np.s_[2::2]
 
694
    slice(2, None, 2)
 
695
    >>> np.index_exp[2::2]
 
696
    (slice(2, None, 2),)
 
697
 
 
698
    >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]
 
699
    array([2, 4])
 
700
 
481
701
    """
482
702
    maxint = sys.maxint
483
703
    def __init__(self, maketuple):
501
721
s_ = IndexExpression(maketuple=False)
502
722
 
503
723
# End contribution from Konrad.
 
724
 
 
725
 
 
726
# The following functions complement those in twodim_base, but are
 
727
# applicable to N-dimensions.
 
728
 
 
729
def fill_diagonal(a, val):
 
730
    """
 
731
    Fill the main diagonal of the given array of any dimensionality.
 
732
 
 
733
    For an array `a` with ``a.ndim > 2``, the diagonal is the list of
 
734
    locations with indices ``a[i, i, ..., i]`` all identical. This function
 
735
    modifies the input array in-place, it does not return a value.
 
736
 
 
737
    Parameters
 
738
    ----------
 
739
    a : array, at least 2-D.
 
740
      Array whose diagonal is to be filled, it gets modified in-place.
 
741
 
 
742
    val : scalar
 
743
      Value to be written on the diagonal, its type must be compatible with
 
744
      that of the array a.
 
745
 
 
746
    See also
 
747
    --------
 
748
    diag_indices, diag_indices_from
 
749
 
 
750
    Notes
 
751
    -----
 
752
    .. versionadded:: 1.4.0
 
753
 
 
754
    This functionality can be obtained via `diag_indices`, but internally
 
755
    this version uses a much faster implementation that never constructs the
 
756
    indices and uses simple slicing.
 
757
 
 
758
    Examples
 
759
    --------
 
760
    >>> a = zeros((3, 3), int)
 
761
    >>> fill_diagonal(a, 5)
 
762
    >>> a
 
763
    array([[5, 0, 0],
 
764
           [0, 5, 0],
 
765
           [0, 0, 5]])
 
766
 
 
767
    The same function can operate on a 4-D array:
 
768
 
 
769
    >>> a = zeros((3, 3, 3, 3), int)
 
770
    >>> fill_diagonal(a, 4)
 
771
 
 
772
    We only show a few blocks for clarity:
 
773
 
 
774
    >>> a[0, 0]
 
775
    array([[4, 0, 0],
 
776
           [0, 0, 0],
 
777
           [0, 0, 0]])
 
778
    >>> a[1, 1]
 
779
    array([[0, 0, 0],
 
780
           [0, 4, 0],
 
781
           [0, 0, 0]])
 
782
    >>> a[2, 2]
 
783
    array([[0, 0, 0],
 
784
           [0, 0, 0],
 
785
           [0, 0, 4]])
 
786
 
 
787
    """
 
788
    if a.ndim < 2:
 
789
        raise ValueError("array must be at least 2-d")
 
790
    if a.ndim == 2:
 
791
        # Explicit, fast formula for the common case.  For 2-d arrays, we
 
792
        # accept rectangular ones.
 
793
        step = a.shape[1] + 1
 
794
    else:
 
795
        # For more than d=2, the strided formula is only valid for arrays with
 
796
        # all dimensions equal, so we check first.
 
797
        if not alltrue(diff(a.shape)==0):
 
798
            raise ValueError("All dimensions of input must be of equal length")
 
799
        step = 1 + (cumprod(a.shape[:-1])).sum()
 
800
 
 
801
    # Write the value out into the diagonal.
 
802
    a.flat[::step] = val
 
803
 
 
804
 
 
805
def diag_indices(n, ndim=2):
 
806
    """
 
807
    Return the indices to access the main diagonal of an array.
 
808
 
 
809
    This returns a tuple of indices that can be used to access the main
 
810
    diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
 
811
    (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for
 
812
    ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``
 
813
    for ``i = [0..n-1]``.
 
814
 
 
815
    Parameters
 
816
    ----------
 
817
    n : int
 
818
      The size, along each dimension, of the arrays for which the returned
 
819
      indices can be used.
 
820
 
 
821
    ndim : int, optional
 
822
      The number of dimensions.
 
823
 
 
824
    See also
 
825
    --------
 
826
    diag_indices_from
 
827
 
 
828
    Notes
 
829
    -----
 
830
    .. versionadded:: 1.4.0
 
831
 
 
832
    Examples
 
833
    --------
 
834
    Create a set of indices to access the diagonal of a (4, 4) array:
 
835
 
 
836
    >>> di = np.diag_indices(4)
 
837
    >>> di
 
838
    (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
 
839
    >>> a = np.arange(16).reshape(4, 4)
 
840
    >>> a
 
841
    array([[ 0,  1,  2,  3],
 
842
           [ 4,  5,  6,  7],
 
843
           [ 8,  9, 10, 11],
 
844
           [12, 13, 14, 15]])
 
845
    >>> a[di] = 100
 
846
    >>> a
 
847
    array([[100,   1,   2,   3],
 
848
           [  4, 100,   6,   7],
 
849
           [  8,   9, 100,  11],
 
850
           [ 12,  13,  14, 100]])
 
851
 
 
852
    Now, we create indices to manipulate a 3-D array:
 
853
 
 
854
    >>> d3 = np.diag_indices(2, 3)
 
855
    >>> d3
 
856
    (array([0, 1]), array([0, 1]), array([0, 1]))
 
857
 
 
858
    And use it to set the diagonal of an array of zeros to 1:
 
859
 
 
860
    >>> a = np.zeros((2, 2, 2), dtype=np.int)
 
861
    >>> a[d3] = 1
 
862
    >>> a
 
863
    array([[[1, 0],
 
864
            [0, 0]],
 
865
           [[0, 0],
 
866
            [0, 1]]])
 
867
 
 
868
    """
 
869
    idx = arange(n)
 
870
    return (idx,) * ndim
 
871
 
 
872
 
 
873
def diag_indices_from(arr):
 
874
    """
 
875
    Return the indices to access the main diagonal of an n-dimensional array.
 
876
 
 
877
    See `diag_indices` for full details.
 
878
 
 
879
    Parameters
 
880
    ----------
 
881
    arr : array, at least 2-D
 
882
 
 
883
    See Also
 
884
    --------
 
885
    diag_indices
 
886
 
 
887
    Notes
 
888
    -----
 
889
    .. versionadded:: 1.4.0
 
890
 
 
891
    """
 
892
 
 
893
    if not arr.ndim >= 2:
 
894
        raise ValueError("input array must be at least 2-d")
 
895
    # For more than d=2, the strided formula is only valid for arrays with
 
896
    # all dimensions equal, so we check first.
 
897
    if not alltrue(diff(arr.shape) == 0):
 
898
        raise ValueError("All dimensions of input must be of equal length")
 
899
 
 
900
    return diag_indices(arr.shape[0], arr.ndim)