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 numpy._core.multiarray import ravel_multi_index, unravel_index
14from numpy._core import overrides, linspace
15from numpy.lib.stride_tricks import as_strided
16from numpy.lib._function_base_impl import diff
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 >>> import numpy as np
69 >>> a = np.arange(10).reshape(2, 5)
70 >>> a
71 array([[0, 1, 2, 3, 4],
72 [5, 6, 7, 8, 9]])
73 >>> ixgrid = np.ix_([0, 1], [2, 4])
74 >>> ixgrid
75 (array([[0],
76 [1]]), array([[2, 4]]))
77 >>> ixgrid[0].shape, ixgrid[1].shape
78 ((2, 1), (1, 2))
79 >>> a[ixgrid]
80 array([[2, 4],
81 [7, 9]])
83 >>> ixgrid = np.ix_([True, True], [2, 4])
84 >>> a[ixgrid]
85 array([[2, 4],
86 [7, 9]])
87 >>> ixgrid = np.ix_([True, True], [False, False, True, False, True])
88 >>> a[ixgrid]
89 array([[2, 4],
90 [7, 9]])
92 """
93 out = []
94 nd = len(args)
95 for k, new in enumerate(args):
96 if not isinstance(new, _nx.ndarray):
97 new = np.asarray(new)
98 if new.size == 0:
99 # Explicitly type empty arrays to avoid float default
100 new = new.astype(_nx.intp)
101 if new.ndim != 1:
102 raise ValueError("Cross index must be 1 dimensional")
103 if issubdtype(new.dtype, _nx.bool):
104 new, = new.nonzero()
105 new = new.reshape((1,)*k + (new.size,) + (1,)*(nd-k-1))
106 out.append(new)
107 return tuple(out)
110class nd_grid:
111 """
112 Construct a multi-dimensional "meshgrid".
114 ``grid = nd_grid()`` creates an instance which will return a mesh-grid
115 when indexed. The dimension and number of the output arrays are equal
116 to the number of indexing dimensions. If the step length is not a
117 complex number, then the stop is not inclusive.
119 However, if the step length is a **complex number** (e.g. 5j), then the
120 integer part of its magnitude is interpreted as specifying the
121 number of points to create between the start and stop values, where
122 the stop value **is inclusive**.
124 If instantiated with an argument of ``sparse=True``, the mesh-grid is
125 open (or not fleshed out) so that only one-dimension of each returned
126 argument is greater than 1.
128 Parameters
129 ----------
130 sparse : bool, optional
131 Whether the grid is sparse or not. Default is False.
133 Notes
134 -----
135 Two instances of `nd_grid` are made available in the NumPy namespace,
136 `mgrid` and `ogrid`, approximately defined as::
138 mgrid = nd_grid(sparse=False)
139 ogrid = nd_grid(sparse=True)
141 Users should use these pre-defined instances instead of using `nd_grid`
142 directly.
143 """
144 __slots__ = ('sparse',)
146 def __init__(self, sparse=False):
147 self.sparse = sparse
149 def __getitem__(self, key):
150 try:
151 size = []
152 # Mimic the behavior of `np.arange` and use a data type
153 # which is at least as large as `np.int_`
154 num_list = [0]
155 for k in range(len(key)):
156 step = key[k].step
157 start = key[k].start
158 stop = key[k].stop
159 if start is None:
160 start = 0
161 if step is None:
162 step = 1
163 if isinstance(step, (_nx.complexfloating, complex)):
164 step = abs(step)
165 size.append(int(step))
166 else:
167 size.append(
168 int(math.ceil((stop - start) / (step*1.0))))
169 num_list += [start, stop, step]
170 typ = _nx.result_type(*num_list)
171 if self.sparse:
172 nn = [_nx.arange(_x, dtype=_t)
173 for _x, _t in zip(size, (typ,)*len(size))]
174 else:
175 nn = _nx.indices(size, typ)
176 for k, kk in enumerate(key):
177 step = kk.step
178 start = kk.start
179 if start is None:
180 start = 0
181 if step is None:
182 step = 1
183 if isinstance(step, (_nx.complexfloating, complex)):
184 step = int(abs(step))
185 if step != 1:
186 step = (kk.stop - start) / float(step - 1)
187 nn[k] = (nn[k]*step+start)
188 if self.sparse:
189 slobj = [_nx.newaxis]*len(size)
190 for k in range(len(size)):
191 slobj[k] = slice(None, None)
192 nn[k] = nn[k][tuple(slobj)]
193 slobj[k] = _nx.newaxis
194 return tuple(nn) # ogrid -> tuple of arrays
195 return nn # mgrid -> ndarray
196 except (IndexError, TypeError):
197 step = key.step
198 stop = key.stop
199 start = key.start
200 if start is None:
201 start = 0
202 if isinstance(step, (_nx.complexfloating, complex)):
203 # Prevent the (potential) creation of integer arrays
204 step_float = abs(step)
205 step = length = int(step_float)
206 if step != 1:
207 step = (key.stop-start)/float(step-1)
208 typ = _nx.result_type(start, stop, step_float)
209 return _nx.arange(0, length, 1, dtype=typ)*step + start
210 else:
211 return _nx.arange(start, stop, step)
214class MGridClass(nd_grid):
215 """
216 An instance which returns a dense multi-dimensional "meshgrid".
218 An instance which returns a dense (or fleshed out) mesh-grid
219 when indexed, so that each returned argument has the same shape.
220 The dimensions and number of the output arrays are equal to the
221 number of indexing dimensions. If the step length is not a complex
222 number, then the stop is not inclusive.
224 However, if the step length is a **complex number** (e.g. 5j), then
225 the integer part of its magnitude is interpreted as specifying the
226 number of points to create between the start and stop values, where
227 the stop value **is inclusive**.
229 Returns
230 -------
231 mesh-grid : ndarray
232 A single array, containing a set of `ndarray`\\ s all of the same
233 dimensions. stacked along the first axis.
235 See Also
236 --------
237 ogrid : like `mgrid` but returns open (not fleshed out) mesh grids
238 meshgrid: return coordinate matrices from coordinate vectors
239 r_ : array concatenator
240 :ref:`how-to-partition`
242 Examples
243 --------
244 >>> import numpy as np
245 >>> np.mgrid[0:5, 0:5]
246 array([[[0, 0, 0, 0, 0],
247 [1, 1, 1, 1, 1],
248 [2, 2, 2, 2, 2],
249 [3, 3, 3, 3, 3],
250 [4, 4, 4, 4, 4]],
251 [[0, 1, 2, 3, 4],
252 [0, 1, 2, 3, 4],
253 [0, 1, 2, 3, 4],
254 [0, 1, 2, 3, 4],
255 [0, 1, 2, 3, 4]]])
256 >>> np.mgrid[-1:1:5j]
257 array([-1. , -0.5, 0. , 0.5, 1. ])
259 >>> np.mgrid[0:4].shape
260 (4,)
261 >>> np.mgrid[0:4, 0:5].shape
262 (2, 4, 5)
263 >>> np.mgrid[0:4, 0:5, 0:6].shape
264 (3, 4, 5, 6)
266 """
267 __slots__ = ()
269 def __init__(self):
270 super().__init__(sparse=False)
273mgrid = MGridClass()
276class OGridClass(nd_grid):
277 """
278 An instance which returns an open multi-dimensional "meshgrid".
280 An instance which returns an open (i.e. not fleshed out) mesh-grid
281 when indexed, so that only one dimension of each returned array is
282 greater than 1. The dimension and number of the output arrays are
283 equal to the number of indexing dimensions. If the step length is
284 not a complex number, then the stop is not inclusive.
286 However, if the step length is a **complex number** (e.g. 5j), then
287 the integer part of its magnitude is interpreted as specifying the
288 number of points to create between the start and stop values, where
289 the stop value **is inclusive**.
291 Returns
292 -------
293 mesh-grid : ndarray or tuple of ndarrays
294 If the input is a single slice, returns an array.
295 If the input is multiple slices, returns a tuple of arrays, with
296 only one dimension not equal to 1.
298 See Also
299 --------
300 mgrid : like `ogrid` but returns dense (or fleshed out) mesh grids
301 meshgrid: return coordinate matrices from coordinate vectors
302 r_ : array concatenator
303 :ref:`how-to-partition`
305 Examples
306 --------
307 >>> from numpy import ogrid
308 >>> ogrid[-1:1:5j]
309 array([-1. , -0.5, 0. , 0.5, 1. ])
310 >>> ogrid[0:5, 0:5]
311 (array([[0],
312 [1],
313 [2],
314 [3],
315 [4]]),
316 array([[0, 1, 2, 3, 4]]))
318 """
319 __slots__ = ()
321 def __init__(self):
322 super().__init__(sparse=True)
325ogrid = OGridClass()
328class AxisConcatenator:
329 """
330 Translates slice objects to concatenation along an axis.
332 For detailed documentation on usage, see `r_`.
333 """
334 __slots__ = ('axis', 'matrix', 'trans1d', 'ndmin')
336 # allow ma.mr_ to override this
337 concatenate = staticmethod(_nx.concatenate)
338 makemat = staticmethod(matrixlib.matrix)
340 def __init__(self, axis=0, matrix=False, ndmin=1, trans1d=-1):
341 self.axis = axis
342 self.matrix = matrix
343 self.trans1d = trans1d
344 self.ndmin = ndmin
346 def __getitem__(self, key):
347 # handle matrix builder syntax
348 if isinstance(key, str):
349 frame = sys._getframe().f_back
350 mymat = matrixlib.bmat(key, frame.f_globals, frame.f_locals)
351 return mymat
353 if not isinstance(key, tuple):
354 key = (key,)
356 # copy attributes, since they can be overridden in the first argument
357 trans1d = self.trans1d
358 ndmin = self.ndmin
359 matrix = self.matrix
360 axis = self.axis
362 objs = []
363 # dtypes or scalars for weak scalar handling in result_type
364 result_type_objs = []
366 for k, item in enumerate(key):
367 scalar = False
368 if isinstance(item, slice):
369 step = item.step
370 start = item.start
371 stop = item.stop
372 if start is None:
373 start = 0
374 if step is None:
375 step = 1
376 if isinstance(step, (_nx.complexfloating, complex)):
377 size = int(abs(step))
378 newobj = linspace(start, stop, num=size)
379 else:
380 newobj = _nx.arange(start, stop, step)
381 if ndmin > 1:
382 newobj = array(newobj, copy=None, ndmin=ndmin)
383 if trans1d != -1:
384 newobj = newobj.swapaxes(-1, trans1d)
385 elif isinstance(item, str):
386 if k != 0:
387 raise ValueError("special directives must be the "
388 "first entry.")
389 if item in ('r', 'c'):
390 matrix = True
391 col = (item == 'c')
392 continue
393 if ',' in item:
394 vec = item.split(',')
395 try:
396 axis, ndmin = [int(x) for x in vec[:2]]
397 if len(vec) == 3:
398 trans1d = int(vec[2])
399 continue
400 except Exception as e:
401 raise ValueError(
402 "unknown special directive {!r}".format(item)
403 ) from e
404 try:
405 axis = int(item)
406 continue
407 except (ValueError, TypeError) as e:
408 raise ValueError("unknown special directive") from e
409 elif type(item) in ScalarType:
410 scalar = True
411 newobj = item
412 else:
413 item_ndim = np.ndim(item)
414 newobj = array(item, copy=None, subok=True, ndmin=ndmin)
415 if trans1d != -1 and item_ndim < ndmin:
416 k2 = ndmin - item_ndim
417 k1 = trans1d
418 if k1 < 0:
419 k1 += k2 + 1
420 defaxes = list(range(ndmin))
421 axes = defaxes[:k1] + defaxes[k2:] + defaxes[k1:k2]
422 newobj = newobj.transpose(axes)
424 objs.append(newobj)
425 if scalar:
426 result_type_objs.append(item)
427 else:
428 result_type_objs.append(newobj.dtype)
430 # Ensure that scalars won't up-cast unless warranted, for 0, drops
431 # through to error in concatenate.
432 if len(result_type_objs) != 0:
433 final_dtype = _nx.result_type(*result_type_objs)
434 # concatenate could do cast, but that can be overridden:
435 objs = [array(obj, copy=None, subok=True,
436 ndmin=ndmin, dtype=final_dtype) for obj in objs]
438 res = self.concatenate(tuple(objs), axis=axis)
440 if matrix:
441 oldndim = res.ndim
442 res = self.makemat(res)
443 if oldndim == 1 and col:
444 res = res.T
445 return res
447 def __len__(self):
448 return 0
450# separate classes are used here instead of just making r_ = concatenator(0),
451# etc. because otherwise we couldn't get the doc string to come out right
452# in help(r_)
455class RClass(AxisConcatenator):
456 """
457 Translates slice objects to concatenation along the first axis.
459 This is a simple way to build up arrays quickly. There are two use cases.
461 1. If the index expression contains comma separated arrays, then stack
462 them along their first axis.
463 2. If the index expression contains slice notation or scalars then create
464 a 1-D array with a range indicated by the slice notation.
466 If slice notation is used, the syntax ``start:stop:step`` is equivalent
467 to ``np.arange(start, stop, step)`` inside of the brackets. However, if
468 ``step`` is an imaginary number (i.e. 100j) then its integer portion is
469 interpreted as a number-of-points desired and the start and stop are
470 inclusive. In other words ``start:stop:stepj`` is interpreted as
471 ``np.linspace(start, stop, step, endpoint=1)`` inside of the brackets.
472 After expansion of slice notation, all comma separated sequences are
473 concatenated together.
475 Optional character strings placed as the first element of the index
476 expression can be used to change the output. The strings 'r' or 'c' result
477 in matrix output. If the result is 1-D and 'r' is specified a 1 x N (row)
478 matrix is produced. If the result is 1-D and 'c' is specified, then a N x 1
479 (column) matrix is produced. If the result is 2-D then both provide the
480 same matrix result.
482 A string integer specifies which axis to stack multiple comma separated
483 arrays along. A string of two comma-separated integers allows indication
484 of the minimum number of dimensions to force each entry into as the
485 second integer (the axis to concatenate along is still the first integer).
487 A string with three comma-separated integers allows specification of the
488 axis to concatenate along, the minimum number of dimensions to force the
489 entries to, and which axis should contain the start of the arrays which
490 are less than the specified number of dimensions. In other words the third
491 integer allows you to specify where the 1's should be placed in the shape
492 of the arrays that have their shapes upgraded. By default, they are placed
493 in the front of the shape tuple. The third argument allows you to specify
494 where the start of the array should be instead. Thus, a third argument of
495 '0' would place the 1's at the end of the array shape. Negative integers
496 specify where in the new shape tuple the last dimension of upgraded arrays
497 should be placed, so the default is '-1'.
499 Parameters
500 ----------
501 Not a function, so takes no parameters
504 Returns
505 -------
506 A concatenated ndarray or matrix.
508 See Also
509 --------
510 concatenate : Join a sequence of arrays along an existing axis.
511 c_ : Translates slice objects to concatenation along the second axis.
513 Examples
514 --------
515 >>> import numpy as np
516 >>> np.r_[np.array([1,2,3]), 0, 0, np.array([4,5,6])]
517 array([1, 2, 3, ..., 4, 5, 6])
518 >>> np.r_[-1:1:6j, [0]*3, 5, 6]
519 array([-1. , -0.6, -0.2, 0.2, 0.6, 1. , 0. , 0. , 0. , 5. , 6. ])
521 String integers specify the axis to concatenate along or the minimum
522 number of dimensions to force entries into.
524 >>> a = np.array([[0, 1, 2], [3, 4, 5]])
525 >>> np.r_['-1', a, a] # concatenate along last axis
526 array([[0, 1, 2, 0, 1, 2],
527 [3, 4, 5, 3, 4, 5]])
528 >>> np.r_['0,2', [1,2,3], [4,5,6]] # concatenate along first axis, dim>=2
529 array([[1, 2, 3],
530 [4, 5, 6]])
532 >>> np.r_['0,2,0', [1,2,3], [4,5,6]]
533 array([[1],
534 [2],
535 [3],
536 [4],
537 [5],
538 [6]])
539 >>> np.r_['1,2,0', [1,2,3], [4,5,6]]
540 array([[1, 4],
541 [2, 5],
542 [3, 6]])
544 Using 'r' or 'c' as a first string argument creates a matrix.
546 >>> np.r_['r',[1,2,3], [4,5,6]]
547 matrix([[1, 2, 3, 4, 5, 6]])
549 """
550 __slots__ = ()
552 def __init__(self):
553 AxisConcatenator.__init__(self, 0)
556r_ = RClass()
559class CClass(AxisConcatenator):
560 """
561 Translates slice objects to concatenation along the second axis.
563 This is short-hand for ``np.r_['-1,2,0', index expression]``, which is
564 useful because of its common occurrence. In particular, arrays will be
565 stacked along their last axis after being upgraded to at least 2-D with
566 1's post-pended to the shape (column vectors made out of 1-D arrays).
568 See Also
569 --------
570 column_stack : Stack 1-D arrays as columns into a 2-D array.
571 r_ : For more detailed documentation.
573 Examples
574 --------
575 >>> import numpy as np
576 >>> np.c_[np.array([1,2,3]), np.array([4,5,6])]
577 array([[1, 4],
578 [2, 5],
579 [3, 6]])
580 >>> np.c_[np.array([[1,2,3]]), 0, 0, np.array([[4,5,6]])]
581 array([[1, 2, 3, ..., 4, 5, 6]])
583 """
584 __slots__ = ()
586 def __init__(self):
587 AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
590c_ = CClass()
593@set_module('numpy')
594class ndenumerate:
595 """
596 Multidimensional index iterator.
598 Return an iterator yielding pairs of array coordinates and values.
600 Parameters
601 ----------
602 arr : ndarray
603 Input array.
605 See Also
606 --------
607 ndindex, flatiter
609 Examples
610 --------
611 >>> import numpy as np
612 >>> a = np.array([[1, 2], [3, 4]])
613 >>> for index, x in np.ndenumerate(a):
614 ... print(index, x)
615 (0, 0) 1
616 (0, 1) 2
617 (1, 0) 3
618 (1, 1) 4
620 """
622 def __init__(self, arr):
623 self.iter = np.asarray(arr).flat
625 def __next__(self):
626 """
627 Standard iterator method, returns the index tuple and array value.
629 Returns
630 -------
631 coords : tuple of ints
632 The indices of the current iteration.
633 val : scalar
634 The array element of the current iteration.
636 """
637 return self.iter.coords, next(self.iter)
639 def __iter__(self):
640 return self
643@set_module('numpy')
644class ndindex:
645 """
646 An N-dimensional iterator object to index arrays.
648 Given the shape of an array, an `ndindex` instance iterates over
649 the N-dimensional index of the array. At each iteration a tuple
650 of indices is returned, the last dimension is iterated over first.
652 Parameters
653 ----------
654 shape : ints, or a single tuple of ints
655 The size of each dimension of the array can be passed as
656 individual parameters or as the elements of a tuple.
658 See Also
659 --------
660 ndenumerate, flatiter
662 Examples
663 --------
664 >>> import numpy as np
666 Dimensions as individual arguments
668 >>> for index in np.ndindex(3, 2, 1):
669 ... print(index)
670 (0, 0, 0)
671 (0, 1, 0)
672 (1, 0, 0)
673 (1, 1, 0)
674 (2, 0, 0)
675 (2, 1, 0)
677 Same dimensions - but in a tuple ``(3, 2, 1)``
679 >>> for index in np.ndindex((3, 2, 1)):
680 ... print(index)
681 (0, 0, 0)
682 (0, 1, 0)
683 (1, 0, 0)
684 (1, 1, 0)
685 (2, 0, 0)
686 (2, 1, 0)
688 """
690 def __init__(self, *shape):
691 if len(shape) == 1 and isinstance(shape[0], tuple):
692 shape = shape[0]
693 x = as_strided(_nx.zeros(1), shape=shape,
694 strides=_nx.zeros_like(shape))
695 self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'],
696 order='C')
698 def __iter__(self):
699 return self
701 def ndincr(self):
702 """
703 Increment the multi-dimensional index by one.
705 This method is for backward compatibility only: do not use.
707 .. deprecated:: 1.20.0
708 This method has been advised against since numpy 1.8.0, but only
709 started emitting DeprecationWarning as of this version.
710 """
711 # NumPy 1.20.0, 2020-09-08
712 warnings.warn(
713 "`ndindex.ndincr()` is deprecated, use `next(ndindex)` instead",
714 DeprecationWarning, stacklevel=2)
715 next(self)
717 def __next__(self):
718 """
719 Standard iterator method, updates the index and returns the index
720 tuple.
722 Returns
723 -------
724 val : tuple of ints
725 Returns a tuple containing the indices of the current
726 iteration.
728 """
729 next(self._it)
730 return self._it.multi_index
733# You can do all this with slice() plus a few special objects,
734# but there's a lot to remember. This version is simpler because
735# it uses the standard array indexing syntax.
736#
737# Written by Konrad Hinsen <hinsen@cnrs-orleans.fr>
738# last revision: 1999-7-23
739#
740# Cosmetic changes by T. Oliphant 2001
741#
742#
744class IndexExpression:
745 """
746 A nicer way to build up index tuples for arrays.
748 .. note::
749 Use one of the two predefined instances ``index_exp`` or `s_`
750 rather than directly using `IndexExpression`.
752 For any index combination, including slicing and axis insertion,
753 ``a[indices]`` is the same as ``a[np.index_exp[indices]]`` for any
754 array `a`. However, ``np.index_exp[indices]`` can be used anywhere
755 in Python code and returns a tuple of slice objects that can be
756 used in the construction of complex index expressions.
758 Parameters
759 ----------
760 maketuple : bool
761 If True, always returns a tuple.
763 See Also
764 --------
765 s_ : Predefined instance without tuple conversion:
766 `s_ = IndexExpression(maketuple=False)`.
767 The ``index_exp`` is another predefined instance that
768 always returns a tuple:
769 `index_exp = IndexExpression(maketuple=True)`.
771 Notes
772 -----
773 You can do all this with :class:`slice` plus a few special objects,
774 but there's a lot to remember and this version is simpler because
775 it uses the standard array indexing syntax.
777 Examples
778 --------
779 >>> import numpy as np
780 >>> np.s_[2::2]
781 slice(2, None, 2)
782 >>> np.index_exp[2::2]
783 (slice(2, None, 2),)
785 >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]
786 array([2, 4])
788 """
789 __slots__ = ('maketuple',)
791 def __init__(self, maketuple):
792 self.maketuple = maketuple
794 def __getitem__(self, item):
795 if self.maketuple and not isinstance(item, tuple):
796 return (item,)
797 else:
798 return item
801index_exp = IndexExpression(maketuple=True)
802s_ = IndexExpression(maketuple=False)
804# End contribution from Konrad.
807# The following functions complement those in twodim_base, but are
808# applicable to N-dimensions.
811def _fill_diagonal_dispatcher(a, val, wrap=None):
812 return (a,)
815@array_function_dispatch(_fill_diagonal_dispatcher)
816def fill_diagonal(a, val, wrap=False):
817 """Fill the main diagonal of the given array of any dimensionality.
819 For an array `a` with ``a.ndim >= 2``, the diagonal is the list of
820 values ``a[i, ..., i]`` with indices ``i`` all identical. This function
821 modifies the input array in-place without returning a value.
823 Parameters
824 ----------
825 a : array, at least 2-D.
826 Array whose diagonal is to be filled in-place.
827 val : scalar or array_like
828 Value(s) to write on the diagonal. If `val` is scalar, the value is
829 written along the diagonal. If array-like, the flattened `val` is
830 written along the diagonal, repeating if necessary to fill all
831 diagonal entries.
833 wrap : bool
834 For tall matrices in NumPy version up to 1.6.2, the
835 diagonal "wrapped" after N columns. You can have this behavior
836 with this option. This affects only tall matrices.
838 See also
839 --------
840 diag_indices, diag_indices_from
842 Notes
843 -----
844 This functionality can be obtained via `diag_indices`, but internally
845 this version uses a much faster implementation that never constructs the
846 indices and uses simple slicing.
848 Examples
849 --------
850 >>> import numpy as np
851 >>> a = np.zeros((3, 3), int)
852 >>> np.fill_diagonal(a, 5)
853 >>> a
854 array([[5, 0, 0],
855 [0, 5, 0],
856 [0, 0, 5]])
858 The same function can operate on a 4-D array:
860 >>> a = np.zeros((3, 3, 3, 3), int)
861 >>> np.fill_diagonal(a, 4)
863 We only show a few blocks for clarity:
865 >>> a[0, 0]
866 array([[4, 0, 0],
867 [0, 0, 0],
868 [0, 0, 0]])
869 >>> a[1, 1]
870 array([[0, 0, 0],
871 [0, 4, 0],
872 [0, 0, 0]])
873 >>> a[2, 2]
874 array([[0, 0, 0],
875 [0, 0, 0],
876 [0, 0, 4]])
878 The wrap option affects only tall matrices:
880 >>> # tall matrices no wrap
881 >>> a = np.zeros((5, 3), int)
882 >>> np.fill_diagonal(a, 4)
883 >>> a
884 array([[4, 0, 0],
885 [0, 4, 0],
886 [0, 0, 4],
887 [0, 0, 0],
888 [0, 0, 0]])
890 >>> # tall matrices wrap
891 >>> a = np.zeros((5, 3), int)
892 >>> np.fill_diagonal(a, 4, wrap=True)
893 >>> a
894 array([[4, 0, 0],
895 [0, 4, 0],
896 [0, 0, 4],
897 [0, 0, 0],
898 [4, 0, 0]])
900 >>> # wide matrices
901 >>> a = np.zeros((3, 5), int)
902 >>> np.fill_diagonal(a, 4, wrap=True)
903 >>> a
904 array([[4, 0, 0, 0, 0],
905 [0, 4, 0, 0, 0],
906 [0, 0, 4, 0, 0]])
908 The anti-diagonal can be filled by reversing the order of elements
909 using either `numpy.flipud` or `numpy.fliplr`.
911 >>> a = np.zeros((3, 3), int);
912 >>> np.fill_diagonal(np.fliplr(a), [1,2,3]) # Horizontal flip
913 >>> a
914 array([[0, 0, 1],
915 [0, 2, 0],
916 [3, 0, 0]])
917 >>> np.fill_diagonal(np.flipud(a), [1,2,3]) # Vertical flip
918 >>> a
919 array([[0, 0, 3],
920 [0, 2, 0],
921 [1, 0, 0]])
923 Note that the order in which the diagonal is filled varies depending
924 on the flip function.
925 """
926 if a.ndim < 2:
927 raise ValueError("array must be at least 2-d")
928 end = None
929 if a.ndim == 2:
930 # Explicit, fast formula for the common case. For 2-d arrays, we
931 # accept rectangular ones.
932 step = a.shape[1] + 1
933 # This is needed to don't have tall matrix have the diagonal wrap.
934 if not wrap:
935 end = a.shape[1] * a.shape[1]
936 else:
937 # For more than d=2, the strided formula is only valid for arrays with
938 # all dimensions equal, so we check first.
939 if not np.all(diff(a.shape) == 0):
940 raise ValueError("All dimensions of input must be of equal length")
941 step = 1 + (np.cumprod(a.shape[:-1])).sum()
943 # Write the value out into the diagonal.
944 a.flat[:end:step] = val
947@set_module('numpy')
948def diag_indices(n, ndim=2):
949 """
950 Return the indices to access the main diagonal of an array.
952 This returns a tuple of indices that can be used to access the main
953 diagonal of an array `a` with ``a.ndim >= 2`` dimensions and shape
954 (n, n, ..., n). For ``a.ndim = 2`` this is the usual diagonal, for
955 ``a.ndim > 2`` this is the set of indices to access ``a[i, i, ..., i]``
956 for ``i = [0..n-1]``.
958 Parameters
959 ----------
960 n : int
961 The size, along each dimension, of the arrays for which the returned
962 indices can be used.
964 ndim : int, optional
965 The number of dimensions.
967 See Also
968 --------
969 diag_indices_from
971 Examples
972 --------
973 >>> import numpy as np
975 Create a set of indices to access the diagonal of a (4, 4) array:
977 >>> di = np.diag_indices(4)
978 >>> di
979 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
980 >>> a = np.arange(16).reshape(4, 4)
981 >>> a
982 array([[ 0, 1, 2, 3],
983 [ 4, 5, 6, 7],
984 [ 8, 9, 10, 11],
985 [12, 13, 14, 15]])
986 >>> a[di] = 100
987 >>> a
988 array([[100, 1, 2, 3],
989 [ 4, 100, 6, 7],
990 [ 8, 9, 100, 11],
991 [ 12, 13, 14, 100]])
993 Now, we create indices to manipulate a 3-D array:
995 >>> d3 = np.diag_indices(2, 3)
996 >>> d3
997 (array([0, 1]), array([0, 1]), array([0, 1]))
999 And use it to set the diagonal of an array of zeros to 1:
1001 >>> a = np.zeros((2, 2, 2), dtype=int)
1002 >>> a[d3] = 1
1003 >>> a
1004 array([[[1, 0],
1005 [0, 0]],
1006 [[0, 0],
1007 [0, 1]]])
1009 """
1010 idx = np.arange(n)
1011 return (idx,) * ndim
1014def _diag_indices_from(arr):
1015 return (arr,)
1018@array_function_dispatch(_diag_indices_from)
1019def diag_indices_from(arr):
1020 """
1021 Return the indices to access the main diagonal of an n-dimensional array.
1023 See `diag_indices` for full details.
1025 Parameters
1026 ----------
1027 arr : array, at least 2-D
1029 See Also
1030 --------
1031 diag_indices
1033 Examples
1034 --------
1035 >>> import numpy as np
1037 Create a 4 by 4 array.
1039 >>> a = np.arange(16).reshape(4, 4)
1040 >>> a
1041 array([[ 0, 1, 2, 3],
1042 [ 4, 5, 6, 7],
1043 [ 8, 9, 10, 11],
1044 [12, 13, 14, 15]])
1046 Get the indices of the diagonal elements.
1048 >>> di = np.diag_indices_from(a)
1049 >>> di
1050 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1052 >>> a[di]
1053 array([ 0, 5, 10, 15])
1055 This is simply syntactic sugar for diag_indices.
1057 >>> np.diag_indices(a.shape[0])
1058 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1060 """
1062 if not arr.ndim >= 2:
1063 raise ValueError("input array must be at least 2-d")
1064 # For more than d=2, the strided formula is only valid for arrays with
1065 # all dimensions equal, so we check first.
1066 if not np.all(diff(arr.shape) == 0):
1067 raise ValueError("All dimensions of input must be of equal length")
1069 return diag_indices(arr.shape[0], arr.ndim)