MAINT: Merge branch 'main' into delete-speedup

1 files changed, 69 insertions, 35 deletions

diff --git a/numpy/lib/index_tricks.py b/numpy/lib/index_tricks.py
index d145477c3..b69226d48 100644
--- a/numpy/lib/index_tricks.py
+++ b/numpy/lib/index_tricks.py
@@ -1,11 +1,12 @@
 import functools
 import sys
 import math
+import warnings
 
 import numpy.core.numeric as _nx
 from numpy.core.numeric import (
     asarray, ScalarType, array, alltrue, cumprod, arange, ndim
-    )
+)
 from numpy.core.numerictypes import find_common_type, issubdtype
 
 import numpy.matrixlib as matrixlib
@@ -24,7 +25,7 @@ __all__ = [
     'ravel_multi_index', 'unravel_index', 'mgrid', 'ogrid', 'r_', 'c_',
     's_', 'index_exp', 'ix_', 'ndenumerate', 'ndindex', 'fill_diagonal',
     'diag_indices', 'diag_indices_from'
-    ]
+]
 
 
 def _ix__dispatcher(*args):
@@ -105,6 +106,7 @@ def ix_(*args):
         out.append(new)
     return tuple(out)
 
+
 class nd_grid:
     """
     Construct a multi-dimensional "meshgrid".
@@ -147,39 +149,39 @@ class nd_grid:
         try:
             size = []
             typ = int
-            for k in range(len(key)):
-                step = key[k].step
-                start = key[k].start
+            for kk in key:
+                step = kk.step
+                start = kk.start
                 if start is None:
                     start = 0
                 if step is None:
                     step = 1
-                if isinstance(step, complex):
+                if isinstance(step, (_nx.complexfloating, complex)):
                     size.append(int(abs(step)))
                     typ = float
                 else:
                     size.append(
-                        int(math.ceil((key[k].stop - start)/(step*1.0))))
-                if (isinstance(step, float) or
-                        isinstance(start, float) or
-                        isinstance(key[k].stop, float)):
+                        int(math.ceil((kk.stop - start) / (step * 1.0))))
+                if (isinstance(step, (_nx.floating, float)) or
+                        isinstance(start, (_nx.floating, float)) or
+                        isinstance(kk.stop, (_nx.floating, float))):
                     typ = float
             if self.sparse:
                 nn = [_nx.arange(_x, dtype=_t)
-                        for _x, _t in zip(size, (typ,)*len(size))]
+                      for _x, _t in zip(size, (typ,)*len(size))]
             else:
                 nn = _nx.indices(size, typ)
-            for k in range(len(size)):
-                step = key[k].step
-                start = key[k].start
+            for k, kk in enumerate(key):
+                step = kk.step
+                start = kk.start
                 if start is None:
                     start = 0
                 if step is None:
                     step = 1
-                if isinstance(step, complex):
+                if isinstance(step, (_nx.complexfloating, complex)):
                     step = int(abs(step))
                     if step != 1:
-                        step = (key[k].stop - start)/float(step-1)
+                        step = (kk.stop - start) / float(step - 1)
                 nn[k] = (nn[k]*step+start)
             if self.sparse:
                 slobj = [_nx.newaxis]*len(size)
@@ -194,12 +196,11 @@ class nd_grid:
             start = key.start
             if start is None:
                 start = 0
-            if isinstance(step, complex):
+            if isinstance(step, (_nx.complexfloating, complex)):
                 step = abs(step)
                 length = int(step)
                 if step != 1:
                     step = (key.stop-start)/float(step-1)
-                stop = key.stop + step
                 return _nx.arange(0, length, 1, float)*step + start
             else:
                 return _nx.arange(start, stop, step)
@@ -221,18 +222,18 @@ class MGridClass(nd_grid):
     the stop value **is inclusive**.
 
     Returns
-    ----------
+    -------
     mesh-grid `ndarrays` all of the same dimensions
 
     See Also
     --------
-    numpy.lib.index_tricks.nd_grid : class of `ogrid` and `mgrid` objects
+    lib.index_tricks.nd_grid : class of `ogrid` and `mgrid` objects
     ogrid : like mgrid but returns open (not fleshed out) mesh grids
     r_ : array concatenator
 
     Examples
     --------
-    >>> np.mgrid[0:5,0:5]
+    >>> np.mgrid[0:5, 0:5]
     array([[[0, 0, 0, 0, 0],
             [1, 1, 1, 1, 1],
             [2, 2, 2, 2, 2],
@@ -247,11 +248,14 @@ class MGridClass(nd_grid):
     array([-1. , -0.5,  0. ,  0.5,  1. ])
 
     """
+
     def __init__(self):
-        super(MGridClass, self).__init__(sparse=False)
+        super().__init__(sparse=False)
+
 
 mgrid = MGridClass()
 
+
 class OGridClass(nd_grid):
     """
     `nd_grid` instance which returns an open multi-dimensional "meshgrid".
@@ -291,8 +295,10 @@ class OGridClass(nd_grid):
             [4]]), array([[0, 1, 2, 3, 4]])]
 
     """
+
     def __init__(self):
-        super(OGridClass, self).__init__(sparse=True)
+        super().__init__(sparse=True)
+
 
 ogrid = OGridClass()
 
@@ -344,7 +350,7 @@ class AxisConcatenator:
                     start = 0
                 if step is None:
                     step = 1
-                if isinstance(step, complex):
+                if isinstance(step, (_nx.complexfloating, complex)):
                     size = int(abs(step))
                     newobj = linspace(start, stop, num=size)
                 else:
@@ -356,7 +362,7 @@ class AxisConcatenator:
             elif isinstance(item, str):
                 if k != 0:
                     raise ValueError("special directives must be the "
-                            "first entry.")
+                                     "first entry.")
                 if item in ('r', 'c'):
                     matrix = True
                     col = (item == 'c')
@@ -375,8 +381,8 @@ class AxisConcatenator:
                 try:
                     axis = int(item)
                     continue
-                except (ValueError, TypeError):
-                    raise ValueError("unknown special directive")
+                except (ValueError, TypeError) as e:
+                    raise ValueError("unknown special directive") from e
             elif type(item) in ScalarType:
                 newobj = array(item, ndmin=ndmin)
                 scalars.append(len(objs))
@@ -419,6 +425,7 @@ class AxisConcatenator:
 # etc. because otherwise we couldn't get the doc string to come out right
 # in help(r_)
 
+
 class RClass(AxisConcatenator):
     """
     Translates slice objects to concatenation along the first axis.
@@ -517,8 +524,10 @@ class RClass(AxisConcatenator):
     def __init__(self):
         AxisConcatenator.__init__(self, 0)
 
+
 r_ = RClass()
 
+
 class CClass(AxisConcatenator):
     """
     Translates slice objects to concatenation along the second axis.
@@ -527,7 +536,7 @@ class CClass(AxisConcatenator):
     useful because of its common occurrence. In particular, arrays will be
     stacked along their last axis after being upgraded to at least 2-D with
     1's post-pended to the shape (column vectors made out of 1-D arrays).
-    
+
     See Also
     --------
     column_stack : Stack 1-D arrays as columns into a 2-D array.
@@ -611,8 +620,9 @@ class ndindex:
 
     Parameters
     ----------
-    `*args` : ints
-      The size of each dimension of the array.
+    shape : ints, or a single tuple of ints
+        The size of each dimension of the array can be passed as 
+        individual parameters or as the elements of a tuple.
 
     See Also
     --------
@@ -620,6 +630,8 @@ class ndindex:
 
     Examples
     --------
+    Dimensions as individual arguments
+    
     >>> for index in np.ndindex(3, 2, 1):
     ...     print(index)
     (0, 0, 0)
@@ -629,6 +641,17 @@ class ndindex:
     (2, 0, 0)
     (2, 1, 0)
 
+    Same dimensions - but in a tuple ``(3, 2, 1)``
+    
+    >>> for index in np.ndindex((3, 2, 1)):
+    ...     print(index)
+    (0, 0, 0)
+    (0, 1, 0)
+    (1, 0, 0)
+    (1, 1, 0)
+    (2, 0, 0)
+    (2, 1, 0)
+
     """
 
     def __init__(self, *shape):
@@ -647,7 +670,15 @@ class ndindex:
         Increment the multi-dimensional index by one.
 
         This method is for backward compatibility only: do not use.
+
+        .. deprecated:: 1.20.0
+            This method has been advised against since numpy 1.8.0, but only
+            started emitting DeprecationWarning as of this version.
         """
+        # NumPy 1.20.0, 2020-09-08
+        warnings.warn(
+            "`ndindex.ndincr()` is deprecated, use `next(ndindex)` instead",
+            DeprecationWarning, stacklevel=2)
         next(self)
 
     def __next__(self):
@@ -730,6 +761,7 @@ class IndexExpression:
         else:
             return item
 
+
 index_exp = IndexExpression(maketuple=True)
 s_ = IndexExpression(maketuple=False)
 
@@ -757,9 +789,11 @@ def fill_diagonal(a, val, wrap=False):
     a : array, at least 2-D.
       Array whose diagonal is to be filled, it gets modified in-place.
 
-    val : scalar
-      Value to be written on the diagonal, its type must be compatible with
-      that of the array a.
+    val : scalar or array_like
+      Value(s) to write on the diagonal. If `val` is scalar, the value is
+      written along the diagonal. If array-like, the flattened `val` is
+      written along the diagonal, repeating if necessary to fill all
+      diagonal entries.
 
     wrap : bool
       For tall matrices in NumPy version up to 1.6.2, the
@@ -862,7 +896,7 @@ def fill_diagonal(a, val, wrap=False):
         # Explicit, fast formula for the common case.  For 2-d arrays, we
         # accept rectangular ones.
         step = a.shape[1] + 1
-        #This is needed to don't have tall matrix have the diagonal wrap.
+        # This is needed to don't have tall matrix have the diagonal wrap.
         if not wrap:
             end = a.shape[1] * a.shape[1]
     else:
@@ -896,7 +930,7 @@ def diag_indices(n, ndim=2):
     ndim : int, optional
       The number of dimensions.
 
-    See also
+    See Also
     --------
     diag_indices_from