Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/numpy/lib/index_tricks.py : 27%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import functools
2import sys
3import math
5import numpy.core.numeric as _nx
6from numpy.core.numeric import (
7 asarray, ScalarType, array, alltrue, cumprod, arange, ndim
8 )
9from numpy.core.numerictypes import find_common_type, issubdtype
11import numpy.matrixlib as matrixlib
12from .function_base import diff
13from numpy.core.multiarray import ravel_multi_index, unravel_index
14from numpy.core.overrides import set_module
15from numpy.core import overrides, linspace
16from numpy.lib.stride_tricks import as_strided
19array_function_dispatch = functools.partial(
20 overrides.array_function_dispatch, module='numpy')
23__all__ = [
24 'ravel_multi_index', 'unravel_index', 'mgrid', 'ogrid', 'r_', 'c_',
25 's_', 'index_exp', 'ix_', 'ndenumerate', 'ndindex', 'fill_diagonal',
26 'diag_indices', 'diag_indices_from'
27 ]
30def _ix__dispatcher(*args):
31 return args
34@array_function_dispatch(_ix__dispatcher)
35def ix_(*args):
36 """
37 Construct an open mesh from multiple sequences.
39 This function takes N 1-D sequences and returns N outputs with N
40 dimensions each, such that the shape is 1 in all but one dimension
41 and the dimension with the non-unit shape value cycles through all
42 N dimensions.
44 Using `ix_` one can quickly construct index arrays that will index
45 the cross product. ``a[np.ix_([1,3],[2,5])]`` returns the array
46 ``[[a[1,2] a[1,5]], [a[3,2] a[3,5]]]``.
48 Parameters
49 ----------
50 args : 1-D sequences
51 Each sequence should be of integer or boolean type.
52 Boolean sequences will be interpreted as boolean masks for the
53 corresponding dimension (equivalent to passing in
54 ``np.nonzero(boolean_sequence)``).
56 Returns
57 -------
58 out : tuple of ndarrays
59 N arrays with N dimensions each, with N the number of input
60 sequences. Together these arrays form an open mesh.
62 See Also
63 --------
64 ogrid, mgrid, meshgrid
66 Examples
67 --------
68 >>> a = np.arange(10).reshape(2, 5)
69 >>> a
70 array([[0, 1, 2, 3, 4],
71 [5, 6, 7, 8, 9]])
72 >>> ixgrid = np.ix_([0, 1], [2, 4])
73 >>> ixgrid
74 (array([[0],
75 [1]]), array([[2, 4]]))
76 >>> ixgrid[0].shape, ixgrid[1].shape
77 ((2, 1), (1, 2))
78 >>> a[ixgrid]
79 array([[2, 4],
80 [7, 9]])
82 >>> ixgrid = np.ix_([True, True], [2, 4])
83 >>> a[ixgrid]
84 array([[2, 4],
85 [7, 9]])
86 >>> ixgrid = np.ix_([True, True], [False, False, True, False, True])
87 >>> a[ixgrid]
88 array([[2, 4],
89 [7, 9]])
91 """
92 out = []
93 nd = len(args)
94 for k, new in enumerate(args):
95 if not isinstance(new, _nx.ndarray):
96 new = asarray(new)
97 if new.size == 0:
98 # Explicitly type empty arrays to avoid float default
99 new = new.astype(_nx.intp)
100 if new.ndim != 1:
101 raise ValueError("Cross index must be 1 dimensional")
102 if issubdtype(new.dtype, _nx.bool_):
103 new, = new.nonzero()
104 new = new.reshape((1,)*k + (new.size,) + (1,)*(nd-k-1))
105 out.append(new)
106 return tuple(out)
108class nd_grid:
109 """
110 Construct a multi-dimensional "meshgrid".
112 ``grid = nd_grid()`` creates an instance which will return a mesh-grid
113 when indexed. The dimension and number of the output arrays are equal
114 to the number of indexing dimensions. If the step length is not a
115 complex number, then the stop is not inclusive.
117 However, if the step length is a **complex number** (e.g. 5j), then the
118 integer part of its magnitude is interpreted as specifying the
119 number of points to create between the start and stop values, where
120 the stop value **is inclusive**.
122 If instantiated with an argument of ``sparse=True``, the mesh-grid is
123 open (or not fleshed out) so that only one-dimension of each returned
124 argument is greater than 1.
126 Parameters
127 ----------
128 sparse : bool, optional
129 Whether the grid is sparse or not. Default is False.
131 Notes
132 -----
133 Two instances of `nd_grid` are made available in the NumPy namespace,
134 `mgrid` and `ogrid`, approximately defined as::
136 mgrid = nd_grid(sparse=False)
137 ogrid = nd_grid(sparse=True)
139 Users should use these pre-defined instances instead of using `nd_grid`
140 directly.
141 """
143 def __init__(self, sparse=False):
144 self.sparse = sparse
146 def __getitem__(self, key):
147 try:
148 size = []
149 typ = int
150 for k in range(len(key)):
151 step = key[k].step
152 start = key[k].start
153 if start is None:
154 start = 0
155 if step is None:
156 step = 1
157 if isinstance(step, complex):
158 size.append(int(abs(step)))
159 typ = float
160 else:
161 size.append(
162 int(math.ceil((key[k].stop - start)/(step*1.0))))
163 if (isinstance(step, float) or
164 isinstance(start, float) or
165 isinstance(key[k].stop, float)):
166 typ = float
167 if self.sparse:
168 nn = [_nx.arange(_x, dtype=_t)
169 for _x, _t in zip(size, (typ,)*len(size))]
170 else:
171 nn = _nx.indices(size, typ)
172 for k in range(len(size)):
173 step = key[k].step
174 start = key[k].start
175 if start is None:
176 start = 0
177 if step is None:
178 step = 1
179 if isinstance(step, complex):
180 step = int(abs(step))
181 if step != 1:
182 step = (key[k].stop - start)/float(step-1)
183 nn[k] = (nn[k]*step+start)
184 if self.sparse:
185 slobj = [_nx.newaxis]*len(size)
186 for k in range(len(size)):
187 slobj[k] = slice(None, None)
188 nn[k] = nn[k][tuple(slobj)]
189 slobj[k] = _nx.newaxis
190 return nn
191 except (IndexError, TypeError):
192 step = key.step
193 stop = key.stop
194 start = key.start
195 if start is None:
196 start = 0
197 if isinstance(step, complex):
198 step = abs(step)
199 length = int(step)
200 if step != 1:
201 step = (key.stop-start)/float(step-1)
202 stop = key.stop + step
203 return _nx.arange(0, length, 1, float)*step + start
204 else:
205 return _nx.arange(start, stop, step)
208class MGridClass(nd_grid):
209 """
210 `nd_grid` instance which returns a dense multi-dimensional "meshgrid".
212 An instance of `numpy.lib.index_tricks.nd_grid` which returns an dense
213 (or fleshed out) mesh-grid when indexed, so that each returned argument
214 has the same shape. The dimensions and number of the output arrays are
215 equal to the number of indexing dimensions. If the step length is not a
216 complex number, then the stop is not inclusive.
218 However, if the step length is a **complex number** (e.g. 5j), then
219 the integer part of its magnitude is interpreted as specifying the
220 number of points to create between the start and stop values, where
221 the stop value **is inclusive**.
223 Returns
224 ----------
225 mesh-grid `ndarrays` all of the same dimensions
227 See Also
228 --------
229 numpy.lib.index_tricks.nd_grid : class of `ogrid` and `mgrid` objects
230 ogrid : like mgrid but returns open (not fleshed out) mesh grids
231 r_ : array concatenator
233 Examples
234 --------
235 >>> np.mgrid[0:5,0:5]
236 array([[[0, 0, 0, 0, 0],
237 [1, 1, 1, 1, 1],
238 [2, 2, 2, 2, 2],
239 [3, 3, 3, 3, 3],
240 [4, 4, 4, 4, 4]],
241 [[0, 1, 2, 3, 4],
242 [0, 1, 2, 3, 4],
243 [0, 1, 2, 3, 4],
244 [0, 1, 2, 3, 4],
245 [0, 1, 2, 3, 4]]])
246 >>> np.mgrid[-1:1:5j]
247 array([-1. , -0.5, 0. , 0.5, 1. ])
249 """
250 def __init__(self):
251 super(MGridClass, self).__init__(sparse=False)
253mgrid = MGridClass()
255class OGridClass(nd_grid):
256 """
257 `nd_grid` instance which returns an open multi-dimensional "meshgrid".
259 An instance of `numpy.lib.index_tricks.nd_grid` which returns an open
260 (i.e. not fleshed out) mesh-grid when indexed, so that only one dimension
261 of each returned array is greater than 1. The dimension and number of the
262 output arrays are equal to the number of indexing dimensions. If the step
263 length is not a complex number, then the stop is not inclusive.
265 However, if the step length is a **complex number** (e.g. 5j), then
266 the integer part of its magnitude is interpreted as specifying the
267 number of points to create between the start and stop values, where
268 the stop value **is inclusive**.
270 Returns
271 -------
272 mesh-grid
273 `ndarrays` with only one dimension not equal to 1
275 See Also
276 --------
277 np.lib.index_tricks.nd_grid : class of `ogrid` and `mgrid` objects
278 mgrid : like `ogrid` but returns dense (or fleshed out) mesh grids
279 r_ : array concatenator
281 Examples
282 --------
283 >>> from numpy import ogrid
284 >>> ogrid[-1:1:5j]
285 array([-1. , -0.5, 0. , 0.5, 1. ])
286 >>> ogrid[0:5,0:5]
287 [array([[0],
288 [1],
289 [2],
290 [3],
291 [4]]), array([[0, 1, 2, 3, 4]])]
293 """
294 def __init__(self):
295 super(OGridClass, self).__init__(sparse=True)
297ogrid = OGridClass()
300class AxisConcatenator:
301 """
302 Translates slice objects to concatenation along an axis.
304 For detailed documentation on usage, see `r_`.
305 """
306 # allow ma.mr_ to override this
307 concatenate = staticmethod(_nx.concatenate)
308 makemat = staticmethod(matrixlib.matrix)
310 def __init__(self, axis=0, matrix=False, ndmin=1, trans1d=-1):
311 self.axis = axis
312 self.matrix = matrix
313 self.trans1d = trans1d
314 self.ndmin = ndmin
316 def __getitem__(self, key):
317 # handle matrix builder syntax
318 if isinstance(key, str):
319 frame = sys._getframe().f_back
320 mymat = matrixlib.bmat(key, frame.f_globals, frame.f_locals)
321 return mymat
323 if not isinstance(key, tuple):
324 key = (key,)
326 # copy attributes, since they can be overridden in the first argument
327 trans1d = self.trans1d
328 ndmin = self.ndmin
329 matrix = self.matrix
330 axis = self.axis
332 objs = []
333 scalars = []
334 arraytypes = []
335 scalartypes = []
337 for k, item in enumerate(key):
338 scalar = False
339 if isinstance(item, slice):
340 step = item.step
341 start = item.start
342 stop = item.stop
343 if start is None:
344 start = 0
345 if step is None:
346 step = 1
347 if isinstance(step, complex):
348 size = int(abs(step))
349 newobj = linspace(start, stop, num=size)
350 else:
351 newobj = _nx.arange(start, stop, step)
352 if ndmin > 1:
353 newobj = array(newobj, copy=False, ndmin=ndmin)
354 if trans1d != -1:
355 newobj = newobj.swapaxes(-1, trans1d)
356 elif isinstance(item, str):
357 if k != 0:
358 raise ValueError("special directives must be the "
359 "first entry.")
360 if item in ('r', 'c'):
361 matrix = True
362 col = (item == 'c')
363 continue
364 if ',' in item:
365 vec = item.split(',')
366 try:
367 axis, ndmin = [int(x) for x in vec[:2]]
368 if len(vec) == 3:
369 trans1d = int(vec[2])
370 continue
371 except Exception as e:
372 raise ValueError(
373 "unknown special directive {!r}".format(item)
374 ) from e
375 try:
376 axis = int(item)
377 continue
378 except (ValueError, TypeError):
379 raise ValueError("unknown special directive")
380 elif type(item) in ScalarType:
381 newobj = array(item, ndmin=ndmin)
382 scalars.append(len(objs))
383 scalar = True
384 scalartypes.append(newobj.dtype)
385 else:
386 item_ndim = ndim(item)
387 newobj = array(item, copy=False, subok=True, ndmin=ndmin)
388 if trans1d != -1 and item_ndim < ndmin:
389 k2 = ndmin - item_ndim
390 k1 = trans1d
391 if k1 < 0:
392 k1 += k2 + 1
393 defaxes = list(range(ndmin))
394 axes = defaxes[:k1] + defaxes[k2:] + defaxes[k1:k2]
395 newobj = newobj.transpose(axes)
396 objs.append(newobj)
397 if not scalar and isinstance(newobj, _nx.ndarray):
398 arraytypes.append(newobj.dtype)
400 # Ensure that scalars won't up-cast unless warranted
401 final_dtype = find_common_type(arraytypes, scalartypes)
402 if final_dtype is not None:
403 for k in scalars:
404 objs[k] = objs[k].astype(final_dtype)
406 res = self.concatenate(tuple(objs), axis=axis)
408 if matrix:
409 oldndim = res.ndim
410 res = self.makemat(res)
411 if oldndim == 1 and col:
412 res = res.T
413 return res
415 def __len__(self):
416 return 0
418# separate classes are used here instead of just making r_ = concatentor(0),
419# etc. because otherwise we couldn't get the doc string to come out right
420# in help(r_)
422class RClass(AxisConcatenator):
423 """
424 Translates slice objects to concatenation along the first axis.
426 This is a simple way to build up arrays quickly. There are two use cases.
428 1. If the index expression contains comma separated arrays, then stack
429 them along their first axis.
430 2. If the index expression contains slice notation or scalars then create
431 a 1-D array with a range indicated by the slice notation.
433 If slice notation is used, the syntax ``start:stop:step`` is equivalent
434 to ``np.arange(start, stop, step)`` inside of the brackets. However, if
435 ``step`` is an imaginary number (i.e. 100j) then its integer portion is
436 interpreted as a number-of-points desired and the start and stop are
437 inclusive. In other words ``start:stop:stepj`` is interpreted as
438 ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.
439 After expansion of slice notation, all comma separated sequences are
440 concatenated together.
442 Optional character strings placed as the first element of the index
443 expression can be used to change the output. The strings 'r' or 'c' result
444 in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)
445 matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1
446 (column) matrix is produced. If the result is 2-D then both provide the
447 same matrix result.
449 A string integer specifies which axis to stack multiple comma separated
450 arrays along. A string of two comma-separated integers allows indication
451 of the minimum number of dimensions to force each entry into as the
452 second integer (the axis to concatenate along is still the first integer).
454 A string with three comma-separated integers allows specification of the
455 axis to concatenate along, the minimum number of dimensions to force the
456 entries to, and which axis should contain the start of the arrays which
457 are less than the specified number of dimensions. In other words the third
458 integer allows you to specify where the 1's should be placed in the shape
459 of the arrays that have their shapes upgraded. By default, they are placed
460 in the front of the shape tuple. The third argument allows you to specify
461 where the start of the array should be instead. Thus, a third argument of
462 '0' would place the 1's at the end of the array shape. Negative integers
463 specify where in the new shape tuple the last dimension of upgraded arrays
464 should be placed, so the default is '-1'.
466 Parameters
467 ----------
468 Not a function, so takes no parameters
471 Returns
472 -------
473 A concatenated ndarray or matrix.
475 See Also
476 --------
477 concatenate : Join a sequence of arrays along an existing axis.
478 c_ : Translates slice objects to concatenation along the second axis.
480 Examples
481 --------
482 >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]
483 array([1, 2, 3, ..., 4, 5, 6])
484 >>> np.r_[-1:1:6j, [0]*3, 5, 6]
485 array([-1. , -0.6, -0.2, 0.2, 0.6, 1. , 0. , 0. , 0. , 5. , 6. ])
487 String integers specify the axis to concatenate along or the minimum
488 number of dimensions to force entries into.
490 >>> a = np.array([[0, 1, 2], [3, 4, 5]])
491 >>> np.r_['-1', a, a] # concatenate along last axis
492 array([[0, 1, 2, 0, 1, 2],
493 [3, 4, 5, 3, 4, 5]])
494 >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2
495 array([[1, 2, 3],
496 [4, 5, 6]])
498 >>> np.r_['0,2,0', [1,2,3], [4,5,6]]
499 array([[1],
500 [2],
501 [3],
502 [4],
503 [5],
504 [6]])
505 >>> np.r_['1,2,0', [1,2,3], [4,5,6]]
506 array([[1, 4],
507 [2, 5],
508 [3, 6]])
510 Using 'r' or 'c' as a first string argument creates a matrix.
512 >>> np.r_['r',[1,2,3], [4,5,6]]
513 matrix([[1, 2, 3, 4, 5, 6]])
515 """
517 def __init__(self):
518 AxisConcatenator.__init__(self, 0)
520r_ = RClass()
522class CClass(AxisConcatenator):
523 """
524 Translates slice objects to concatenation along the second axis.
526 This is short-hand for ``np.r_['-1,2,0', index expression]``, which is
527 useful because of its common occurrence. In particular, arrays will be
528 stacked along their last axis after being upgraded to at least 2-D with
529 1's post-pended to the shape (column vectors made out of 1-D arrays).
531 See Also
532 --------
533 column_stack : Stack 1-D arrays as columns into a 2-D array.
534 r_ : For more detailed documentation.
536 Examples
537 --------
538 >>> np.c_[np.array([1,2,3]), np.array([4,5,6])]
539 array([[1, 4],
540 [2, 5],
541 [3, 6]])
542 >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]
543 array([[1, 2, 3, ..., 4, 5, 6]])
545 """
547 def __init__(self):
548 AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
551c_ = CClass()
554@set_module('numpy')
555class ndenumerate:
556 """
557 Multidimensional index iterator.
559 Return an iterator yielding pairs of array coordinates and values.
561 Parameters
562 ----------
563 arr : ndarray
564 Input array.
566 See Also
567 --------
568 ndindex, flatiter
570 Examples
571 --------
572 >>> a = np.array([[1, 2], [3, 4]])
573 >>> for index, x in np.ndenumerate(a):
574 ... print(index, x)
575 (0, 0) 1
576 (0, 1) 2
577 (1, 0) 3
578 (1, 1) 4
580 """
582 def __init__(self, arr):
583 self.iter = asarray(arr).flat
585 def __next__(self):
586 """
587 Standard iterator method, returns the index tuple and array value.
589 Returns
590 -------
591 coords : tuple of ints
592 The indices of the current iteration.
593 val : scalar
594 The array element of the current iteration.
596 """
597 return self.iter.coords, next(self.iter)
599 def __iter__(self):
600 return self
603@set_module('numpy')
604class ndindex:
605 """
606 An N-dimensional iterator object to index arrays.
608 Given the shape of an array, an `ndindex` instance iterates over
609 the N-dimensional index of the array. At each iteration a tuple
610 of indices is returned, the last dimension is iterated over first.
612 Parameters
613 ----------
614 `*args` : ints
615 The size of each dimension of the array.
617 See Also
618 --------
619 ndenumerate, flatiter
621 Examples
622 --------
623 >>> for index in np.ndindex(3, 2, 1):
624 ... print(index)
625 (0, 0, 0)
626 (0, 1, 0)
627 (1, 0, 0)
628 (1, 1, 0)
629 (2, 0, 0)
630 (2, 1, 0)
632 """
634 def __init__(self, *shape):
635 if len(shape) == 1 and isinstance(shape[0], tuple):
636 shape = shape[0]
637 x = as_strided(_nx.zeros(1), shape=shape,
638 strides=_nx.zeros_like(shape))
639 self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
640 order='C')
642 def __iter__(self):
643 return self
645 def ndincr(self):
646 """
647 Increment the multi-dimensional index by one.
649 This method is for backward compatibility only: do not use.
650 """
651 next(self)
653 def __next__(self):
654 """
655 Standard iterator method, updates the index and returns the index
656 tuple.
658 Returns
659 -------
660 val : tuple of ints
661 Returns a tuple containing the indices of the current
662 iteration.
664 """
665 next(self._it)
666 return self._it.multi_index
669# You can do all this with slice() plus a few special objects,
670# but there's a lot to remember. This version is simpler because
671# it uses the standard array indexing syntax.
672#
673# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
674# last revision: 1999-7-23
675#
676# Cosmetic changes by T. Oliphant 2001
677#
678#
680class IndexExpression:
681 """
682 A nicer way to build up index tuples for arrays.
684 .. note::
685 Use one of the two predefined instances `index_exp` or `s_`
686 rather than directly using `IndexExpression`.
688 For any index combination, including slicing and axis insertion,
689 ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any
690 array `a`. However, ``np.index_exp[indices]`` can be used anywhere
691 in Python code and returns a tuple of slice objects that can be
692 used in the construction of complex index expressions.
694 Parameters
695 ----------
696 maketuple : bool
697 If True, always returns a tuple.
699 See Also
700 --------
701 index_exp : Predefined instance that always returns a tuple:
702 `index_exp = IndexExpression(maketuple=True)`.
703 s_ : Predefined instance without tuple conversion:
704 `s_ = IndexExpression(maketuple=False)`.
706 Notes
707 -----
708 You can do all this with `slice()` plus a few special objects,
709 but there's a lot to remember and this version is simpler because
710 it uses the standard array indexing syntax.
712 Examples
713 --------
714 >>> np.s_[2::2]
715 slice(2, None, 2)
716 >>> np.index_exp[2::2]
717 (slice(2, None, 2),)
719 >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]
720 array([2, 4])
722 """
724 def __init__(self, maketuple):
725 self.maketuple = maketuple
727 def __getitem__(self, item):
728 if self.maketuple and not isinstance(item, tuple):
729 return (item,)
730 else:
731 return item
733index_exp = IndexExpression(maketuple=True)
734s_ = IndexExpression(maketuple=False)
736# End contribution from Konrad.
739# The following functions complement those in twodim_base, but are
740# applicable to N-dimensions.
743def _fill_diagonal_dispatcher(a, val, wrap=None):
744 return (a,)
747@array_function_dispatch(_fill_diagonal_dispatcher)
748def fill_diagonal(a, val, wrap=False):
749 """Fill the main diagonal of the given array of any dimensionality.
751 For an array `a` with ``a.ndim >= 2``, the diagonal is the list of
752 locations with indices ``a[i, ..., i]`` all identical. This function
753 modifies the input array in-place, it does not return a value.
755 Parameters
756 ----------
757 a : array, at least 2-D.
758 Array whose diagonal is to be filled, it gets modified in-place.
760 val : scalar
761 Value to be written on the diagonal, its type must be compatible with
762 that of the array a.
764 wrap : bool
765 For tall matrices in NumPy version up to 1.6.2, the
766 diagonal "wrapped" after N columns. You can have this behavior
767 with this option. This affects only tall matrices.
769 See also
770 --------
771 diag_indices, diag_indices_from
773 Notes
774 -----
775 .. versionadded:: 1.4.0
777 This functionality can be obtained via `diag_indices`, but internally
778 this version uses a much faster implementation that never constructs the
779 indices and uses simple slicing.
781 Examples
782 --------
783 >>> a = np.zeros((3, 3), int)
784 >>> np.fill_diagonal(a, 5)
785 >>> a
786 array([[5, 0, 0],
787 [0, 5, 0],
788 [0, 0, 5]])
790 The same function can operate on a 4-D array:
792 >>> a = np.zeros((3, 3, 3, 3), int)
793 >>> np.fill_diagonal(a, 4)
795 We only show a few blocks for clarity:
797 >>> a[0, 0]
798 array([[4, 0, 0],
799 [0, 0, 0],
800 [0, 0, 0]])
801 >>> a[1, 1]
802 array([[0, 0, 0],
803 [0, 4, 0],
804 [0, 0, 0]])
805 >>> a[2, 2]
806 array([[0, 0, 0],
807 [0, 0, 0],
808 [0, 0, 4]])
810 The wrap option affects only tall matrices:
812 >>> # tall matrices no wrap
813 >>> a = np.zeros((5, 3), int)
814 >>> np.fill_diagonal(a, 4)
815 >>> a
816 array([[4, 0, 0],
817 [0, 4, 0],
818 [0, 0, 4],
819 [0, 0, 0],
820 [0, 0, 0]])
822 >>> # tall matrices wrap
823 >>> a = np.zeros((5, 3), int)
824 >>> np.fill_diagonal(a, 4, wrap=True)
825 >>> a
826 array([[4, 0, 0],
827 [0, 4, 0],
828 [0, 0, 4],
829 [0, 0, 0],
830 [4, 0, 0]])
832 >>> # wide matrices
833 >>> a = np.zeros((3, 5), int)
834 >>> np.fill_diagonal(a, 4, wrap=True)
835 >>> a
836 array([[4, 0, 0, 0, 0],
837 [0, 4, 0, 0, 0],
838 [0, 0, 4, 0, 0]])
840 The anti-diagonal can be filled by reversing the order of elements
841 using either `numpy.flipud` or `numpy.fliplr`.
843 >>> a = np.zeros((3, 3), int);
844 >>> np.fill_diagonal(np.fliplr(a), [1,2,3]) # Horizontal flip
845 >>> a
846 array([[0, 0, 1],
847 [0, 2, 0],
848 [3, 0, 0]])
849 >>> np.fill_diagonal(np.flipud(a), [1,2,3]) # Vertical flip
850 >>> a
851 array([[0, 0, 3],
852 [0, 2, 0],
853 [1, 0, 0]])
855 Note that the order in which the diagonal is filled varies depending
856 on the flip function.
857 """
858 if a.ndim < 2:
859 raise ValueError("array must be at least 2-d")
860 end = None
861 if a.ndim == 2:
862 # Explicit, fast formula for the common case. For 2-d arrays, we
863 # accept rectangular ones.
864 step = a.shape[1] + 1
865 #This is needed to don't have tall matrix have the diagonal wrap.
866 if not wrap:
867 end = a.shape[1] * a.shape[1]
868 else:
869 # For more than d=2, the strided formula is only valid for arrays with
870 # all dimensions equal, so we check first.
871 if not alltrue(diff(a.shape) == 0):
872 raise ValueError("All dimensions of input must be of equal length")
873 step = 1 + (cumprod(a.shape[:-1])).sum()
875 # Write the value out into the diagonal.
876 a.flat[:end:step] = val
879@set_module('numpy')
880def diag_indices(n, ndim=2):
881 """
882 Return the indices to access the main diagonal of an array.
884 This returns a tuple of indices that can be used to access the main
885 diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
886 (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for
887 ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``
888 for ``i = [0..n-1]``.
890 Parameters
891 ----------
892 n : int
893 The size, along each dimension, of the arrays for which the returned
894 indices can be used.
896 ndim : int, optional
897 The number of dimensions.
899 See also
900 --------
901 diag_indices_from
903 Notes
904 -----
905 .. versionadded:: 1.4.0
907 Examples
908 --------
909 Create a set of indices to access the diagonal of a (4, 4) array:
911 >>> di = np.diag_indices(4)
912 >>> di
913 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
914 >>> a = np.arange(16).reshape(4, 4)
915 >>> a
916 array([[ 0, 1, 2, 3],
917 [ 4, 5, 6, 7],
918 [ 8, 9, 10, 11],
919 [12, 13, 14, 15]])
920 >>> a[di] = 100
921 >>> a
922 array([[100, 1, 2, 3],
923 [ 4, 100, 6, 7],
924 [ 8, 9, 100, 11],
925 [ 12, 13, 14, 100]])
927 Now, we create indices to manipulate a 3-D array:
929 >>> d3 = np.diag_indices(2, 3)
930 >>> d3
931 (array([0, 1]), array([0, 1]), array([0, 1]))
933 And use it to set the diagonal of an array of zeros to 1:
935 >>> a = np.zeros((2, 2, 2), dtype=int)
936 >>> a[d3] = 1
937 >>> a
938 array([[[1, 0],
939 [0, 0]],
940 [[0, 0],
941 [0, 1]]])
943 """
944 idx = arange(n)
945 return (idx,) * ndim
948def _diag_indices_from(arr):
949 return (arr,)
952@array_function_dispatch(_diag_indices_from)
953def diag_indices_from(arr):
954 """
955 Return the indices to access the main diagonal of an n-dimensional array.
957 See `diag_indices` for full details.
959 Parameters
960 ----------
961 arr : array, at least 2-D
963 See Also
964 --------
965 diag_indices
967 Notes
968 -----
969 .. versionadded:: 1.4.0
971 """
973 if not arr.ndim >= 2:
974 raise ValueError("input array must be at least 2-d")
975 # For more than d=2, the strided formula is only valid for arrays with
976 # all dimensions equal, so we check first.
977 if not alltrue(diff(arr.shape) == 0):
978 raise ValueError("All dimensions of input must be of equal length")
980 return diag_indices(arr.shape[0], arr.ndim)