Make colormap loading lazier

In most cases, 1 or possibly 2 colormaps are ever used in a plot, so
loading all of them is a waste. Instead just record all the possible
names, and load whichever is requested as necessary.

Author: QuLogic

lib/matplotlib/_cm_bivar.py

@@ -2,9 +2,8 @@
 # date: 2024-05-24
 
 import numpy as np
-from matplotlib.colors import SegmentedBivarColormap
 
-BiPeak = np.array(
+BiPeak = (
     [0.000, 0.674, 0.931, 0.000, 0.680, 0.922, 0.000, 0.685, 0.914, 0.000,
      0.691, 0.906, 0.000, 0.696, 0.898, 0.000, 0.701, 0.890, 0.000, 0.706,
      0.882, 0.000, 0.711, 0.875, 0.000, 0.715, 0.867, 0.000, 0.720, 0.860,
@@ -1273,9 +1272,10 @@
      0.282, 0.946, 0.422, 0.263, 0.943, 0.417, 0.243, 0.939, 0.411, 0.223,
      0.935, 0.405, 0.202, 0.931, 0.399, 0.181, 0.927, 0.393, 0.158, 0.923,
      0.387, 0.134, 0.918, 0.381, 0.107,
-     ]).reshape((65, 65, 3))
+     ],
+    (65, 65, 3))
 
-BiOrangeBlue = np.array(
+BiOrangeBlue = (
     [0.000, 0.000, 0.000, 0.000, 0.062, 0.125, 0.000, 0.125, 0.250, 0.000,
      0.188, 0.375, 0.000, 0.250, 0.500, 0.000, 0.312, 0.625, 0.000, 0.375,
      0.750, 0.000, 0.438, 0.875, 0.000, 0.500, 1.000, 0.125, 0.062, 0.000,
@@ -1301,12 +1301,18 @@
      0.562, 0.125, 1.000, 0.625, 0.250, 1.000, 0.688, 0.375, 1.000, 0.750,
      0.500, 1.000, 0.812, 0.625, 1.000, 0.875, 0.750, 1.000, 0.938, 0.875,
      1.000, 1.000, 1.000,
-     ]).reshape((9, 9, 3))
+     ],
+    (9, 9, 3))
 
 cmaps = {
-    "BiPeak": SegmentedBivarColormap(
-        BiPeak, 256, "square", (.5, .5), name="BiPeak"),
-    "BiOrangeBlue": SegmentedBivarColormap(
-        BiOrangeBlue, 256, "square", (0, 0), name="BiOrangeBlue"),
-    "BiCone": SegmentedBivarColormap(BiPeak, 256, "circle", (.5, .5), name="BiCone"),
+    "BiPeak": (*BiPeak, 256, "square", (.5, .5)),
+    "BiOrangeBlue": (*BiOrangeBlue, 256, "square", (0, 0)),
+    "BiCone": (*BiPeak, 256, "circle", (.5, .5)),
 }
+
+
+def cmap_init(name, cmap_spec):
+    from matplotlib.colors import SegmentedBivarColormap
+    data, shape, *args = cmap_spec
+    return SegmentedBivarColormap(np.array(data).reshape(shape), *args,
+                                  name=name)

lib/matplotlib/_cm_listed.py

@@ -1,5 +1,3 @@
-from .colors import ListedColormap
-
 _magma_data = [[0.001462, 0.000466, 0.013866],
                [0.002258, 0.001295, 0.018331],
                [0.003279, 0.002305, 0.023708],
@@ -2832,16 +2830,15 @@
     [0.74458, 0.99141, 0.64748]]
 
 cmaps = {
-    name: ListedColormap(data, name=name) for name, data in [
-        ('magma', _magma_data),
-        ('inferno', _inferno_data),
-        ('plasma', _plasma_data),
-        ('viridis', _viridis_data),
-        ('cividis', _cividis_data),
-        ('twilight', _twilight_data),
-        ('twilight_shifted', _twilight_shifted_data),
-        ('turbo', _turbo_data),
-        ('berlin', _berlin_data),
-        ('managua', _managua_data),
-        ('vanimo', _vanimo_data),
-    ]}
+    'magma': _magma_data,
+    'inferno': _inferno_data,
+    'plasma': _plasma_data,
+    'viridis': _viridis_data,
+    'cividis': _cividis_data,
+    'twilight': _twilight_data,
+    'twilight_shifted': _twilight_shifted_data,
+    'turbo': _turbo_data,
+    'berlin': _berlin_data,
+    'managua': _managua_data,
+    'vanimo': _vanimo_data,
+}

lib/matplotlib/_cm_multivar.py

@@ -1,7 +1,6 @@
 # auto-generated by https://github.com/trygvrad/multivariate_colormaps
 # date: 2024-05-28
 
-from .colors import LinearSegmentedColormap, MultivarColormap
 import matplotlib as mpl
 _LUTSIZE = mpl.rcParams['image.lut']
 
@@ -145,22 +144,17 @@
                    [0.615, 0.175, 0.341],
                    [0.649, 0.178, 0.355]]
 
-cmaps = {
-    name: LinearSegmentedColormap.from_list(name, data, _LUTSIZE) for name, data in [
-        ('2VarAddA0', _2VarAddA0_data),
-        ('2VarAddA1', _2VarAddA1_data),
-        ('2VarSubA0', _2VarSubA0_data),
-        ('2VarSubA1', _2VarSubA1_data),
-        ('3VarAddA0', _3VarAddA0_data),
-        ('3VarAddA1', _3VarAddA1_data),
-        ('3VarAddA2', _3VarAddA2_data),
-    ]}
-
 cmap_families = {
-    '2VarAddA': MultivarColormap([cmaps[f'2VarAddA{i}'] for i in range(2)],
-                                 'sRGB_add', name='2VarAddA'),
-    '2VarSubA': MultivarColormap([cmaps[f'2VarSubA{i}'] for i in range(2)],
-                                 'sRGB_sub', name='2VarSubA'),
-    '3VarAddA': MultivarColormap([cmaps[f'3VarAddA{i}'] for i in range(3)],
-                                 'sRGB_add', name='3VarAddA'),
+    '2VarAddA': ((_2VarAddA0_data, _2VarAddA1_data), 'sRGB_add'),
+    '2VarSubA': ((_2VarSubA0_data, _2VarSubA1_data), 'sRGB_sub'),
+    '3VarAddA': ((_3VarAddA0_data, _3VarAddA1_data, _3VarAddA2_data), 'sRGB_add'),
 }
+
+
+def cmap_init(name, cmap_spec):
+    from .colors import LinearSegmentedColormap, MultivarColormap
+    specs, combination_mode = cmap_spec
+    return MultivarColormap(
+        [LinearSegmentedColormap.from_list(f'{name}{i}', data, _LUTSIZE)
+         for i, data in enumerate(specs)],
+        combination_mode, name=name)

lib/matplotlib/cm.py

@@ -15,15 +15,17 @@
 """
 
 from collections.abc import Mapping
+import functools
 
 import matplotlib as mpl
 from matplotlib import _api, colors
 # TODO make this warn on access
 from matplotlib.colorizer import _ScalarMappable as ScalarMappable  # noqa
 from matplotlib._cm import datad
 from matplotlib._cm_listed import cmaps as cmaps_listed
-from matplotlib._cm_multivar import cmap_families as multivar_cmaps
-from matplotlib._cm_bivar import cmaps as bivar_cmaps
+from matplotlib._cm_multivar import (
+    cmap_families as multivar_cmaps, cmap_init as _multivar_cmap_init)
+from matplotlib._cm_bivar import cmaps as bivar_cmaps, cmap_init as _bivar_cmap_init
 
 
 _LUTSIZE = mpl.rcParams['image.lut']
@@ -34,15 +36,6 @@ def _gen_cmap_registry():
     Generate a dict mapping standard colormap names to standard colormaps, as
     well as the reversed colormaps.
     """
-    cmap_d = {**cmaps_listed}
-    for name, spec in datad.items():
-        cmap_d[name] = (  # Precache the cmaps at a fixed lutsize..
-            colors.LinearSegmentedColormap(name, spec, _LUTSIZE)
-            if 'red' in spec else
-            colors.ListedColormap(spec['listed'], name)
-            if 'listed' in spec else
-            colors.LinearSegmentedColormap.from_list(name, spec, _LUTSIZE))
-
     # Register colormap aliases for gray and grey.
     aliases = {
         # alias -> original name
@@ -51,16 +44,30 @@ def _gen_cmap_registry():
         'gist_yerg': 'gist_yarg',
         'Grays': 'Greys',
     }
-    for alias, original_name in aliases.items():
-        cmap = cmap_d[original_name].copy()
-        cmap.name = alias
-        cmap_d[alias] = cmap
+    cmap_d = {**cmaps_listed, **datad, **aliases}
+    # Reversed cmaps.
+    for name in list(cmap_d.keys()):
+        cmap_d[f'{name}_r'] = name
+
+    def cmap_init(name, cmap_spec):
+        if name in datad:
+            spec = cmap_spec
+            return (  # Precache the cmaps at a fixed lutsize..
+                colors.LinearSegmentedColormap(name, spec, _LUTSIZE)
+                if 'red' in spec else
+                colors.ListedColormap(spec['listed'], name)
+                if 'listed' in spec else
+                colors.LinearSegmentedColormap.from_list(name, spec, _LUTSIZE))
+        if name in aliases:
+            cmap = cmap_init(cmap_spec, cmap_d[cmap_spec]).copy()
+            cmap.name = name
+            return cmap
+        if name.endswith('_r'):
+            # Generate reversed cmaps.
+            return cmap_init(cmap_spec, cmap_d[cmap_spec]).reversed()
+        return colors.ListedColormap(cmap_spec, name=name)
 
-    # Generate reversed cmaps.
-    for cmap in list(cmap_d.values()):
-        rmap = cmap.reversed()
-        cmap_d[rmap.name] = rmap
-    return cmap_d
+    return cmap_d, cmap_init
 
 
 class ColormapRegistry(Mapping):
@@ -87,13 +94,18 @@ class ColormapRegistry(Mapping):
         from matplotlib import colormaps
         list(colormaps)
     """
-    def __init__(self, cmaps):
-        self._cmaps = cmaps
-        self._builtin_cmaps = tuple(cmaps)
+    def __init__(self, cmaps, init_func):
+        self._cmaps = {name: None for name in cmaps}
+        self._init_func = init_func
+        self._builtin_cmaps = cmaps
 
     def __getitem__(self, item):
         try:
-            return self._cmaps[item].copy()
+            cmap = self._cmaps[item]
+            if cmap is None:
+                cmap = self._cmaps[item] = self._init_func(item,
+                                                           self._builtin_cmaps[item])
+            return cmap.copy()
         except KeyError:
             raise KeyError(f"{item!r} is not a known colormap name") from None
 
@@ -235,12 +247,11 @@ def get_cmap(self, cmap):
 # public access to the colormaps should be via `matplotlib.colormaps`. For now,
 # we still create the registry here, but that should stay an implementation
 # detail.
-_colormaps = ColormapRegistry(_gen_cmap_registry())
-globals().update(_colormaps)
+_colormaps = ColormapRegistry(*_gen_cmap_registry())
 
-_multivar_colormaps = ColormapRegistry(multivar_cmaps)
+_multivar_colormaps = ColormapRegistry(multivar_cmaps, _multivar_cmap_init)
 
-_bivar_colormaps = ColormapRegistry(bivar_cmaps)
+_bivar_colormaps = ColormapRegistry(bivar_cmaps, _bivar_cmap_init)
 
 
 # This is an exact copy of pyplot.get_cmap(). It was removed in 3.9, but apparently
@@ -307,3 +318,10 @@ def _ensure_cmap(cmap):
     if cmap_name not in _colormaps:
         _api.check_in_list(sorted(_colormaps), cmap=cmap_name)
     return mpl.colormaps[cmap_name]
+
+
+@functools.cache
+def __getattr__(name):
+    if name in _colormaps._builtin_cmaps:
+        return _colormaps[name]
+    raise AttributeError(f"module 'matplotlib.cm' has no attribute {name!r}")