-
Notifications
You must be signed in to change notification settings - Fork 18
Lazy loading colormaps #343
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 11 commits
fbf976f
f996653
e550e39
b16a1e1
ab2ddef
6e3278d
9c7027d
f2dae18
5c74d25
cab7433
95dc502
6fbf452
6a7537e
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3090,13 +3090,10 @@ def __init__(self, kwargs): | |
| kwargs : dict-like | ||
| The source dictionary. | ||
| """ | ||
| super().__init__(kwargs) | ||
| # The colormap is initialized with all the base colormaps | ||
| # We have to change the classes internally to Perceptual, Continuous or Discrete | ||
| # such that ultraplot knows what these objects are. We piggy back on the registering mechanism | ||
| # by overriding matplotlib's behavior | ||
| for name in tuple(self._cmaps.keys()): | ||
| self.register(self._cmaps[name], name=name) | ||
| super().__init__({k.lower(): v for k, v in kwargs.items()}) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Why switch to
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We perform a lower command internally later when the colormap is registered. |
||
| # The colormap is initialized with all the base colormaps. | ||
| # These are converted to ultraplot's own colormap objects | ||
| # on the fly when they are first accessed. | ||
|
|
||
| def _translate_deprecated(self, key): | ||
| """ | ||
|
|
@@ -3188,9 +3185,42 @@ def __getitem__(self, key): | |
|
|
||
| if reverse: | ||
| key = key.removesuffix("_r") | ||
|
|
||
| # Retrieve colormap | ||
| if self._has_item(key): | ||
| value = self._cmaps[key].copy() | ||
| value = self._cmaps[key] | ||
|
|
||
| # Lazy loading from file | ||
| if isinstance(value, dict) and value.get("is_lazy"): | ||
| path = value["path"] | ||
| type = value["type"] | ||
| is_default = value.get("is_default", False) | ||
| if type == "continuous": | ||
| cmap = ContinuousColormap.from_file(path, warn_on_failure=True) | ||
| elif type == "discrete": | ||
| cmap = DiscreteColormap.from_file(path, warn_on_failure=True) | ||
| else: | ||
| raise ValueError( | ||
| f"Invalid colormap type {type!r} for key {key!r} in file {path!r}. " | ||
| "Expected 'continuous' or 'discrete'." | ||
| ) | ||
|
|
||
| if cmap: | ||
| if is_default and cmap.name.lower() in CMAPS_CYCLIC: | ||
| cmap.set_cyclic(True) | ||
| self._cmaps[key] = cmap | ||
| value = cmap | ||
| else: # failed to load | ||
| # remove from registry to avoid trying again | ||
| del self._cmaps[key] | ||
| raise KeyError(f"Failed to load colormap {key!r} from {path!r}") | ||
|
||
|
|
||
| # Lazy loading for builtin matplotlib cmaps | ||
| if not isinstance(value, (ContinuousColormap, DiscreteColormap)): | ||
| value = _translate_cmap(value) | ||
| self._cmaps[key] = value | ||
|
|
||
| value = value.copy() | ||
| else: | ||
| raise KeyError( | ||
| f"Invalid colormap or color cycle name {key!r}. Options are: " | ||
|
|
@@ -3220,7 +3250,19 @@ def register(self, cmap, *, name=None, force=False): | |
| # surpress warning if the colormap is not generate by ultraplot | ||
| if name not in self._builtin_cmaps: | ||
| print(f"Overwriting {name!r} that was already registered") | ||
| self._cmaps[name] = cmap.copy() | ||
| self._cmaps[name] = cmap.copy(name=name) | ||
|
|
||
| def register_lazy(self, name, path, type, is_default=False): | ||
| """ | ||
| Register a colormap to be loaded lazily from a file. | ||
| """ | ||
| name = self._translate_key(name, mirror=False) | ||
| self._cmaps[name] = { | ||
| "path": path, | ||
| "type": type, | ||
| "is_default": is_default, | ||
| "is_lazy": True, | ||
| } | ||
|
|
||
|
|
||
| # Initialize databases | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,129 @@ | ||
| import os | ||
| import pytest | ||
| import numpy as np | ||
| import matplotlib.colors as mcolors | ||
|
|
||
| from ultraplot import colors as pcolors | ||
| from ultraplot import config | ||
|
|
||
|
|
||
| @pytest.fixture(autouse=True) | ||
| def setup_teardown(): | ||
| """ | ||
| Reset the colormap database before and after each test. | ||
| """ | ||
| # This ensures a clean state for each test. | ||
| # The singleton instance is replaced with a new one. | ||
| pcolors._cmap_database = pcolors._init_cmap_database() | ||
| config.register_cmaps(default=True) | ||
| config.register_cycles(default=True) | ||
| yield | ||
|
|
||
|
|
||
| def test_lazy_loading_builtin(): | ||
| """ | ||
| Test that built-in colormaps are lazy-loaded. | ||
| """ | ||
| # Before access, it should be a matplotlib colormap | ||
| cmap_raw = pcolors._cmap_database._cmaps["viridis"] | ||
| assert isinstance( | ||
| cmap_raw, | ||
| ( | ||
| pcolors.ContinuousColormap, | ||
| pcolors.DiscreteColormap, | ||
| mcolors.ListedColormap, | ||
| ), | ||
| ) | ||
|
|
||
| # After access, it should be an ultraplot colormap | ||
| cmap_get = pcolors._cmap_database.get_cmap("viridis") | ||
| assert isinstance(cmap_get, pcolors.ContinuousColormap) | ||
|
|
||
| # The internal representation should also be updated | ||
| cmap_raw_after = pcolors._cmap_database._cmaps["viridis"] | ||
| assert isinstance(cmap_raw_after, pcolors.ContinuousColormap) | ||
|
|
||
|
|
||
| def test_case_insensitivity(): | ||
| """ | ||
| Test that colormap lookup is case-insensitive. | ||
| """ | ||
| cmap1 = pcolors._cmap_database.get_cmap("ViRiDiS") | ||
| cmap2 = pcolors._cmap_database.get_cmap("viridis") | ||
| assert cmap1.name.lower().startswith("_viridis") | ||
| assert cmap2.name.lower().startswith("_viridis") | ||
|
|
||
|
|
||
| def test_reversed_shifted(): | ||
| """ | ||
| Test reversed and shifted colormaps. | ||
| """ | ||
| # Create a simple colormap to test the reversal logic | ||
| # This avoids dependency on the exact definition of 'viridis' in matplotlib | ||
| colors_list = [(1, 0, 0), (0, 1, 0), (0, 0, 1)] # Red, Green, Blue | ||
| test_cmap = pcolors.ContinuousColormap.from_list("test_cmap", colors_list) | ||
| pcolors._cmap_database.register(test_cmap) | ||
|
|
||
| cmap = pcolors._cmap_database.get_cmap("test_cmap") | ||
| cmap_r = pcolors._cmap_database.get_cmap("test_cmap_r") | ||
|
|
||
| # Check name | ||
| assert cmap_r.name == "_test_cmap_copy_r" | ||
| # Check colors | ||
| # Start of original should be end of reversed | ||
| assert np.allclose(cmap(0.0), cmap_r(1.0)) | ||
| # End of original should be start of reversed | ||
| assert np.allclose(cmap(1.0), cmap_r(0.0)) | ||
| # Middle should be the same | ||
| assert np.allclose(cmap(0.5)[:3], cmap_r(0.5)[:3][::-1]) | ||
|
|
||
|
|
||
| def test_grays_translation(): | ||
| """ | ||
| Test that 'Grays' is translated to 'greys'. | ||
| """ | ||
| cmap_grays = pcolors._cmap_database.get_cmap("Grays") | ||
| assert cmap_grays.name.lower().startswith("_greys") | ||
|
|
||
|
|
||
| def test_lazy_loading_file(tmp_path): | ||
| """ | ||
| Test that colormaps from files are lazy-loaded. | ||
| """ | ||
| # Create a dummy colormap file | ||
| cmap_data = "1, 0, 0\n0, 1, 0\n0, 0, 1" | ||
| cmap_file = tmp_path / "my_test_cmap.rgb" | ||
| cmap_file.write_text(cmap_data) | ||
|
|
||
| # Register it lazily | ||
| pcolors._cmap_database.register_lazy("my_test_cmap", str(cmap_file), "continuous") | ||
|
|
||
| # Before access, it should be a lazy-load dict | ||
| cmap_raw = pcolors._cmap_database._cmaps["my_test_cmap"] | ||
| assert isinstance(cmap_raw, dict) | ||
| assert cmap_raw["is_lazy"] | ||
|
|
||
| # After access, it should be an ultraplot colormap | ||
| cmap_get = pcolors._cmap_database.get_cmap("my_test_cmap") | ||
| assert isinstance(cmap_get, pcolors.ContinuousColormap) | ||
| assert cmap_get.name.lower().startswith("_my_test_cmap") | ||
|
|
||
| # The internal representation should also be updated | ||
| cmap_raw_after = pcolors._cmap_database._cmaps["my_test_cmap"] | ||
| assert isinstance(cmap_raw_after, pcolors.ContinuousColormap) | ||
|
|
||
|
|
||
| def test_register_new(): | ||
| """ | ||
| Test registering a new colormap. | ||
| """ | ||
| colors_list = [(0, 0, 0), (1, 1, 1)] | ||
| new_cmap = pcolors.DiscreteColormap(colors_list, name="my_new_cmap") | ||
| pcolors._cmap_database.register(new_cmap) | ||
|
|
||
| # Check it was registered | ||
| cmap_get = pcolors._cmap_database.get_cmap("my_new_cmap") | ||
| assert cmap_get.name.lower().startswith( | ||
| "_my_new_cmap" | ||
| ), f"Received {cmap_get.name.lower()} expected _my_new_cmap" | ||
| assert len(cmap_get.colors) == 2 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I guess we are no longer using the
registermethod?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah good catch. We do need it as we override the
ColormapRegistery. Moved the lazy loading from file to a private function to make this distinction clearer and still rely onregister