From 01d3635764ca03a7e8381f06eae58f28c1b7c076 Mon Sep 17 00:00:00 2001 From: Antony Lee Date: Wed, 13 Jan 2021 14:14:59 +0100 Subject: [PATCH] Speedup LinearSegmentedColormap.from_list. ... 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.) --- lib/matplotlib/colors.py | 31 ++++++++++++++++++++----------- 1 file changed, 20 insertions(+), 11 deletions(-) diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py index 761504e8ca7b..e84c029f5a80 100644 --- a/lib/matplotlib/colors.py +++ b/lib/matplotlib/colors.py @@ -361,11 +361,20 @@ def to_rgba_array(c, alpha=None): if len(c) == 0: return np.zeros((0, 4), float) + + # Quick path if the whole sequence can be directly converted to a numpy + # array in one shot. + 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: - if np.iterable(alpha): - return np.array([to_rgba(cc, aa) for cc, aa in zip(c, alpha)]) - 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 +923,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)