Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1import functools
2import math
3import sys
4from itertools import product
6import numpy as np
7import numpy._core.numeric as _nx
8import numpy.matrixlib as matrixlib
9from numpy._core import linspace, overrides
10from numpy._core.multiarray import ravel_multi_index, unravel_index
11from numpy._core.numeric import ScalarType, array
12from numpy._core.numerictypes import issubdtype
13from numpy._utils import set_module
14from numpy.lib._function_base_impl import diff
16array_function_dispatch = functools.partial(
17 overrides.array_function_dispatch, module='numpy')
20__all__ = [
21 'ravel_multi_index', 'unravel_index', 'mgrid', 'ogrid', 'r_', 'c_',
22 's_', 'index_exp', 'ix_', 'ndenumerate', 'ndindex', 'fill_diagonal',
23 'diag_indices', 'diag_indices_from'
24]
27def _ix__dispatcher(*args):
28 return args
31@array_function_dispatch(_ix__dispatcher)
32def ix_(*args):
33 """
34 Construct an open mesh from multiple sequences.
36 This function takes N 1-D sequences and returns N outputs with N
37 dimensions each, such that the shape is 1 in all but one dimension
38 and the dimension with the non-unit shape value cycles through all
39 N dimensions.
41 Using `ix_` one can quickly construct index arrays that will index
42 the cross product. ``a[np.ix_([1,3],[2,5])]`` returns the array
43 ``[[a[1,2] a[1,5]], [a[3,2] a[3,5]]]``.
45 Parameters
46 ----------
47 args : 1-D sequences
48 Each sequence should be of integer or boolean type.
49 Boolean sequences will be interpreted as boolean masks for the
50 corresponding dimension (equivalent to passing in
51 ``np.nonzero(boolean_sequence)``).
53 Returns
54 -------
55 out : tuple of ndarrays
56 N arrays with N dimensions each, with N the number of input
57 sequences. Together these arrays form an open mesh.
59 See Also
60 --------
61 ogrid, mgrid, meshgrid
63 Examples
64 --------
65 >>> import numpy as np
66 >>> a = np.arange(10).reshape(2, 5)
67 >>> a
68 array([[0, 1, 2, 3, 4],
69 [5, 6, 7, 8, 9]])
70 >>> ixgrid = np.ix_([0, 1], [2, 4])
71 >>> ixgrid
72 (array([[0],
73 [1]]), array([[2, 4]]))
74 >>> ixgrid[0].shape, ixgrid[1].shape
75 ((2, 1), (1, 2))
76 >>> a[ixgrid]
77 array([[2, 4],
78 [7, 9]])
80 >>> ixgrid = np.ix_([True, True], [2, 4])
81 >>> a[ixgrid]
82 array([[2, 4],
83 [7, 9]])
84 >>> ixgrid = np.ix_([True, True], [False, False, True, False, True])
85 >>> a[ixgrid]
86 array([[2, 4],
87 [7, 9]])
89 """
90 out = []
91 nd = len(args)
92 for k, new in enumerate(args):
93 if not isinstance(new, _nx.ndarray):
94 new = np.asarray(new)
95 if new.size == 0:
96 # Explicitly type empty arrays to avoid float default
97 new = new.astype(_nx.intp)
98 if new.ndim != 1:
99 raise ValueError("Cross index must be 1 dimensional")
100 if issubdtype(new.dtype, _nx.bool):
101 new, = new.nonzero()
102 new = new.reshape((1,) * k + (new.size,) + (1,) * (nd - k - 1))
103 out.append(new)
104 return tuple(out)
107class nd_grid:
108 """
109 Construct a multi-dimensional "meshgrid".
111 ``grid = nd_grid()`` creates an instance which will return a mesh-grid
112 when indexed. The dimension and number of the output arrays are equal
113 to the number of indexing dimensions. If the step length is not a
114 complex number, then the stop is not inclusive.
116 However, if the step length is a **complex number** (e.g. 5j), then the
117 integer part of its magnitude is interpreted as specifying the
118 number of points to create between the start and stop values, where
119 the stop value **is inclusive**.
121 If instantiated with an argument of ``sparse=True``, the mesh-grid is
122 open (or not fleshed out) so that only one-dimension of each returned
123 argument is greater than 1.
125 Parameters
126 ----------
127 sparse : bool, optional
128 Whether the grid is sparse or not. Default is False.
130 Notes
131 -----
132 Two instances of `nd_grid` are made available in the NumPy namespace,
133 `mgrid` and `ogrid`, approximately defined as::
135 mgrid = nd_grid(sparse=False)
136 ogrid = nd_grid(sparse=True)
138 Users should use these pre-defined instances instead of using `nd_grid`
139 directly.
140 """
141 __slots__ = ('sparse',)
143 def __init__(self, sparse=False):
144 self.sparse = sparse
146 def __getitem__(self, key):
147 try:
148 size = []
149 # Mimic the behavior of `np.arange` and use a data type
150 # which is at least as large as `np.int_`
151 num_list = [0]
152 for k in range(len(key)):
153 step = key[k].step
154 start = key[k].start
155 stop = key[k].stop
156 if start is None:
157 start = 0
158 if step is None:
159 step = 1
160 if isinstance(step, (_nx.complexfloating, complex)):
161 step = abs(step)
162 size.append(int(step))
163 else:
164 size.append(
165 math.ceil((stop - start) / step))
166 num_list += [start, stop, step]
167 typ = _nx.result_type(*num_list)
168 if self.sparse:
169 nn = [_nx.arange(_x, dtype=_t)
170 for _x, _t in zip(size, (typ,) * len(size))]
171 else:
172 nn = _nx.indices(size, typ)
173 for k, kk in enumerate(key):
174 step = kk.step
175 start = kk.start
176 if start is None:
177 start = 0
178 if step is None:
179 step = 1
180 if isinstance(step, (_nx.complexfloating, complex)):
181 step = int(abs(step))
182 if step != 1:
183 step = (kk.stop - start) / float(step - 1)
184 nn[k] = (nn[k] * step + start)
185 if self.sparse:
186 slobj = [_nx.newaxis] * len(size)
187 for k in range(len(size)):
188 slobj[k] = slice(None, None)
189 nn[k] = nn[k][tuple(slobj)]
190 slobj[k] = _nx.newaxis
191 return tuple(nn) # ogrid -> tuple of arrays
192 return nn # mgrid -> ndarray
193 except (IndexError, TypeError):
194 step = key.step
195 stop = key.stop
196 start = key.start
197 if start is None:
198 start = 0
199 if isinstance(step, (_nx.complexfloating, complex)):
200 # Prevent the (potential) creation of integer arrays
201 step_float = abs(step)
202 step = length = int(step_float)
203 if step != 1:
204 step = (key.stop - start) / float(step - 1)
205 typ = _nx.result_type(start, stop, step_float)
206 return _nx.arange(0, length, 1, dtype=typ) * step + start
207 else:
208 return _nx.arange(start, stop, step)
211class MGridClass(nd_grid):
212 """
213 An instance which returns a dense multi-dimensional "meshgrid".
215 An instance which returns a dense (or fleshed out) mesh-grid
216 when indexed, so that each returned argument has the same shape.
217 The dimensions and number of the output arrays are equal to the
218 number of indexing dimensions. If the step length is not a complex
219 number, then the stop is not inclusive.
221 However, if the step length is a **complex number** (e.g. 5j), then
222 the integer part of its magnitude is interpreted as specifying the
223 number of points to create between the start and stop values, where
224 the stop value **is inclusive**.
226 Returns
227 -------
228 mesh-grid : ndarray
229 A single array, containing a set of `ndarray`\\ s all of the same
230 dimensions. stacked along the first axis.
232 See Also
233 --------
234 ogrid : like `mgrid` but returns open (not fleshed out) mesh grids
235 meshgrid: return coordinate matrices from coordinate vectors
236 r_ : array concatenator
237 :ref:`how-to-partition`
239 Examples
240 --------
241 >>> import numpy as np
242 >>> np.mgrid[0:5, 0:5]
243 array([[[0, 0, 0, 0, 0],
244 [1, 1, 1, 1, 1],
245 [2, 2, 2, 2, 2],
246 [3, 3, 3, 3, 3],
247 [4, 4, 4, 4, 4]],
248 [[0, 1, 2, 3, 4],
249 [0, 1, 2, 3, 4],
250 [0, 1, 2, 3, 4],
251 [0, 1, 2, 3, 4],
252 [0, 1, 2, 3, 4]]])
253 >>> np.mgrid[-1:1:5j]
254 array([-1. , -0.5, 0. , 0.5, 1. ])
256 >>> np.mgrid[0:4].shape
257 (4,)
258 >>> np.mgrid[0:4, 0:5].shape
259 (2, 4, 5)
260 >>> np.mgrid[0:4, 0:5, 0:6].shape
261 (3, 4, 5, 6)
263 """
264 __slots__ = ()
266 def __init__(self):
267 super().__init__(sparse=False)
270mgrid = MGridClass()
273class OGridClass(nd_grid):
274 """
275 An instance which returns an open multi-dimensional "meshgrid".
277 An instance which returns an open (i.e. not fleshed out) mesh-grid
278 when indexed, so that only one dimension of each returned array is
279 greater than 1. The dimension and number of the output arrays are
280 equal to the number of indexing dimensions. If the step length is
281 not a complex number, then the stop is not inclusive.
283 However, if the step length is a **complex number** (e.g. 5j), then
284 the integer part of its magnitude is interpreted as specifying the
285 number of points to create between the start and stop values, where
286 the stop value **is inclusive**.
288 Returns
289 -------
290 mesh-grid : ndarray or tuple of ndarrays
291 If the input is a single slice, returns an array.
292 If the input is multiple slices, returns a tuple of arrays, with
293 only one dimension not equal to 1.
295 See Also
296 --------
297 mgrid : like `ogrid` but returns dense (or fleshed out) mesh grids
298 meshgrid: return coordinate matrices from coordinate vectors
299 r_ : array concatenator
300 :ref:`how-to-partition`
302 Examples
303 --------
304 >>> from numpy import ogrid
305 >>> ogrid[-1:1:5j]
306 array([-1. , -0.5, 0. , 0.5, 1. ])
307 >>> ogrid[0:5, 0:5]
308 (array([[0],
309 [1],
310 [2],
311 [3],
312 [4]]),
313 array([[0, 1, 2, 3, 4]]))
315 """
316 __slots__ = ()
318 def __init__(self):
319 super().__init__(sparse=True)
322ogrid = OGridClass()
325class AxisConcatenator:
326 """
327 Translates slice objects to concatenation along an axis.
329 For detailed documentation on usage, see `r_`.
330 """
331 __slots__ = ('axis', 'matrix', 'ndmin', 'trans1d')
333 # allow ma.mr_ to override this
334 concatenate = staticmethod(_nx.concatenate)
335 makemat = staticmethod(matrixlib.matrix)
337 def __init__(self, axis=0, matrix=False, ndmin=1, trans1d=-1):
338 self.axis = axis
339 self.matrix = matrix
340 self.trans1d = trans1d
341 self.ndmin = ndmin
343 def __getitem__(self, key):
344 # handle matrix builder syntax
345 if isinstance(key, str):
346 frame = sys._getframe().f_back
347 mymat = matrixlib.bmat(key, frame.f_globals, frame.f_locals)
348 return mymat
350 if not isinstance(key, tuple):
351 key = (key,)
353 # copy attributes, since they can be overridden in the first argument
354 trans1d = self.trans1d
355 ndmin = self.ndmin
356 matrix = self.matrix
357 axis = self.axis
359 objs = []
360 # dtypes or scalars for weak scalar handling in result_type
361 result_type_objs = []
363 for k, item in enumerate(key):
364 scalar = False
365 if isinstance(item, slice):
366 step = item.step
367 start = item.start
368 stop = item.stop
369 if start is None:
370 start = 0
371 if step is None:
372 step = 1
373 if isinstance(step, (_nx.complexfloating, complex)):
374 size = int(abs(step))
375 newobj = linspace(start, stop, num=size)
376 else:
377 newobj = _nx.arange(start, stop, step)
378 if ndmin > 1:
379 newobj = array(newobj, copy=None, ndmin=ndmin)
380 if trans1d != -1:
381 newobj = newobj.swapaxes(-1, trans1d)
382 elif isinstance(item, str):
383 if k != 0:
384 raise ValueError("special directives must be the "
385 "first entry.")
386 if item in ('r', 'c'):
387 matrix = True
388 col = (item == 'c')
389 continue
390 if ',' in item:
391 vec = item.split(',')
392 try:
393 axis, ndmin = [int(x) for x in vec[:2]]
394 if len(vec) == 3:
395 trans1d = int(vec[2])
396 continue
397 except Exception as e:
398 raise ValueError(
399 f"unknown special directive {item!r}"
400 ) from e
401 try:
402 axis = int(item)
403 continue
404 except (ValueError, TypeError) as e:
405 raise ValueError("unknown special directive") from e
406 elif type(item) in ScalarType:
407 scalar = True
408 newobj = item
409 else:
410 item_ndim = np.ndim(item)
411 newobj = array(item, copy=None, subok=True, ndmin=ndmin)
412 if trans1d != -1 and item_ndim < ndmin:
413 k2 = ndmin - item_ndim
414 k1 = trans1d
415 if k1 < 0:
416 k1 += k2 + 1
417 defaxes = list(range(ndmin))
418 axes = defaxes[:k1] + defaxes[k2:] + defaxes[k1:k2]
419 newobj = newobj.transpose(axes)
421 objs.append(newobj)
422 if scalar:
423 result_type_objs.append(item)
424 else:
425 result_type_objs.append(newobj.dtype)
427 # Ensure that scalars won't up-cast unless warranted, for 0, drops
428 # through to error in concatenate.
429 if len(result_type_objs) != 0:
430 final_dtype = _nx.result_type(*result_type_objs)
431 # concatenate could do cast, but that can be overridden:
432 objs = [array(obj, copy=None, subok=True,
433 ndmin=ndmin, dtype=final_dtype) for obj in objs]
435 res = self.concatenate(tuple(objs), axis=axis)
437 if matrix:
438 oldndim = res.ndim
439 res = self.makemat(res)
440 if oldndim == 1 and col:
441 res = res.T
442 return res
444 def __len__(self):
445 return 0
447# separate classes are used here instead of just making r_ = concatenator(0),
448# etc. because otherwise we couldn't get the doc string to come out right
449# in help(r_)
452class RClass(AxisConcatenator):
453 """
454 Translates slice objects to concatenation along the first axis.
456 This is a simple way to build up arrays quickly. There are two use cases.
458 1. If the index expression contains comma separated arrays, then stack
459 them along their first axis.
460 2. If the index expression contains slice notation or scalars then create
461 a 1-D array with a range indicated by the slice notation.
463 If slice notation is used, the syntax ``start:stop:step`` is equivalent
464 to ``np.arange(start, stop, step)`` inside of the brackets. However, if
465 ``step`` is an imaginary number (i.e. 100j) then its integer portion is
466 interpreted as a number-of-points desired and the start and stop are
467 inclusive. In other words ``start:stop:stepj`` is interpreted as
468 ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.
469 After expansion of slice notation, all comma separated sequences are
470 concatenated together.
472 Optional character strings placed as the first element of the index
473 expression can be used to change the output. The strings 'r' or 'c' result
474 in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)
475 matrix is produced. If the result is 1-D and 'c' is specified, then
476 an N x 1 (column) matrix is produced.
477 If the result is 2-D then both provide the same matrix result.
479 A string integer specifies which axis to stack multiple comma separated
480 arrays along. A string of two comma-separated integers allows indication
481 of the minimum number of dimensions to force each entry into as the
482 second integer (the axis to concatenate along is still the first integer).
484 A string with three comma-separated integers allows specification of the
485 axis to concatenate along, the minimum number of dimensions to force the
486 entries to, and which axis should contain the start of the arrays which
487 are less than the specified number of dimensions. In other words the third
488 integer allows you to specify where the 1's should be placed in the shape
489 of the arrays that have their shapes upgraded. By default, they are placed
490 in the front of the shape tuple. The third argument allows you to specify
491 where the start of the array should be instead. Thus, a third argument of
492 '0' would place the 1's at the end of the array shape. Negative integers
493 specify where in the new shape tuple the last dimension of upgraded arrays
494 should be placed, so the default is '-1'.
496 Parameters
497 ----------
498 Not a function, so takes no parameters
501 Returns
502 -------
503 A concatenated ndarray or matrix.
505 See Also
506 --------
507 concatenate : Join a sequence of arrays along an existing axis.
508 c_ : Translates slice objects to concatenation along the second axis.
510 Examples
511 --------
512 >>> import numpy as np
513 >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]
514 array([1, 2, 3, ..., 4, 5, 6])
515 >>> np.r_[-1:1:6j, [0]*3, 5, 6]
516 array([-1. , -0.6, -0.2, 0.2, 0.6, 1. , 0. , 0. , 0. , 5. , 6. ])
518 String integers specify the axis to concatenate along or the minimum
519 number of dimensions to force entries into.
521 >>> a = np.array([[0, 1, 2], [3, 4, 5]])
522 >>> np.r_['-1', a, a] # concatenate along last axis
523 array([[0, 1, 2, 0, 1, 2],
524 [3, 4, 5, 3, 4, 5]])
525 >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2
526 array([[1, 2, 3],
527 [4, 5, 6]])
529 >>> np.r_['0,2,0', [1,2,3], [4,5,6]]
530 array([[1],
531 [2],
532 [3],
533 [4],
534 [5],
535 [6]])
536 >>> np.r_['1,2,0', [1,2,3], [4,5,6]]
537 array([[1, 4],
538 [2, 5],
539 [3, 6]])
541 Using 'r' or 'c' as a first string argument creates a matrix.
543 >>> np.r_['r',[1,2,3], [4,5,6]]
544 matrix([[1, 2, 3, 4, 5, 6]])
546 """
547 __slots__ = ()
549 def __init__(self):
550 AxisConcatenator.__init__(self, 0)
553r_ = RClass()
556class CClass(AxisConcatenator):
557 """
558 Translates slice objects to concatenation along the second axis.
560 This is short-hand for ``np.r_['-1,2,0', index expression]``, which is
561 useful because of its common occurrence. In particular, arrays will be
562 stacked along their last axis after being upgraded to at least 2-D with
563 1's post-pended to the shape (column vectors made out of 1-D arrays).
565 See Also
566 --------
567 column_stack : Stack 1-D arrays as columns into a 2-D array.
568 r_ : For more detailed documentation.
570 Examples
571 --------
572 >>> import numpy as np
573 >>> np.c_[np.array([1,2,3]), np.array([4,5,6])]
574 array([[1, 4],
575 [2, 5],
576 [3, 6]])
577 >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]
578 array([[1, 2, 3, ..., 4, 5, 6]])
580 """
581 __slots__ = ()
583 def __init__(self):
584 AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
587c_ = CClass()
590@set_module('numpy')
591class ndenumerate:
592 """
593 Multidimensional index iterator.
595 Return an iterator yielding pairs of array coordinates and values.
597 Parameters
598 ----------
599 arr : ndarray
600 Input array.
602 See Also
603 --------
604 ndindex, flatiter
606 Examples
607 --------
608 >>> import numpy as np
609 >>> a = np.array([[1, 2], [3, 4]])
610 >>> for index, x in np.ndenumerate(a):
611 ... print(index, x)
612 (0, 0) 1
613 (0, 1) 2
614 (1, 0) 3
615 (1, 1) 4
617 """
619 def __init__(self, arr):
620 self.iter = np.asarray(arr).flat
622 def __next__(self):
623 """
624 Standard iterator method, returns the index tuple and array value.
626 Returns
627 -------
628 coords : tuple of ints
629 The indices of the current iteration.
630 val : scalar
631 The array element of the current iteration.
633 """
634 return self.iter.coords, next(self.iter)
636 def __iter__(self):
637 return self
640@set_module('numpy')
641class ndindex:
642 """
643 An N-dimensional iterator object to index arrays.
645 Given the shape of an array, an `ndindex` instance iterates over
646 the N-dimensional index of the array. At each iteration a tuple
647 of indices is returned, the last dimension is iterated over first.
649 Parameters
650 ----------
651 shape : ints, or a single tuple of ints
652 The size of each dimension of the array can be passed as
653 individual parameters or as the elements of a tuple.
655 See Also
656 --------
657 ndenumerate, flatiter
659 Examples
660 --------
661 >>> import numpy as np
663 Dimensions as individual arguments
665 >>> for index in np.ndindex(3, 2, 1):
666 ... print(index)
667 (0, 0, 0)
668 (0, 1, 0)
669 (1, 0, 0)
670 (1, 1, 0)
671 (2, 0, 0)
672 (2, 1, 0)
674 Same dimensions - but in a tuple ``(3, 2, 1)``
676 >>> for index in np.ndindex((3, 2, 1)):
677 ... print(index)
678 (0, 0, 0)
679 (0, 1, 0)
680 (1, 0, 0)
681 (1, 1, 0)
682 (2, 0, 0)
683 (2, 1, 0)
685 """
687 def __init__(self, *shape):
688 if len(shape) == 1 and isinstance(shape[0], tuple):
689 shape = shape[0]
690 if min(shape, default=0) < 0:
691 raise ValueError("negative dimensions are not allowed")
692 self._iter = product(*map(range, shape))
694 def __iter__(self):
695 return self
697 def __next__(self):
698 """
699 Standard iterator method, updates the index and returns the index
700 tuple.
702 Returns
703 -------
704 val : tuple of ints
705 Returns a tuple containing the indices of the current
706 iteration.
708 """
709 return next(self._iter)
712# You can do all this with slice() plus a few special objects,
713# but there's a lot to remember. This version is simpler because
714# it uses the standard array indexing syntax.
715#
716# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
717# last revision: 1999-7-23
718#
719# Cosmetic changes by T. Oliphant 2001
720#
721#
723class IndexExpression:
724 """
725 A nicer way to build up index tuples for arrays.
727 .. note::
728 Use one of the two predefined instances ``index_exp`` or `s_`
729 rather than directly using `IndexExpression`.
731 For any index combination, including slicing and axis insertion,
732 ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any
733 array `a`. However, ``np.index_exp[indices]`` can be used anywhere
734 in Python code and returns a tuple of slice objects that can be
735 used in the construction of complex index expressions.
737 Parameters
738 ----------
739 maketuple : bool
740 If True, always returns a tuple.
742 See Also
743 --------
744 s_ : Predefined instance without tuple conversion:
745 `s_ = IndexExpression(maketuple=False)`.
746 The ``index_exp`` is another predefined instance that
747 always returns a tuple:
748 `index_exp = IndexExpression(maketuple=True)`.
750 Notes
751 -----
752 You can do all this with :class:`slice` plus a few special objects,
753 but there's a lot to remember and this version is simpler because
754 it uses the standard array indexing syntax.
756 Examples
757 --------
758 >>> import numpy as np
759 >>> np.s_[2::2]
760 slice(2, None, 2)
761 >>> np.index_exp[2::2]
762 (slice(2, None, 2),)
764 >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]
765 array([2, 4])
767 """
768 __slots__ = ('maketuple',)
770 def __init__(self, maketuple):
771 self.maketuple = maketuple
773 def __getitem__(self, item):
774 if self.maketuple and not isinstance(item, tuple):
775 return (item,)
776 else:
777 return item
780index_exp = IndexExpression(maketuple=True)
781s_ = IndexExpression(maketuple=False)
783# End contribution from Konrad.
786# The following functions complement those in twodim_base, but are
787# applicable to N-dimensions.
790def _fill_diagonal_dispatcher(a, val, wrap=None):
791 return (a,)
794@array_function_dispatch(_fill_diagonal_dispatcher)
795def fill_diagonal(a, val, wrap=False):
796 """Fill the main diagonal of the given array of any dimensionality.
798 For an array `a` with ``a.ndim >= 2``, the diagonal is the list of
799 values ``a[i, ..., i]`` with indices ``i`` all identical. This function
800 modifies the input array in-place without returning a value.
802 Parameters
803 ----------
804 a : array, at least 2-D.
805 Array whose diagonal is to be filled in-place.
806 val : scalar or array_like
807 Value(s) to write on the diagonal. If `val` is scalar, the value is
808 written along the diagonal. If array-like, the flattened `val` is
809 written along the diagonal, repeating if necessary to fill all
810 diagonal entries.
812 wrap : bool
813 For tall matrices in NumPy version up to 1.6.2, the
814 diagonal "wrapped" after N columns. You can have this behavior
815 with this option. This affects only tall matrices.
817 See also
818 --------
819 diag_indices, diag_indices_from
821 Notes
822 -----
823 This functionality can be obtained via `diag_indices`, but internally
824 this version uses a much faster implementation that never constructs the
825 indices and uses simple slicing.
827 Examples
828 --------
829 >>> import numpy as np
830 >>> a = np.zeros((3, 3), int)
831 >>> np.fill_diagonal(a, 5)
832 >>> a
833 array([[5, 0, 0],
834 [0, 5, 0],
835 [0, 0, 5]])
837 The same function can operate on a 4-D array:
839 >>> a = np.zeros((3, 3, 3, 3), int)
840 >>> np.fill_diagonal(a, 4)
842 We only show a few blocks for clarity:
844 >>> a[0, 0]
845 array([[4, 0, 0],
846 [0, 0, 0],
847 [0, 0, 0]])
848 >>> a[1, 1]
849 array([[0, 0, 0],
850 [0, 4, 0],
851 [0, 0, 0]])
852 >>> a[2, 2]
853 array([[0, 0, 0],
854 [0, 0, 0],
855 [0, 0, 4]])
857 The wrap option affects only tall matrices:
859 >>> # tall matrices no wrap
860 >>> a = np.zeros((5, 3), int)
861 >>> np.fill_diagonal(a, 4)
862 >>> a
863 array([[4, 0, 0],
864 [0, 4, 0],
865 [0, 0, 4],
866 [0, 0, 0],
867 [0, 0, 0]])
869 >>> # tall matrices wrap
870 >>> a = np.zeros((5, 3), int)
871 >>> np.fill_diagonal(a, 4, wrap=True)
872 >>> a
873 array([[4, 0, 0],
874 [0, 4, 0],
875 [0, 0, 4],
876 [0, 0, 0],
877 [4, 0, 0]])
879 >>> # wide matrices
880 >>> a = np.zeros((3, 5), int)
881 >>> np.fill_diagonal(a, 4, wrap=True)
882 >>> a
883 array([[4, 0, 0, 0, 0],
884 [0, 4, 0, 0, 0],
885 [0, 0, 4, 0, 0]])
887 The anti-diagonal can be filled by reversing the order of elements
888 using either `numpy.flipud` or `numpy.fliplr`.
890 >>> a = np.zeros((3, 3), int);
891 >>> np.fill_diagonal(np.fliplr(a), [1,2,3]) # Horizontal flip
892 >>> a
893 array([[0, 0, 1],
894 [0, 2, 0],
895 [3, 0, 0]])
896 >>> np.fill_diagonal(np.flipud(a), [1,2,3]) # Vertical flip
897 >>> a
898 array([[0, 0, 3],
899 [0, 2, 0],
900 [1, 0, 0]])
902 Note that the order in which the diagonal is filled varies depending
903 on the flip function.
904 """
905 if a.ndim < 2:
906 raise ValueError("array must be at least 2-d")
907 end = None
908 if a.ndim == 2:
909 # Explicit, fast formula for the common case. For 2-d arrays, we
910 # accept rectangular ones.
911 step = a.shape[1] + 1
912 # This is needed to don't have tall matrix have the diagonal wrap.
913 if not wrap:
914 end = a.shape[1] * a.shape[1]
915 else:
916 # For more than d=2, the strided formula is only valid for arrays with
917 # all dimensions equal, so we check first.
918 if not np.all(diff(a.shape) == 0):
919 raise ValueError("All dimensions of input must be of equal length")
920 step = 1 + (np.cumprod(a.shape[:-1])).sum()
922 # Write the value out into the diagonal.
923 a.flat[:end:step] = val
926@set_module('numpy')
927def diag_indices(n, ndim=2):
928 """
929 Return the indices to access the main diagonal of an array.
931 This returns a tuple of indices that can be used to access the main
932 diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
933 (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for
934 ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``
935 for ``i = [0..n-1]``.
937 Parameters
938 ----------
939 n : int
940 The size, along each dimension, of the arrays for which the returned
941 indices can be used.
943 ndim : int, optional
944 The number of dimensions.
946 See Also
947 --------
948 diag_indices_from
950 Examples
951 --------
952 >>> import numpy as np
954 Create a set of indices to access the diagonal of a (4, 4) array:
956 >>> di = np.diag_indices(4)
957 >>> di
958 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
959 >>> a = np.arange(16).reshape(4, 4)
960 >>> a
961 array([[ 0, 1, 2, 3],
962 [ 4, 5, 6, 7],
963 [ 8, 9, 10, 11],
964 [12, 13, 14, 15]])
965 >>> a[di] = 100
966 >>> a
967 array([[100, 1, 2, 3],
968 [ 4, 100, 6, 7],
969 [ 8, 9, 100, 11],
970 [ 12, 13, 14, 100]])
972 Now, we create indices to manipulate a 3-D array:
974 >>> d3 = np.diag_indices(2, 3)
975 >>> d3
976 (array([0, 1]), array([0, 1]), array([0, 1]))
978 And use it to set the diagonal of an array of zeros to 1:
980 >>> a = np.zeros((2, 2, 2), dtype=int)
981 >>> a[d3] = 1
982 >>> a
983 array([[[1, 0],
984 [0, 0]],
985 [[0, 0],
986 [0, 1]]])
988 """
989 idx = np.arange(n)
990 return (idx,) * ndim
993def _diag_indices_from(arr):
994 return (arr,)
997@array_function_dispatch(_diag_indices_from)
998def diag_indices_from(arr):
999 """
1000 Return the indices to access the main diagonal of an n-dimensional array.
1002 See `diag_indices` for full details.
1004 Parameters
1005 ----------
1006 arr : array, at least 2-D
1008 See Also
1009 --------
1010 diag_indices
1012 Examples
1013 --------
1014 >>> import numpy as np
1016 Create a 4 by 4 array.
1018 >>> a = np.arange(16).reshape(4, 4)
1019 >>> a
1020 array([[ 0, 1, 2, 3],
1021 [ 4, 5, 6, 7],
1022 [ 8, 9, 10, 11],
1023 [12, 13, 14, 15]])
1025 Get the indices of the diagonal elements.
1027 >>> di = np.diag_indices_from(a)
1028 >>> di
1029 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1031 >>> a[di]
1032 array([ 0, 5, 10, 15])
1034 This is simply syntactic sugar for diag_indices.
1036 >>> np.diag_indices(a.shape[0])
1037 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1039 """
1041 if not arr.ndim >= 2:
1042 raise ValueError("input array must be at least 2-d")
1043 # For more than d=2, the strided formula is only valid for arrays with
1044 # all dimensions equal, so we check first.
1045 if not np.all(diff(arr.shape) == 0):
1046 raise ValueError("All dimensions of input must be of equal length")
1048 return diag_indices(arr.shape[0], arr.ndim)