codekingpro/portable-devtools
114k
1import functools
2import math
3import sys
4from itertools import product
5
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
15
16array_function_dispatch = functools.partial(
17 overrides.array_function_dispatch, module='numpy')
18
19
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]
25
26
27def _ix__dispatcher(*args):
28 return args
29
30
31@array_function_dispatch(_ix__dispatcher)
32def ix_(*args):
33 """
34 Construct an open mesh from multiple sequences.
35
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.
40
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]]]``.
44
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)``).
52
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.
58
59 See Also
60 --------
61 ogrid, mgrid, meshgrid
62
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]])
79
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]])
88
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)
105
106
107class nd_grid:
108 """
109 Construct a multi-dimensional "meshgrid".
110
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.
115
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**.
120
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.
124
125 Parameters
126 ----------
127 sparse : bool, optional
128 Whether the grid is sparse or not. Default is False.
129
130 Notes
131 -----
132 Two instances of `nd_grid` are made available in the NumPy namespace,
133 `mgrid` and `ogrid`, approximately defined as::
134
135 mgrid = nd_grid(sparse=False)
136 ogrid = nd_grid(sparse=True)
137
138 Users should use these pre-defined instances instead of using `nd_grid`
139 directly.
140 """
141 __slots__ = ('sparse',)
142
143 def __init__(self, sparse=False):
144 self.sparse = sparse
145
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)
209
210
211class MGridClass(nd_grid):
212 """
213 An instance which returns a dense multi-dimensional "meshgrid".
214
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.
220
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**.
225
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.
231
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`
238
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. ])
255
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)
262
263 """
264 __slots__ = ()
265
266 def __init__(self):
267 super().__init__(sparse=False)
268
269
270mgrid = MGridClass()
271
272
273class OGridClass(nd_grid):
274 """
275 An instance which returns an open multi-dimensional "meshgrid".
276
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.
282
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**.
287
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.
294
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`
301
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]]))
314
315 """
316 __slots__ = ()
317
318 def __init__(self):
319 super().__init__(sparse=True)
320
321
322ogrid = OGridClass()
323
324
325class AxisConcatenator:
326 """
327 Translates slice objects to concatenation along an axis.
328
329 For detailed documentation on usage, see `r_`.
330 """
331 __slots__ = ('axis', 'matrix', 'ndmin', 'trans1d')
332
333 # allow ma.mr_ to override this
334 concatenate = staticmethod(_nx.concatenate)
335 makemat = staticmethod(matrixlib.matrix)
336
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
342
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
349
350 if not isinstance(key, tuple):
351 key = (key,)
352
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
358
359 objs = []
360 # dtypes or scalars for weak scalar handling in result_type
361 result_type_objs = []
362
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)
420
421 objs.append(newobj)
422 if scalar:
423 result_type_objs.append(item)
424 else:
425 result_type_objs.append(newobj.dtype)
426
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]
434
435 res = self.concatenate(tuple(objs), axis=axis)
436
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
443
444 def __len__(self):
445 return 0
446
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_)
450
451
452class RClass(AxisConcatenator):
453 """
454 Translates slice objects to concatenation along the first axis.
455
456 This is a simple way to build up arrays quickly. There are two use cases.
457
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.
462
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.
471
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.
478
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).
483
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'.
495
496 Parameters
497 ----------
498 Not a function, so takes no parameters
499
500
501 Returns
502 -------
503 A concatenated ndarray or matrix.
504
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.
509
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. ])
517
518 String integers specify the axis to concatenate along or the minimum
519 number of dimensions to force entries into.
520
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]])
528
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]])
540
541 Using 'r' or 'c' as a first string argument creates a matrix.
542
543 >>> np.r_['r',[1,2,3], [4,5,6]]
544 matrix([[1, 2, 3, 4, 5, 6]])
545
546 """
547 __slots__ = ()
548
549 def __init__(self):
550 AxisConcatenator.__init__(self, 0)
551
552
553r_ = RClass()
554
555
556class CClass(AxisConcatenator):
557 """
558 Translates slice objects to concatenation along the second axis.
559
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).
564
565 See Also
566 --------
567 column_stack : Stack 1-D arrays as columns into a 2-D array.
568 r_ : For more detailed documentation.
569
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]])
579
580 """
581 __slots__ = ()
582
583 def __init__(self):
584 AxisConcatenator.__init__(self, -1, ndmin=2, trans1d=0)
585
586
587c_ = CClass()
588
589
590@set_module('numpy')
591class ndenumerate:
592 """
593 Multidimensional index iterator.
594
595 Return an iterator yielding pairs of array coordinates and values.
596
597 Parameters
598 ----------
599 arr : ndarray
600 Input array.
601
602 See Also
603 --------
604 ndindex, flatiter
605
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
616
617 """
618
619 def __init__(self, arr):
620 self.iter = np.asarray(arr).flat
621
622 def __next__(self):
623 """
624 Standard iterator method, returns the index tuple and array value.
625
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.
632
633 """
634 return self.iter.coords, next(self.iter)
635
636 def __iter__(self):
637 return self
638
639
640@set_module('numpy')
641class ndindex:
642 """
643 An N-dimensional iterator object to index arrays.
644
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.
648
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.
654
655 See Also
656 --------
657 ndenumerate, flatiter
658
659 Examples
660 --------
661 >>> import numpy as np
662
663 Dimensions as individual arguments
664
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)
673
674 Same dimensions - but in a tuple ``(3, 2, 1)``
675
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)
684
685 """
686
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))
693
694 def __iter__(self):
695 return self
696
697 def __next__(self):
698 """
699 Standard iterator method, updates the index and returns the index
700 tuple.
701
702 Returns
703 -------
704 val : tuple of ints
705 Returns a tuple containing the indices of the current
706 iteration.
707
708 """
709 return next(self._iter)
710
711
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#
722
723class IndexExpression:
724 """
725 A nicer way to build up index tuples for arrays.
726
727 .. note::
728 Use one of the two predefined instances ``index_exp`` or `s_`
729 rather than directly using `IndexExpression`.
730
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.
736
737 Parameters
738 ----------
739 maketuple : bool
740 If True, always returns a tuple.
741
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)`.
749
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.
755
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),)
763
764 >>> np.array([0, 1, 2, 3, 4])[np.s_[2::2]]
765 array([2, 4])
766
767 """
768 __slots__ = ('maketuple',)
769
770 def __init__(self, maketuple):
771 self.maketuple = maketuple
772
773 def __getitem__(self, item):
774 if self.maketuple and not isinstance(item, tuple):
775 return (item,)
776 else:
777 return item
778
779
780index_exp = IndexExpression(maketuple=True)
781s_ = IndexExpression(maketuple=False)
782
783# End contribution from Konrad.
784
785
786# The following functions complement those in twodim_base, but are
787# applicable to N-dimensions.
788
789
790def _fill_diagonal_dispatcher(a, val, wrap=None):
791 return (a,)
792
793
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.
797
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.
801
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.
811
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.
816
817 See also
818 --------
819 diag_indices, diag_indices_from
820
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.
826
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]])
836
837 The same function can operate on a 4-D array:
838
839 >>> a = np.zeros((3, 3, 3, 3), int)
840 >>> np.fill_diagonal(a, 4)
841
842 We only show a few blocks for clarity:
843
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]])
856
857 The wrap option affects only tall matrices:
858
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]])
868
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]])
878
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]])
886
887 The anti-diagonal can be filled by reversing the order of elements
888 using either `numpy.flipud` or `numpy.fliplr`.
889
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]])
901
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()
921
922 # Write the value out into the diagonal.
923 a.flat[:end:step] = val
924
925
926@set_module('numpy')
927def diag_indices(n, ndim=2):
928 """
929 Return the indices to access the main diagonal of an array.
930
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]``.
936
937 Parameters
938 ----------
939 n : int
940 The size, along each dimension, of the arrays for which the returned
941 indices can be used.
942
943 ndim : int, optional
944 The number of dimensions.
945
946 See Also
947 --------
948 diag_indices_from
949
950 Examples
951 --------
952 >>> import numpy as np
953
954 Create a set of indices to access the diagonal of a (4, 4) array:
955
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]])
971
972 Now, we create indices to manipulate a 3-D array:
973
974 >>> d3 = np.diag_indices(2, 3)
975 >>> d3
976 (array([0, 1]), array([0, 1]), array([0, 1]))
977
978 And use it to set the diagonal of an array of zeros to 1:
979
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]]])
987
988 """
989 idx = np.arange(n)
990 return (idx,) * ndim
991
992
993def _diag_indices_from(arr):
994 return (arr,)
995
996
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.
1001
1002 See `diag_indices` for full details.
1003
1004 Parameters
1005 ----------
1006 arr : array, at least 2-D
1007
1008 See Also
1009 --------
1010 diag_indices
1011
1012 Examples
1013 --------
1014 >>> import numpy as np
1015
1016 Create a 4 by 4 array.
1017
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]])
1024
1025 Get the indices of the diagonal elements.
1026
1027 >>> di = np.diag_indices_from(a)
1028 >>> di
1029 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1030
1031 >>> a[di]
1032 array([ 0, 5, 10, 15])
1033
1034 This is simply syntactic sugar for diag_indices.
1035
1036 >>> np.diag_indices(a.shape[0])
1037 (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
1038
1039 """
1040
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")
1047
1048 return diag_indices(arr.shape[0], arr.ndim)
1049