ENH: Speed up trim_zeros (#16911)

* Added a benchmark for `trim_zeros()`

* Improve the performance of `np.trim_zeros()`

* Increase the variety of the tests

Fall back to the old `np.trim_zeros()` implementation if an exception is encountered.
Emit a `DeprecationWarning` in such case.


* DEP,REL: Added a deprecation release note

5 files changed, 146 insertions, 21 deletions

diff --git a/benchmarks/benchmarks/bench_trim_zeros.py b/benchmarks/benchmarks/bench_trim_zeros.py
new file mode 100644
index 000000000..4e25a8b02
--- /dev/null
+++ b/benchmarks/benchmarks/bench_trim_zeros.py
@@ -0,0 +1,27 @@
+from .common import Benchmark
+
+import numpy as np
+
+_FLOAT = np.dtype('float64')
+_COMPLEX = np.dtype('complex128')
+_INT = np.dtype('int64')
+_BOOL = np.dtype('bool')
+
+
+class TrimZeros(Benchmark):
+    param_names = ["dtype", "size"]
+    params = [
+        [_INT, _FLOAT, _COMPLEX, _BOOL],
+        [3000, 30_000, 300_000]
+    ]
+
+    def setup(self, dtype, size):
+        n = size // 3
+        self.array = np.hstack([
+            np.zeros(n),
+            np.random.uniform(size=n),
+            np.zeros(n),
+        ]).astype(dtype)
+
+    def time_trim_zeros(self, dtype, size):
+        np.trim_zeros(self.array)
diff --git a/doc/release/upcoming_changes/16911.deprecation.rst b/doc/release/upcoming_changes/16911.deprecation.rst
new file mode 100644
index 000000000..d4dcb629c
--- /dev/null
+++ b/doc/release/upcoming_changes/16911.deprecation.rst
@@ -0,0 +1,7 @@
+``trim_zeros`` now requires a 1D array compatible with ``ndarray.astype(bool)``
+-------------------------------------------------------------------------------
+The ``trim_zeros`` function will, in the future, require an array with the
+following two properties:
+
+* It must be 1D.
+* It must be convertable into a boolean array.
diff --git a/numpy/core/tests/test_deprecations.py b/numpy/core/tests/test_deprecations.py
index 47258ed08..9004bef30 100644
--- a/numpy/core/tests/test_deprecations.py
+++ b/numpy/core/tests/test_deprecations.py
@@ -615,7 +615,7 @@ class BuiltInRoundComplexDType(_DeprecationTestCase):
             self.assert_deprecated(round, args=(scalar,))
             self.assert_deprecated(round, args=(scalar, 0))
             self.assert_deprecated(round, args=(scalar,), kwargs={'ndigits': 0})
-    
+
     def test_not_deprecated(self):
         for scalar_type in self.not_deprecated_types:
             scalar = scalar_type(0)
@@ -706,3 +706,21 @@ class TestRaggedArray(_DeprecationTestCase):
         # And when it is an assignment into a lower dimensional subarray:
         self.assert_deprecated(lambda: np.array([arr, [0]], dtype=np.float64))
         self.assert_deprecated(lambda: np.array([[0], arr], dtype=np.float64))
+
+
+class TestTrimZeros(_DeprecationTestCase):
+    # Numpy 1.20.0, 2020-07-31
+    @pytest.mark.parametrize("arr", [np.random.rand(10, 10).tolist(),
+                                     np.random.rand(10).astype(str)])
+    def test_deprecated(self, arr):
+        with warnings.catch_warnings():
+            warnings.simplefilter('error', DeprecationWarning)
+            try:
+                np.trim_zeros(arr)
+            except DeprecationWarning as ex:
+                assert_(isinstance(ex.__cause__, ValueError))
+            else:
+                raise AssertionError("No error raised during function call")
+
+        out = np.lib.function_base._trim_zeros_old(arr)
+        assert_array_equal(arr, out)
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 6ea9cc4de..cd8862c94 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -433,7 +433,7 @@ def asarray_chkfinite(a, dtype=None, order=None):
         By default, the data-type is inferred from the input data.
     order : {'C', 'F', 'A', 'K'}, optional
         Memory layout.  'A' and 'K' depend on the order of input array a.
-        'C' row-major (C-style), 
+        'C' row-major (C-style),
         'F' column-major (Fortran-style) memory representation.
         'A' (any) means 'F' if `a` is Fortran contiguous, 'C' otherwise
         'K' (keep) preserve input order
@@ -1625,6 +1625,57 @@ def trim_zeros(filt, trim='fb'):
     [1, 2]
 
     """
+    try:
+        return _trim_zeros_new(filt, trim)
+    except Exception as ex:
+        # Numpy 1.20.0, 2020-07-31
+        warning = DeprecationWarning(
+            "in the future trim_zeros will require a 1-D array as input "
+            "that is compatible with ndarray.astype(bool)"
+        )
+        warning.__cause__ = ex
+        warnings.warn(warning, stacklevel=3)
+
+        # Fall back to the old implementation if an exception is encountered
+        # Note that the same exception may or may not be raised here as well
+        return _trim_zeros_old(filt, trim)
+
+
+def _trim_zeros_new(filt, trim='fb'):
+    """Newer optimized implementation of ``trim_zeros()``."""
+    arr = np.asanyarray(filt).astype(bool, copy=False)
+
+    if arr.ndim != 1:
+        raise ValueError('trim_zeros requires an array of exactly one dimension')
+    elif not len(arr):
+        return filt
+
+    trim_upper = trim.upper()
+    first = last = None
+
+    if 'F' in trim_upper:
+        first = arr.argmax()
+        # If `arr[first] is False` then so are all other elements
+        if not arr[first]:
+            return filt[:0]
+
+    if 'B' in trim_upper:
+        last = len(arr) - arr[::-1].argmax()
+        # If `arr[last - 1] is False` then so are all other elements
+        if not arr[last - 1]:
+            return filt[:0]
+
+    return filt[first:last]
+
+
+def _trim_zeros_old(filt, trim='fb'):
+    """
+    Older unoptimized implementation of ``trim_zeros()``.
+
+    Used as fallback in case an exception is encountered
+    in ``_trim_zeros_new()``.
+
+    """
     first = 0
     trim = trim.upper()
     if 'F' in trim:
@@ -2546,11 +2597,11 @@ def corrcoef(x, y=None, rowvar=True, bias=np._NoValue, ddof=np._NoValue):
     for backwards compatibility with previous versions of this function.  These
     arguments had no effect on the return values of the function and can be
     safely ignored in this and previous versions of numpy.
-    
+
     Examples
-    --------   
+    --------
     In this example we generate two random arrays, ``xarr`` and ``yarr``, and
-    compute the row-wise and column-wise Pearson correlation coefficients, 
+    compute the row-wise and column-wise Pearson correlation coefficients,
     ``R``. Since ``rowvar`` is  true by  default, we first find the row-wise
     Pearson correlation coefficients between the variables of ``xarr``.
 
@@ -2566,11 +2617,11 @@ def corrcoef(x, y=None, rowvar=True, bias=np._NoValue, ddof=np._NoValue):
     array([[ 1.        ,  0.99256089, -0.68080986],
            [ 0.99256089,  1.        , -0.76492172],
            [-0.68080986, -0.76492172,  1.        ]])
-    
-    If we add another set of variables and observations ``yarr``, we can 
+
+    If we add another set of variables and observations ``yarr``, we can
     compute the row-wise Pearson correlation coefficients between the
     variables in ``xarr`` and ``yarr``.
-   
+
     >>> yarr = rng.random((3, 3))
     >>> yarr
     array([[0.45038594, 0.37079802, 0.92676499],
@@ -2592,7 +2643,7 @@ def corrcoef(x, y=None, rowvar=True, bias=np._NoValue, ddof=np._NoValue):
              1.        ]])
 
     Finally if we use the option ``rowvar=False``, the columns are now
-    being treated as the variables and we will find the column-wise Pearson 
+    being treated as the variables and we will find the column-wise Pearson
     correlation coefficients between variables in ``xarr`` and ``yarr``.
 
     >>> R3 = np.corrcoef(xarr, yarr, rowvar=False)
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index eb2fc3311..89c1a2d9b 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1166,25 +1166,47 @@ class TestAngle:
 
 class TestTrimZeros:
 
-    """
-    Only testing for integer splits.
+    a = np.array([0, 0, 1, 0, 2, 3, 4, 0])
+    b = a.astype(float)
+    c = a.astype(complex)
+    d = np.array([None, [], 1, False, 'b', 3.0, range(4), b''], dtype=object)
 
-    """
+    def values(self):
+        attr_names = ('a', 'b', 'c', 'd')
+        return (getattr(self, name) for name in attr_names)
 
     def test_basic(self):
-        a = np.array([0, 0, 1, 2, 3, 4, 0])
-        res = trim_zeros(a)
-        assert_array_equal(res, np.array([1, 2, 3, 4]))
+        slc = np.s_[2:-1]
+        for arr in self.values():
+            res = trim_zeros(arr)
+            assert_array_equal(res, arr[slc])
 
     def test_leading_skip(self):
-        a = np.array([0, 0, 1, 0, 2, 3, 4, 0])
-        res = trim_zeros(a)
-        assert_array_equal(res, np.array([1, 0, 2, 3, 4]))
+        slc = np.s_[:-1]
+        for arr in self.values():
+            res = trim_zeros(arr, trim='b')
+            assert_array_equal(res, arr[slc])
 
     def test_trailing_skip(self):
-        a = np.array([0, 0, 1, 0, 2, 3, 0, 4, 0])
-        res = trim_zeros(a)
-        assert_array_equal(res, np.array([1, 0, 2, 3, 0, 4]))
+        slc = np.s_[2:]
+        for arr in self.values():
+            res = trim_zeros(arr, trim='F')
+            assert_array_equal(res, arr[slc])
+
+    def test_all_zero(self):
+        for _arr in self.values():
+            arr = np.zeros_like(_arr, dtype=_arr.dtype)
+
+            res1 = trim_zeros(arr, trim='B')
+            assert len(res1) == 0
+
+            res2 = trim_zeros(arr, trim='f')
+            assert len(res2) == 0
+
+    def test_size_zero(self):
+        arr = np.zeros(0)
+        res = trim_zeros(arr)
+        assert_array_equal(arr, res)
 
 
 class TestExtins: