Speedup LinearSegmentedColormap.from_list.

anntzer · anntzer · commit 29dc7eb8b018 · 2021-01-13T14:16:50.000+01:00
... by special-casing lists-of-rgb(a) tuples in to_rgba_array.

This speeds up colorcet's import nearly twofold, from ~430ms to ~250ms
on my machine.  (colorcet generates a lot of colormaps at import time.)
diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
@@ -362,10 +362,16 @@ def to_rgba_array(c, alpha=None):
     if len(c) == 0:
         return np.zeros((0, 4), float)
     else:
-        if np.iterable(alpha):
-            return np.array([to_rgba(cc, aa) for cc, aa in zip(c, alpha)])
+        lens = {len(cc) if isinstance(cc, (list, tuple)) else -1 for cc in c}
+        if lens == {3}:
+            rgba = np.column_stack([c, np.ones(len(c))])
+        elif lens == {4}:
+            rgba = np.array(c)
         else:
-            return np.array([to_rgba(cc, alpha) for cc in c])
+            rgba = np.array([to_rgba(cc) for cc in c])
+        if alpha is not None:
+            rgba[:, 3] = alpha
+        return rgba
 
 
 def to_rgb(c):
@@ -914,13 +920,13 @@ def from_list(name, colors, N=256, gamma=1.0):
         else:
             vals = np.linspace(0, 1, len(colors))
 
-        cdict = dict(red=[], green=[], blue=[], alpha=[])
-        for val, color in zip(vals, colors):
-            r, g, b, a = to_rgba(color)
-            cdict['red'].append((val, r, r))
-            cdict['green'].append((val, g, g))
-            cdict['blue'].append((val, b, b))
-            cdict['alpha'].append((val, a, a))
+        r, g, b, a = to_rgba_array(colors).T
+        cdict = {
+            "red": np.column_stack([vals, r, r]),
+            "green": np.column_stack([vals, g, g]),
+            "blue": np.column_stack([vals, b, b]),
+            "alpha": np.column_stack([vals, a, a]),
+        }
 
         return LinearSegmentedColormap(name, cdict, N, gamma)
 
