From 24ad119a2e9599e44be911a930b2db262b333ba3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= <cr52@protonmail.com>
Date: Sun, 25 Jun 2023 16:03:20 +0200
Subject: [PATCH] exp: avoid building a colormap registry and copies

---
 yt/visualization/color_maps.py            | 29 ++++++-----------------
 yt/visualization/tests/test_color_maps.py | 11 ++++-----
 2 files changed, 12 insertions(+), 28 deletions(-)

diff --git a/yt/visualization/color_maps.py b/yt/visualization/color_maps.py
index 48edbc9af19..a6dc7054dea 100644
--- a/yt/visualization/color_maps.py
+++ b/yt/visualization/color_maps.py
@@ -11,8 +11,6 @@
 
 from . import _colormap_data as _cm
 
-yt_colormaps = {}
-
 
 def add_colormap(name, cdict):
     """
@@ -25,19 +23,7 @@ def add_colormap(name, cdict):
         since="4.3",
         stacklevel=3,
     )
-    # Note: this function modifies the global variable 'yt_colormaps'
-    yt_colormaps[name] = LinearSegmentedColormap(name, cdict, 256)
-    mpl.colormaps.register(yt_colormaps[name])
-
-
-def _add_colormap(name, colors):
-    """
-    Adds a colormap to the colormaps available in yt for this session
-    """
-    # Note: this function modifies the global variable 'yt_colormaps'
-    cmap = ListedColormap(colors, name=name, N=256)
-    yt_colormaps[cmap.name] = cmap
-    mpl.colormaps.register(cmap)
+    mpl.colormaps.register(LinearSegmentedColormap(name, cdict, 256))
 
 
 # YTEP-0040 backward compatibility layer
@@ -66,20 +52,19 @@ def register_yt_colormaps_from_cmyt():
     """
 
     for hist_name, alias in _HISTORICAL_ALIASES.items():
-        cmap = mpl.colormaps[alias].copy()
-        cmap.name = hist_name
-        mpl.colormaps.register(cmap)
-        mpl.colormaps.register(cmap.reversed())
+        cmap = mpl.colormaps[alias]
+        cmap_r = mpl.colormaps[f"{alias}_r"]
+
+        mpl.colormaps.register(cmap, name=hist_name)
+        mpl.colormaps.register(cmap_r, name=f"{hist_name}_r")
 
 
 register_yt_colormaps_from_cmyt()
 
 # Add colormaps in _colormap_data.py that weren't defined here
 for k, v in _cm.color_map_luts.items():
-    if k in yt_colormaps:
-        continue
     try:
-        _add_colormap(k, v)
+        mpl.colormaps.register(ListedColormap(v, name=k, N=256))
     except ValueError:
         # expected if another map with identical name was already registered
         mylog.warning("cannot register colormap '%s' (naming collision)", k)
diff --git a/yt/visualization/tests/test_color_maps.py b/yt/visualization/tests/test_color_maps.py
index 6c99d9e8e7a..a823fd249f3 100644
--- a/yt/visualization/tests/test_color_maps.py
+++ b/yt/visualization/tests/test_color_maps.py
@@ -61,12 +61,11 @@ def test_make_colormap(self):
 
 
 def test_cmyt_integration():
-    for name in ["algae", "bds_highcontrast", "kelp", "arbre", "octarine", "kamae"]:
-        cmap = plt.get_cmap(name)
-        assert cmap.name == name
-        name_r = name + "_r"
-        cmap_r = plt.get_cmap(name_r)
-        assert cmap_r.name == name_r
+    from yt.visualization.color_maps import _HISTORICAL_ALIASES
+
+    for hist_name, alias in _HISTORICAL_ALIASES.items():
+        assert plt.get_cmap(hist_name) == plt.get_cmap(alias)
+        assert plt.get_cmap(f"{hist_name}_r") == plt.get_cmap(f"{alias}_r")
 
     for name in ["algae", "kelp", "arbre", "octarine", "pastel"]:
         cmap = plt.get_cmap("cmyt." + name)