Merge pull request #121 from neutrinoceros/startup_speedup

yt-project · Jun 19, 2023 · 5d3a1f4 · 5d3a1f4
2 parents a473b6e + 7d5a1ff
commit 5d3a1f4
Show file tree

Hide file tree

Showing 9 changed files with 55 additions and 80 deletions.
diff --git a/cmyt/__init__.py b/cmyt/__init__.py
@@ -1,3 +1,3 @@
 from .cm import *
 
-__version__ = "1.2.0"
+__version__ = "1.2.1"
diff --git a/cmyt/cm.py b/cmyt/cm.py
@@ -1,32 +1,18 @@
 from importlib import import_module
 
-import numpy as np
-
-from cmyt.utils import ColorDict, cmyt_cmaps, register_colormap
-
-
-def _luts_to_cdict(luts: np.ndarray) -> ColorDict:
-    _vs = np.linspace(0, 1, 256)
-
-    return {
-        "red": np.transpose([_vs, luts[0], luts[0]]),
-        "green": np.transpose([_vs, luts[1], luts[1]]),
-        "blue": np.transpose([_vs, luts[2], luts[2]]),
-    }
-
+from cmyt.utils import cmyt_cmaps, register_colormap
 
 for name in cmyt_cmaps:
     # register to MPL
     mod = import_module(f"cmyt.colormaps.{name}")
     if hasattr(mod, "data"):
-        data = mod.data
+        cmap, cmap_r = register_colormap(name, color_dict=mod.data)
     elif hasattr(mod, "luts"):
-        data = _luts_to_cdict(mod.luts)
+        cmap, cmap_r = register_colormap(name, colors=mod.luts)
     else:
         raise RuntimeError(
             f"colormap module '{name}' doesn't contain necessary data for registration."
         )
-    cmap, cmap_r = register_colormap(name, data)
 
     globals()[cmap.name] = cmap
     globals()[cmap_r.name] = cmap_r

diff --git a/cmyt/colormaps/apricity.py b/cmyt/colormaps/apricity.py
@@ -29,8 +29,8 @@
     "spline_method": "CatmulClark",
 }
 
-luts = (
-    np.array(
+luts = np.transpose(
+    (
         [
             0.349019,
             0.349019,
@@ -288,9 +288,7 @@
             0.952941,
             0.956862,
             0.964705,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.411764,
             0.4156864,
@@ -548,9 +546,7 @@
             0.964705,
             0.964705,
             0.964705,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.686274,
             0.682352,
@@ -808,7 +804,6 @@
             0.231372,
             0.223529,
             0.215686,
-        ]
-    ),
-    np.ones(256),
+        ],
+    )
 )
diff --git a/cmyt/colormaps/arbre.py b/cmyt/colormaps/arbre.py
@@ -21,8 +21,8 @@
     "max_Jp": 94.7771696638,
 }
 
-luts = (
-    np.array(
+luts = np.transpose(
+    (
         [
             0.44131774,
             0.44370177,
@@ -280,9 +280,7 @@
             0.97790967,
             0.986926,
             0.99464852,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.05626182,
             0.06063603,
@@ -540,9 +538,7 @@
             0.94575723,
             0.94660423,
             0.94788672,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.04951266,
             0.06083219,
@@ -800,7 +796,6 @@
             0.2341978,
             0.25223614,
             0.2722682,
-        ]
-    ),
-    np.ones(256),
+        ],
+    )
 )
diff --git a/cmyt/colormaps/dusk.py b/cmyt/colormaps/dusk.py
@@ -31,8 +31,8 @@
     "max_Jp": 96.5975103734,
 }
 
-luts = (
-    np.array(
+luts = np.transpose(
+    (
         [
             0.02379297,
             0.0261157,
@@ -290,9 +290,7 @@
             0.9722811,
             0.97360166,
             0.97499094,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.01131879,
             0.01391783,
@@ -550,9 +548,7 @@
             0.93539672,
             0.94119732,
             0.94701577,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.02001135,
             0.02410049,
@@ -810,7 +806,6 @@
             0.92805132,
             0.93660318,
             0.9448656,
-        ]
-    ),
-    np.ones(256),
+        ],
+    )
 )
diff --git a/cmyt/colormaps/kelp.py b/cmyt/colormaps/kelp.py
@@ -21,8 +21,8 @@
     "max_Jp": 95,
 }
 
-luts = (
-    np.array(
+luts = np.transpose(
+    (
         [
             0.07873808,
             0.08503098,
@@ -280,9 +280,7 @@
             0.96451696,
             0.97201416,
             0.97959794,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.02380049,
             0.02762946,
@@ -540,9 +538,7 @@
             0.94151285,
             0.9446259,
             0.94771918,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.45890713,
             0.46137905,
@@ -800,7 +796,6 @@
             0.52734442,
             0.52433316,
             0.52118636,
-        ]
-    ),
-    np.ones(256),
+        ],
+    )
 )
diff --git a/cmyt/colormaps/octarine.py b/cmyt/colormaps/octarine.py
@@ -22,8 +22,8 @@
     "max_Jp": 82.1875,
 }
 
-luts = (
-    np.array(
+luts = np.transpose(
+    (
         [
             0.01845663,
             0.01940818,
@@ -281,9 +281,7 @@
             0.91657895,
             0.92518702,
             0.93347579,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.14549808,
             0.14959758,
@@ -541,9 +539,7 @@
             0.67282641,
             0.67266658,
             0.67286793,
-        ]
-    ),
-    np.array(
+        ],
         [
             0.31784919,
             0.31399639,
@@ -801,7 +797,6 @@
             0.97849108,
             0.98057884,
             0.98147471,
-        ]
-    ),
-    np.ones(256),
+        ],
+    )
 )
diff --git a/cmyt/utils.py b/cmyt/utils.py
@@ -5,7 +5,7 @@
 
 import matplotlib
 import numpy as np
-from matplotlib.colors import Colormap, LinearSegmentedColormap
+from matplotlib.colors import Colormap, LinearSegmentedColormap, ListedColormap
 from more_itertools import always_iterable
 
 # type aliases
@@ -95,20 +95,34 @@ def _register_mpl_cmap(cmap: Colormap) -> None:
 
 def register_colormap(
     name: str,
-    color_dict: ColorDict,
+    *,
+    color_dict: Optional[ColorDict] = None,
+    colors: Optional[np.ndarray] = None,
 ) -> Tuple[LinearSegmentedColormap, LinearSegmentedColormap]:
     name = prefix_name(name)
 
+    if color_dict is not None and colors is not None:
+        raise TypeError("Either color_dict or colors must be provided, but not both")
     # register to MPL
-    mpl_cmap = LinearSegmentedColormap(name=name, segmentdata=color_dict, N=256)
+    if color_dict is not None:
+        mpl_cmap = LinearSegmentedColormap(name=name, segmentdata=color_dict, N=256)
+    elif colors is not None:
+        mpl_cmap = ListedColormap(colors, name=name, N=256)
+    else:
+        raise TypeError("color_dict or colors must be provided")
     mpl_cmap_r = mpl_cmap.reversed()
     _register_mpl_cmap(mpl_cmap)
     _register_mpl_cmap(mpl_cmap_r)
 
     # return cmaps with unprefixed names for registration as importable objects
-    cmap = LinearSegmentedColormap(
-        name=unprefix_name(name), segmentdata=color_dict, N=256
-    )
+    if MPL_VERSION >= (3, 4):
+        cmap = mpl_cmap.copy()
+    else:
+        if color_dict is not None:
+            cmap = LinearSegmentedColormap(name=name, segmentdata=color_dict, N=256)
+        elif colors is not None:
+            cmap = ListedColormap(colors, name=name, N=256)
+    cmap.name = unprefix_name(name)
     cmap_r = cmap.reversed()
 
     return cmap, cmap_r

diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cmyt"
-version = "1.2.0"
+version = "1.2.1"
 description = "A collection of Matplotlib colormaps from the yt project"
 authors = [
     { name = "The yt project", email = "[email protected]" },