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 sys
3import math
4import warnings
6import numpy as np
7from .._utils import set_module
8import numpy.core.numeric as _nx
9from numpy.core.numeric import ScalarType, array
10from numpy.core.numerictypes import issubdtype
12import numpy.matrixlib as matrixlib
13from .function_base import diff
14from numpy.core.multiarray import ravel_multi_index, unravel_index
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 = np.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)
109class nd_grid:
110 """
111 Construct a multi-dimensional "meshgrid".
113 ``grid = nd_grid()`` creates an instance which will return a mesh-grid
114 when indexed. The dimension and number of the output arrays are equal
115 to the number of indexing dimensions. If the step length is not a
116 complex number, then the stop is not inclusive.
118 However, if the step length is a **complex number** (e.g. 5j), then the
119 integer part of its magnitude is interpreted as specifying the
120 number of points to create between the start and stop values, where
121 the stop value **is inclusive**.
123 If instantiated with an argument of ``sparse=True``, the mesh-grid is
124 open (or not fleshed out) so that only one-dimension of each returned
125 argument is greater than 1.
127 Parameters
128 ----------
129 sparse : bool, optional
130 Whether the grid is sparse or not. Default is False.
132 Notes
133 -----
134 Two instances of `nd_grid` are made available in the NumPy namespace,
135 `mgrid` and `ogrid`, approximately defined as::
137 mgrid = nd_grid(sparse=False)
138 ogrid = nd_grid(sparse=True)
140 Users should use these pre-defined instances instead of using `nd_grid`
141 directly.
142 """
144 def __init__(self, sparse=False):
145 self.sparse = sparse
147 def __getitem__(self, key):
148 try:
149 size = []
150 # Mimic the behavior of `np.arange` and use a data type
151 # which is at least as large as `np.int_`
152 num_list = [0]
153 for k in range(len(key)):
154 step = key[k].step
155 start = key[k].start
156 stop = key[k].stop
157 if start is None:
158 start = 0
159 if step is None:
160 step = 1
161 if isinstance(step, (_nx.complexfloating, complex)):
162 step = abs(step)
163 size.append(int(step))
164 else:
165 size.append(
166 int(math.ceil((stop - start) / (step*1.0))))
167 num_list += [start, stop, step]
168 typ = _nx.result_type(*num_list)
169 if self.sparse:
170 nn = [_nx.arange(_x, dtype=_t)
171 for _x, _t in zip(size, (typ,)*len(size))]
172 else:
173 nn = _nx.indices(size, typ)
174 for k, kk in enumerate(key):
175 step = kk.step
176 start = kk.start
177 if start is None:
178 start = 0
179 if step is None:
180 step = 1
181 if isinstance(step, (_nx.complexfloating, complex)):
182 step = int(abs(step))
183 if step != 1:
184 step = (kk.stop - start) / float(step - 1)
185 nn[k] = (nn[k]*step+start)
186 if self.sparse:
187 slobj = [_nx.newaxis]*len(size)
188 for k in range(len(size)):
189 slobj[k] = slice(None, None)
190 nn[k] = nn[k][tuple(slobj)]
191 slobj[k] = _nx.newaxis
192 return nn
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 `ndarrays` all of the same dimensions
230 See Also
231 --------
232 ogrid : like `mgrid` but returns open (not fleshed out) mesh grids
233 meshgrid: return coordinate matrices from coordinate vectors
234 r_ : array concatenator
235 :ref:`how-to-partition`
237 Examples
238 --------
239 >>> np.mgrid[0:5, 0:5]
240 array([[[0, 0, 0, 0, 0],
241 [1, 1, 1, 1, 1],
242 [2, 2, 2, 2, 2],
243 [3, 3, 3, 3, 3],
244 [4, 4, 4, 4, 4]],
245 [[0, 1, 2, 3, 4],
246 [0, 1, 2, 3, 4],
247 [0, 1, 2, 3, 4],
248 [0, 1, 2, 3, 4],
249 [0, 1, 2, 3, 4]]])
250 >>> np.mgrid[-1:1:5j]
251 array([-1. , -0.5, 0. , 0.5, 1. ])
253 """
255 def __init__(self):
256 super().__init__(sparse=False)
259mgrid = MGridClass()
262class OGridClass(nd_grid):
263 """
264 An instance which returns an open multi-dimensional "meshgrid".
266 An instance which returns an open (i.e. not fleshed out) mesh-grid
267 when indexed, so that only one dimension of each returned array is
268 greater than 1. The dimension and number of the output arrays are
269 equal to the number of indexing dimensions. If the step length is
270 not a complex number, then the stop is not inclusive.
272 However, if the step length is a **complex number** (e.g. 5j), then
273 the integer part of its magnitude is interpreted as specifying the
274 number of points to create between the start and stop values, where
275 the stop value **is inclusive**.
277 Returns
278 -------
279 mesh-grid
280 `ndarrays` with only one dimension not equal to 1
282 See Also
283 --------
284 mgrid : like `ogrid` but returns dense (or fleshed out) mesh grids
285 meshgrid: return coordinate matrices from coordinate vectors
286 r_ : array concatenator
287 :ref:`how-to-partition`
289 Examples
290 --------
291 >>> from numpy import ogrid
292 >>> ogrid[-1:1:5j]
293 array([-1. , -0.5, 0. , 0.5, 1. ])
294 >>> ogrid[0:5,0:5]
295 [array([[0],
296 [1],
297 [2],
298 [3],
299 [4]]), array([[0, 1, 2, 3, 4]])]
301 """
303 def __init__(self):
304 super().__init__(sparse=True)
307ogrid = OGridClass()
310class AxisConcatenator:
311 """
312 Translates slice objects to concatenation along an axis.
314 For detailed documentation on usage, see `r_`.
315 """
316 # allow ma.mr_ to override this
317 concatenate = staticmethod(_nx.concatenate)
318 makemat = staticmethod(matrixlib.matrix)
320 def __init__(self, axis=0, matrix=False, ndmin=1, trans1d=-1):
321 self.axis = axis
322 self.matrix = matrix
323 self.trans1d = trans1d
324 self.ndmin = ndmin
326 def __getitem__(self, key):
327 # handle matrix builder syntax
328 if isinstance(key, str):
329 frame = sys._getframe().f_back
330 mymat = matrixlib.bmat(key, frame.f_globals, frame.f_locals)
331 return mymat
333 if not isinstance(key, tuple):
334 key = (key,)
336 # copy attributes, since they can be overridden in the first argument
337 trans1d = self.trans1d
338 ndmin = self.ndmin
339 matrix = self.matrix
340 axis = self.axis
342 objs = []
343 # dtypes or scalars for weak scalar handling in result_type
344 result_type_objs = []
346 for k, item in enumerate(key):
347 scalar = False
348 if isinstance(item, slice):
349 step = item.step
350 start = item.start
351 stop = item.stop
352 if start is None:
353 start = 0
354 if step is None:
355 step = 1
356 if isinstance(step, (_nx.complexfloating, complex)):
357 size = int(abs(step))
358 newobj = linspace(start, stop, num=size)
359 else:
360 newobj = _nx.arange(start, stop, step)
361 if ndmin > 1:
362 newobj = array(newobj, copy=False, ndmin=ndmin)
363 if trans1d != -1:
364 newobj = newobj.swapaxes(-1, trans1d)
365 elif isinstance(item, str):
366 if k != 0:
367 raise ValueError("special directives must be the "
368 "first entry.")
369 if item in ('r', 'c'):
370 matrix = True
371 col = (item == 'c')
372 continue
373 if ',' in item:
374 vec = item.split(',')
375 try:
376 axis, ndmin = [int(x) for x in vec[:2]]
377 if len(vec) == 3:
378 trans1d = int(vec[2])
379 continue
380 except Exception as e:
381 raise ValueError(
382 "unknown special directive {!r}".format(item)
383 ) from e
384 try:
385 axis = int(item)
386 continue
387 except (ValueError, TypeError) as e:
388 raise ValueError("unknown special directive") from e
389 elif type(item) in ScalarType:
390 scalar = True
391 newobj = item
392 else:
393 item_ndim = np.ndim(item)
394 newobj = array(item, copy=False, subok=True, ndmin=ndmin)
395 if trans1d != -1 and item_ndim < ndmin:
396 k2 = ndmin - item_ndim
397 k1 = trans1d
398 if k1 < 0:
399 k1 += k2 + 1
400 defaxes = list(range(ndmin))
401 axes = defaxes[:k1] + defaxes[k2:] + defaxes[k1:k2]
402 newobj = newobj.transpose(axes)
404 objs.append(newobj)
405 if scalar:
406 result_type_objs.append(item)
407 else:
408 result_type_objs.append(newobj.dtype)
410 # Ensure that scalars won't up-cast unless warranted, for 0, drops
411 # through to error in concatenate.
412 if len(result_type_objs) != 0:
413 final_dtype = _nx.result_type(*result_type_objs)
414 # concatenate could do cast, but that can be overriden:
415 objs = [array(obj, copy=False, subok=True,
416 ndmin=ndmin, dtype=final_dtype) for obj in objs]
418 res = self.concatenate(tuple(objs), axis=axis)
420 if matrix:
421 oldndim = res.ndim
422 res = self.makemat(res)
423 if oldndim == 1 and col:
424 res = res.T
425 return res
427 def __len__(self):
428 return 0
430# separate classes are used here instead of just making r_ = concatentor(0),
431# etc. because otherwise we couldn't get the doc string to come out right
432# in help(r_)
435class RClass(AxisConcatenator):
436 """
437 Translates slice objects to concatenation along the first axis.
439 This is a simple way to build up arrays quickly. There are two use cases.
441 1. If the index expression contains comma separated arrays, then stack
442 them along their first axis.
443 2. If the index expression contains slice notation or scalars then create
444 a 1-D array with a range indicated by the slice notation.
446 If slice notation is used, the syntax ``start:stop:step`` is equivalent
447 to ``np.arange(start, stop, step)`` inside of the brackets. However, if
448 ``step`` is an imaginary number (i.e. 100j) then its integer portion is
449 interpreted as a number-of-points desired and the start and stop are
450 inclusive. In other words ``start:stop:stepj`` is interpreted as
451 ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.
452 After expansion of slice notation, all comma separated sequences are
453 concatenated together.
455 Optional character strings placed as the first element of the index
456 expression can be used to change the output. The strings 'r' or 'c' result
457 in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)
458 matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1
459 (column) matrix is produced. If the result is 2-D then both provide the
460 same matrix result.
462 A string integer specifies which axis to stack multiple comma separated
463 arrays along. A string of two comma-separated integers allows indication
464 of the minimum number of dimensions to force each entry into as the
465 second integer (the axis to concatenate along is still the first integer).
467 A string with three comma-separated integers allows specification of the
468 axis to concatenate along, the minimum number of dimensions to force the
469 entries to, and which axis should contain the start of the arrays which
470 are less than the specified number of dimensions. In other words the third
471 integer allows you to specify where the 1's should be placed in the shape
472 of the arrays that have their shapes upgraded. By default, they are placed
473 in the front of the shape tuple. The third argument allows you to specify
474 where the start of the array should be instead. Thus, a third argument of
475 '0' would place the 1's at the end of the array shape. Negative integers
476 specify where in the new shape tuple the last dimension of upgraded arrays
477 should be placed, so the default is '-1'.
479 Parameters
480 ----------
481 Not a function, so takes no parameters
484 Returns
485 -------
486 A concatenated ndarray or matrix.
488 See Also
489 --------
490 concatenate : Join a sequence of arrays along an existing axis.
491 c_ : Translates slice objects to concatenation along the second axis.
493 Examples
494 --------
495 >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]
496 array([1, 2, 3, ..., 4, 5, 6])
497 >>> np.r_[-1:1:6j, [0]*3, 5, 6]
498 array([-1. , -0.6, -0.2, 0.2, 0.6, 1. , 0. , 0. , 0. , 5. , 6. ])
500 String integers specify the axis to concatenate along or the minimum
501 number of dimensions to force entries into.
503 >>> a = np.array([[0, 1, 2], [3, 4, 5]])
504 >>> np.r_['-1', a, a] # concatenate along last axis
505 array([[0, 1, 2, 0, 1, 2],
506 [3, 4, 5, 3, 4, 5]])
507 >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2
508 array([[1, 2, 3],
509 [4, 5, 6]])
511 >>> np.r_['0,2,0', [1,2,3], [4,5,6]]
512 array([[1],
513 [2],
514 [3],
515 [4],
516 [5],
517 [6]])
518 >>> np.r_['1,2,0', [1,2,3], [4,5,6]]
519 array([[1, 4],
520 [2, 5],
521 [3, 6]])
523 Using 'r' or 'c' as a first string argument creates a matrix.
525 >>> np.r_['r',[1,2,3], [4,5,6]]
526 matrix([[1, 2, 3, 4, 5, 6]])
528 """
530 def __init__(self):
531 AxisConcatenator.__init__(self, 0)
534r_ = RClass()
537class CClass(AxisConcatenator):
538 """
539 Translates slice objects to concatenation along the second axis.
541 This is short-hand for ``np.r_['-1,2,0', index expression]``, which is
542 useful because of its common occurrence. In particular, arrays will be
543 stacked along their last axis after being upgraded to at least 2-D with
544 1's post-pended to the shape (column vectors made out of 1-D arrays).
546 See Also
547 --------
548 column_stack : Stack 1-D arrays as columns into a 2-D array.
549 r_ : For more detailed documentation.
551 Examples
552 --------
553 >>> np.c_[np.array([1,2,3]), np.array([4,5,6])]
554 array([[1, 4],
555 [2, 5],
556 [3, 6]])
557 >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]
558 array([[1, 2, 3, ..., 4, 5, 6]])
560 """
562 def __init__(self):
563 AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
566c_ = CClass()
569@set_module('numpy')
570class ndenumerate:
571 """
572 Multidimensional index iterator.
574 Return an iterator yielding pairs of array coordinates and values.
576 Parameters
577 ----------
578 arr : ndarray
579 Input array.
581 See Also
582 --------
583 ndindex, flatiter
585 Examples
586 --------
587 >>> a = np.array([[1, 2], [3, 4]])
588 >>> for index, x in np.ndenumerate(a):
589 ... print(index, x)
590 (0, 0) 1
591 (0, 1) 2
592 (1, 0) 3
593 (1, 1) 4
595 """
597 def __init__(self, arr):
598 self.iter = np.asarray(arr).flat
600 def __next__(self):
601 """
602 Standard iterator method, returns the index tuple and array value.
604 Returns
605 -------
606 coords : tuple of ints
607 The indices of the current iteration.
608 val : scalar
609 The array element of the current iteration.
611 """
612 return self.iter.coords, next(self.iter)
614 def __iter__(self):
615 return self
618@set_module('numpy')
619class ndindex:
620 """
621 An N-dimensional iterator object to index arrays.
623 Given the shape of an array, an `ndindex` instance iterates over
624 the N-dimensional index of the array. At each iteration a tuple
625 of indices is returned, the last dimension is iterated over first.
627 Parameters
628 ----------
629 shape : ints, or a single tuple of ints
630 The size of each dimension of the array can be passed as
631 individual parameters or as the elements of a tuple.
633 See Also
634 --------
635 ndenumerate, flatiter
637 Examples
638 --------
639 Dimensions as individual arguments
641 >>> for index in np.ndindex(3, 2, 1):
642 ... print(index)
643 (0, 0, 0)
644 (0, 1, 0)
645 (1, 0, 0)
646 (1, 1, 0)
647 (2, 0, 0)
648 (2, 1, 0)
650 Same dimensions - but in a tuple ``(3, 2, 1)``
652 >>> for index in np.ndindex((3, 2, 1)):
653 ... print(index)
654 (0, 0, 0)
655 (0, 1, 0)
656 (1, 0, 0)
657 (1, 1, 0)
658 (2, 0, 0)
659 (2, 1, 0)
661 """
663 def __init__(self, *shape):
664 if len(shape) == 1 and isinstance(shape[0], tuple):
665 shape = shape[0]
666 x = as_strided(_nx.zeros(1), shape=shape,
667 strides=_nx.zeros_like(shape))
668 self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
669 order='C')
671 def __iter__(self):
672 return self
674 def ndincr(self):
675 """
676 Increment the multi-dimensional index by one.
678 This method is for backward compatibility only: do not use.
680 .. deprecated:: 1.20.0
681 This method has been advised against since numpy 1.8.0, but only
682 started emitting DeprecationWarning as of this version.
683 """
684 # NumPy 1.20.0, 2020-09-08
685 warnings.warn(
686 "`ndindex.ndincr()` is deprecated, use `next(ndindex)` instead",
687 DeprecationWarning, stacklevel=2)
688 next(self)
690 def __next__(self):
691 """
692 Standard iterator method, updates the index and returns the index
693 tuple.
695 Returns
696 -------
697 val : tuple of ints
698 Returns a tuple containing the indices of the current
699 iteration.
701 """
702 next(self._it)
703 return self._it.multi_index
706# You can do all this with slice() plus a few special objects,
707# but there's a lot to remember. This version is simpler because
708# it uses the standard array indexing syntax.
709#
710# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
711# last revision: 1999-7-23
712#
713# Cosmetic changes by T. Oliphant 2001
714#
715#
717class IndexExpression:
718 """
719 A nicer way to build up index tuples for arrays.
721 .. note::
722 Use one of the two predefined instances `index_exp` or `s_`
723 rather than directly using `IndexExpression`.
725 For any index combination, including slicing and axis insertion,
726 ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any
727 array `a`. However, ``np.index_exp[indices]`` can be used anywhere
728 in Python code and returns a tuple of slice objects that can be
729 used in the construction of complex index expressions.
731 Parameters
732 ----------
733 maketuple : bool
734 If True, always returns a tuple.
736 See Also
737 --------
738 index_exp : Predefined instance that always returns a tuple:
739 `index_exp = IndexExpression(maketuple=True)`.
740 s_ : Predefined instance without tuple conversion:
741 `s_ = IndexExpression(maketuple=False)`.
743 Notes
744 -----
745 You can do all this with `slice()` plus a few special objects,
746 but there's a lot to remember and this version is simpler because
747 it uses the standard array indexing syntax.
749 Examples
750 --------
751 >>> np.s_[2::2]
752 slice(2, None, 2)
753 >>> np.index_exp[2::2]
754 (slice(2, None, 2),)
756 >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]
757 array([2, 4])
759 """
761 def __init__(self, maketuple):
762 self.maketuple = maketuple
764 def __getitem__(self, item):
765 if self.maketuple and not isinstance(item, tuple):
766 return (item,)
767 else:
768 return item
771index_exp = IndexExpression(maketuple=True)
772s_ = IndexExpression(maketuple=False)
774# End contribution from Konrad.
777# The following functions complement those in twodim_base, but are
778# applicable to N-dimensions.
781def _fill_diagonal_dispatcher(a, val, wrap=None):
782 return (a,)
785@array_function_dispatch(_fill_diagonal_dispatcher)
786def fill_diagonal(a, val, wrap=False):
787 """Fill the main diagonal of the given array of any dimensionality.
789 For an array `a` with ``a.ndim >= 2``, the diagonal is the list of
790 locations with indices ``a[i, ..., i]`` all identical. This function
791 modifies the input array in-place, it does not return a value.
793 Parameters
794 ----------
795 a : array, at least 2-D.
796 Array whose diagonal is to be filled, it gets modified in-place.
798 val : scalar or array_like
799 Value(s) to write on the diagonal. If `val` is scalar, the value is
800 written along the diagonal. If array-like, the flattened `val` is
801 written along the diagonal, repeating if necessary to fill all
802 diagonal entries.
804 wrap : bool
805 For tall matrices in NumPy version up to 1.6.2, the
806 diagonal "wrapped" after N columns. You can have this behavior
807 with this option. This affects only tall matrices.
809 See also
810 --------
811 diag_indices, diag_indices_from
813 Notes
814 -----
815 .. versionadded:: 1.4.0
817 This functionality can be obtained via `diag_indices`, but internally
818 this version uses a much faster implementation that never constructs the
819 indices and uses simple slicing.
821 Examples
822 --------
823 >>> a = np.zeros((3, 3), int)
824 >>> np.fill_diagonal(a, 5)
825 >>> a
826 array([[5, 0, 0],
827 [0, 5, 0],
828 [0, 0, 5]])
830 The same function can operate on a 4-D array:
832 >>> a = np.zeros((3, 3, 3, 3), int)
833 >>> np.fill_diagonal(a, 4)
835 We only show a few blocks for clarity:
837 >>> a[0, 0]
838 array([[4, 0, 0],
839 [0, 0, 0],
840 [0, 0, 0]])
841 >>> a[1, 1]
842 array([[0, 0, 0],
843 [0, 4, 0],
844 [0, 0, 0]])
845 >>> a[2, 2]
846 array([[0, 0, 0],
847 [0, 0, 0],
848 [0, 0, 4]])
850 The wrap option affects only tall matrices:
852 >>> # tall matrices no wrap
853 >>> a = np.zeros((5, 3), int)
854 >>> np.fill_diagonal(a, 4)
855 >>> a
856 array([[4, 0, 0],
857 [0, 4, 0],
858 [0, 0, 4],
859 [0, 0, 0],
860 [0, 0, 0]])
862 >>> # tall matrices wrap
863 >>> a = np.zeros((5, 3), int)
864 >>> np.fill_diagonal(a, 4, wrap=True)
865 >>> a
866 array([[4, 0, 0],
867 [0, 4, 0],
868 [0, 0, 4],
869 [0, 0, 0],
870 [4, 0, 0]])
872 >>> # wide matrices
873 >>> a = np.zeros((3, 5), int)
874 >>> np.fill_diagonal(a, 4, wrap=True)
875 >>> a
876 array([[4, 0, 0, 0, 0],
877 [0, 4, 0, 0, 0],
878 [0, 0, 4, 0, 0]])
880 The anti-diagonal can be filled by reversing the order of elements
881 using either `numpy.flipud` or `numpy.fliplr`.
883 >>> a = np.zeros((3, 3), int);
884 >>> np.fill_diagonal(np.fliplr(a), [1,2,3]) # Horizontal flip
885 >>> a
886 array([[0, 0, 1],
887 [0, 2, 0],
888 [3, 0, 0]])
889 >>> np.fill_diagonal(np.flipud(a), [1,2,3]) # Vertical flip
890 >>> a
891 array([[0, 0, 3],
892 [0, 2, 0],
893 [1, 0, 0]])
895 Note that the order in which the diagonal is filled varies depending
896 on the flip function.
897 """
898 if a.ndim < 2:
899 raise ValueError("array must be at least 2-d")
900 end = None
901 if a.ndim == 2:
902 # Explicit, fast formula for the common case. For 2-d arrays, we
903 # accept rectangular ones.
904 step = a.shape[1] + 1
905 # This is needed to don't have tall matrix have the diagonal wrap.
906 if not wrap:
907 end = a.shape[1] * a.shape[1]
908 else:
909 # For more than d=2, the strided formula is only valid for arrays with
910 # all dimensions equal, so we check first.
911 if not np.all(diff(a.shape) == 0):
912 raise ValueError("All dimensions of input must be of equal length")
913 step = 1 + (np.cumprod(a.shape[:-1])).sum()
915 # Write the value out into the diagonal.
916 a.flat[:end:step] = val
919@set_module('numpy')
920def diag_indices(n, ndim=2):
921 """
922 Return the indices to access the main diagonal of an array.
924 This returns a tuple of indices that can be used to access the main
925 diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
926 (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for
927 ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``
928 for ``i = [0..n-1]``.
930 Parameters
931 ----------
932 n : int
933 The size, along each dimension, of the arrays for which the returned
934 indices can be used.
936 ndim : int, optional
937 The number of dimensions.
939 See Also
940 --------
941 diag_indices_from
943 Notes
944 -----
945 .. versionadded:: 1.4.0
947 Examples
948 --------
949 Create a set of indices to access the diagonal of a (4, 4) array:
951 >>> di = np.diag_indices(4)
952 >>> di
953 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
954 >>> a = np.arange(16).reshape(4, 4)
955 >>> a
956 array([[ 0, 1, 2, 3],
957 [ 4, 5, 6, 7],
958 [ 8, 9, 10, 11],
959 [12, 13, 14, 15]])
960 >>> a[di] = 100
961 >>> a
962 array([[100, 1, 2, 3],
963 [ 4, 100, 6, 7],
964 [ 8, 9, 100, 11],
965 [ 12, 13, 14, 100]])
967 Now, we create indices to manipulate a 3-D array:
969 >>> d3 = np.diag_indices(2, 3)
970 >>> d3
971 (array([0, 1]), array([0, 1]), array([0, 1]))
973 And use it to set the diagonal of an array of zeros to 1:
975 >>> a = np.zeros((2, 2, 2), dtype=int)
976 >>> a[d3] = 1
977 >>> a
978 array([[[1, 0],
979 [0, 0]],
980 [[0, 0],
981 [0, 1]]])
983 """
984 idx = np.arange(n)
985 return (idx,) * ndim
988def _diag_indices_from(arr):
989 return (arr,)
992@array_function_dispatch(_diag_indices_from)
993def diag_indices_from(arr):
994 """
995 Return the indices to access the main diagonal of an n-dimensional array.
997 See `diag_indices` for full details.
999 Parameters
1000 ----------
1001 arr : array, at least 2-D
1003 See Also
1004 --------
1005 diag_indices
1007 Notes
1008 -----
1009 .. versionadded:: 1.4.0
1011 Examples
1012 --------
1014 Create a 4 by 4 array.
1016 >>> a = np.arange(16).reshape(4, 4)
1017 >>> a
1018 array([[ 0, 1, 2, 3],
1019 [ 4, 5, 6, 7],
1020 [ 8, 9, 10, 11],
1021 [12, 13, 14, 15]])
1023 Get the indices of the diagonal elements.
1025 >>> di = np.diag_indices_from(a)
1026 >>> di
1027 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1029 >>> a[di]
1030 array([ 0, 5, 10, 15])
1032 This is simply syntactic sugar for diag_indices.
1034 >>> np.diag_indices(a.shape[0])
1035 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1037 """
1039 if not arr.ndim >= 2:
1040 raise ValueError("input array must be at least 2-d")
1041 # For more than d=2, the strided formula is only valid for arrays with
1042 # all dimensions equal, so we check first.
1043 if not np.all(diff(arr.shape) == 0):
1044 raise ValueError("All dimensions of input must be of equal length")
1046 return diag_indices(arr.shape[0], arr.ndim)