From 2827015e7f7299009aa4b22d953cd90f3e0aeddd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Robert?= Date: Sun, 24 Nov 2024 15:53:30 +0100 Subject: [PATCH] ENH: (CLI): support colormap extension packages other than cblind (cmasher, cmocean, cmyt) --- nonos/main.py | 61 +++++++++++++++++++---------- requirements/tests_all.txt | 3 ++ tests/test_plotting.py | 78 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 123 insertions(+), 19 deletions(-) diff --git a/nonos/main.py b/nonos/main.py index 610deecc..5cb03543 100644 --- a/nonos/main.py +++ b/nonos/main.py @@ -11,12 +11,13 @@ import re import time from collections import ChainMap +from importlib import import_module from importlib.metadata import version +from importlib.util import find_spec from multiprocessing import Pool from pathlib import Path from typing import TYPE_CHECKING, Any, Optional -import cblind # noqa import inifix import numpy as np from packaging.version import Version @@ -50,6 +51,13 @@ NONOS_VERSION = version("nonos") INIFIX_GE_5_0 = Version(version("inifix")) >= Version("5.0.0") +KNOWN_CMAP_PACKAGE_PREFIXES = { + "cb": "cblind", + "cmo": "cmocean", + "cmr": "cmasher", + "cmyt": "cmyt", +} + def get_non_interactive_figure(fmt: str) -> "Figure": from matplotlib.backends.backend_agg import FigureCanvasAgg @@ -91,7 +99,7 @@ def process_field( vmin, vmax, scaling: float, - cmap, + cmap: str, title, unit_conversion: int, datadir, @@ -155,17 +163,39 @@ def process_field( else: fig = get_non_interactive_figure(fmt) + plot_kwargs = { + "log": log, + "vmin": vmin, + "vmax": vmax, + "cmap": cmap, + "title": f"${title}$", + "unit_conversion": unit_conversion, + } + + if cm_prefix := cmap.rpartition(".")[0]: + if (cm_package := KNOWN_CMAP_PACKAGE_PREFIXES.get(cm_prefix)) is None: + print_warn( + f"requested colormap {cmap!r} with the unknown prefix {cm_prefix!r}. " + "The default colormap will be used instead." + ) + plot_kwargs.pop("cmap") + elif not find_spec(cm_package): + print_warn( + f"requested colormap {cmap!r}, but {cm_package} is not installed. " + "The default colormap will be used instead." + ) + plot_kwargs.pop("cmap") + else: + # all known colormap packages work by registering their colormaps + # at import time. + import_module(cm_package) + ax = fig.add_subplot(111, polar=False) if dim == 1: - dsop.map(plane[0], rotate_with=planet_file).plot( - fig, - ax, - log=log, - vmin=vmin, - vmax=vmax, - title=f"${title}$", - unit_conversion=unit_conversion, - ) + if "cmap" in plot_kwargs: + plot_kwargs.pop("cmap") + + dsop.map(plane[0], rotate_with=planet_file).plot(fig, ax, **plot_kwargs) akey = dsop.map(plane[0], rotate_with=planet_file).dict_plotable["abscissa"] avalue = dsop.map(plane[0], rotate_with=planet_file).dict_plotable[akey] extent = parse_range(extent, dim=dim) @@ -173,14 +203,7 @@ def process_field( ax.set_xlim(extent[0], extent[1]) elif dim == 2: dsop.map(plane[0], plane[1], rotate_with=planet_file).plot( - fig, - ax, - log=log, - vmin=vmin, - vmax=vmax, - cmap=cmap, - title=f"${title}$", - unit_conversion=unit_conversion, + fig, ax, **plot_kwargs ) akey = dsop.map(plane[0], plane[1], rotate_with=planet_file).dict_plotable[ "abscissa" diff --git a/requirements/tests_all.txt b/requirements/tests_all.txt index af480764..b1db7509 100644 --- a/requirements/tests_all.txt +++ b/requirements/tests_all.txt @@ -1,2 +1,5 @@ -r tests_min.txt +cblind>=2.3.0 +cmocean>=3.0.3 +cmyt>=2.0.0 numexpr>=2.8.3 diff --git a/tests/test_plotting.py b/tests/test_plotting.py index d65eec3d..f2b2b089 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -1,4 +1,5 @@ import os +from importlib.util import find_spec import numpy.testing as npt import pytest @@ -199,3 +200,80 @@ def test_reg(test_data_dir, map_args): fig = Figure() ax = fig.add_subplot() ds["RHO"].map(*map_args).plot(fig, ax) + + +@pytest.mark.parametrize( + "cmap", + [ + pytest.param( + cmap_name, + marks=pytest.mark.skipif( + not find_spec(pkg_name), reason=f"{pkg_name} is not installed" + ), + ) + for pkg_name, cmap_name in [ + ("cblind", "cb.rainbow"), + ("cmocean", "cmo.thermal"), + ("cmasher", "cmr.dusk"), + ("cmyt", "cmyt.arbre"), + ] + ], +) +def test_colormap_extensions_integration(cmap, capsys, test_data_dir, tmp_path): + simdir = str(test_data_dir / "idefix_planet3d") + os.chdir(tmp_path) + ret = main(["-dir", simdir, "-geometry", "polar", "-cmap", cmap]) + assert ret == 0 + out, err = capsys.readouterr() + assert out == "" + assert err == "" + + +@pytest.mark.parametrize( + "cmap, pkg", + [ + pytest.param( + cmap_name, + pkg_name, + marks=pytest.mark.skipif( + find_spec(pkg_name), reason=f"{pkg_name} is installed" + ), + ) + for pkg_name, cmap_name in [ + ("cblind", "cb.rainbow"), + ("cmocean", "cmo.thermal"), + ("cmasher", "cmr.dusk"), + ("cmyt", "cmyt.arbre"), + ] + ], +) +def test_colormap_extensions_missing_package( + cmap, pkg, capsys, test_data_dir, tmp_path +): + simdir = str(test_data_dir / "idefix_planet3d") + os.chdir(tmp_path) + ret = main(["-dir", simdir, "-geometry", "polar", "-cmap", cmap]) + assert ret == 0 + out, err = capsys.readouterr() + assert out == "" + assert err.replace("\n", "") == ( + f"🦴 Warning requested colormap {cmap!r}, but {pkg} is not installed. " + "The default colormap will be used instead." + ) + + +def test_unknown_colormap_package_prefix(capsys, test_data_dir, tmp_path): + simdir = str(test_data_dir / "idefix_planet3d") + os.chdir(tmp_path) + ret = main( + ["-dir", simdir, "-geometry", "polar", "-cmap", "cmunknown.thismapdoesnexist"] + ) + assert ret == 0 + + out, err = capsys.readouterr() + assert out == "" + assert err.replace("\n", "") == ( + "🦴 Warning requested colormap 'cmunknown.thismapdoesnexist' " + "with the unknown prefix 'cmunknown'. " + "The default colormap will be used instead." + )