ENH: speed-up mlab.contiguous_regions using numpy

lib/matplotlib/mlab.py

@@ -3874,23 +3874,26 @@ def contiguous_regions(mask):
     """
     return a list of (ind0, ind1) such that mask[ind0:ind1].all() is
     True and we cover all such regions
-
-    TODO: this is a pure python implementation which probably has a much
-    faster numpy impl
     """
+    mask = np.asarray(mask, dtype=bool)
+
+    if not mask.size:
+        return []
+
+    # Find the indices of region changes, and correct offset
+    idx, = np.nonzero(mask[:-1] != mask[1:])
+    idx += 1
+
+    # List operations are faster for moderately sized arrays
+    idx = idx.tolist()
 
-    in_region = None
-    boundaries = []
-    for i, val in enumerate(mask):
-        if in_region is None and val:
-            in_region = i
-        elif in_region is not None and not val:
-            boundaries.append((in_region, i))
-            in_region = None
+    # Add first and/or last index if needed
+    if mask[0]:
+        idx = [0] + idx
+    if mask[-1]:
+        idx.append(len(mask))
 
-    if in_region is not None:
-        boundaries.append((in_region, i+1))
-    return boundaries
+    return list(zip(idx[::2], idx[1::2]))
 
 
 def cross_from_below(x, threshold):

lib/matplotlib/tests/test_mlab.py

@@ -2963,6 +2963,30 @@ def test_evaluate_equal_dim_and_num_lt(self):
 
 
 #*****************************************************************
+
+def test_contiguous_regions():
+    a, b, c = 3, 4, 5
+    # Starts and ends with True
+    mask = [True]*a + [False]*b + [True]*c
+    expected = [(0, a), (a+b, a+b+c)]
+    assert_equal(mlab.contiguous_regions(mask), expected)
+    d, e = 6, 7
+    # Starts with True ends with False
+    mask = mask + [False]*e
+    assert_equal(mlab.contiguous_regions(mask), expected)
+    # Starts with False ends with True
+    mask = [False]*d + mask[:-e]
+    expected = [(d, d+a), (d+a+b, d+a+b+c)]
+    assert_equal(mlab.contiguous_regions(mask), expected)
+    # Starts and ends with False
+    mask = mask + [False]*e
+    assert_equal(mlab.contiguous_regions(mask), expected)
+    # No True in mask
+    assert_equal(mlab.contiguous_regions([False]*5), [])
+    # Empty mask
+    assert_equal(mlab.contiguous_regions([]), [])
+
+
 #*****************************************************************
 
 if __name__ == '__main__':