Merge pull request #1141 from MetOffice/colourbar_name_fallback

Search for all varnames and allow user colorbar override

src/CSET/operators/plot.py

@@ -33,7 +33,13 @@
 import numpy as np
 from markdown_it import MarkdownIt
 
-from CSET._common import get_recipe_metadata, iter_maybe, render_file, slugify
+from CSET._common import (
+    combine_dicts,
+    get_recipe_metadata,
+    iter_maybe,
+    render_file,
+    slugify,
+)
 from CSET.operators._utils import get_cube_yxcoordname, is_transect
 
 # Use a non-interactive plotting backend.
@@ -135,37 +141,40 @@ def _make_plot_html_page(plots: list):
 
 
 @functools.cache
-def _load_colorbar_map() -> dict:
+def _load_colorbar_map(user_colorbar_file: str = None) -> dict:
     """Load the colorbar definitions from a file.
 
     This is a separate function to make it cacheable.
     """
-    # Grab the colorbar file from the recipe global metadata.
-    try:
-        colorbar_file = get_recipe_metadata()["style_file_path"]
-        logging.debug("Colour bar file: %s", colorbar_file)
-        with open(colorbar_file, "rt", encoding="UTF-8") as fp:
-            colorbar = json.load(fp)
-    except (FileNotFoundError, KeyError):
-        logging.info("Colorbar file does not exist. Using default values.")
-        operator_files = _py312_importlib_resources_files_shim()
-        colorbar_def_file = operator_files.joinpath("_colorbar_definition.json")
-        with open(colorbar_def_file, "rt", encoding="UTF-8") as fp:
-            colorbar = json.load(fp)
+    colorbar_file = _py312_importlib_resources_files_shim().joinpath(
+        "_colorbar_definition.json"
+    )
+    with open(colorbar_file, "rt", encoding="UTF-8") as fp:
+        colorbar = json.load(fp)
+
+    logging.debug("User colour bar file: %s", user_colorbar_file)
+    override_colorbar = {}
+    if user_colorbar_file:
+        try:
+            with open(user_colorbar_file, "rt", encoding="UTF-8") as fp:
+                override_colorbar = json.load(fp)
+        except FileNotFoundError:
+            logging.warning("Colorbar file does not exist. Using default values.")
+
+    # Overwrite values with the user supplied colorbar definition.
+    colorbar = combine_dicts(colorbar, override_colorbar)
     return colorbar
 
 
-def _colorbar_map_levels(varname: str, **kwargs):
+def _colorbar_map_levels(cube: iris.cube.Cube):
     """Specify the color map and levels.
 
     For the given variable name, from a colorbar dictionary file.
 
     Parameters
     ----------
-    colorbar_file: str
-        Filename of the colorbar dictionary to read.
-    varname: str
-        Variable name to extract from the dictionary
+    cube: Cube
+        Cube of variable for which the colorbar information is desired.
 
     Returns
     -------
@@ -177,33 +186,33 @@ def _colorbar_map_levels(varname: str, **kwargs):
     norm:
         BoundaryNorm information.
     """
-    colorbar = _load_colorbar_map()
-
-    # Get the colormap for this variable.
-    try:
-        cmap = colorbar[varname]["cmap"]
-        logging.debug("From colorbar dictionary: Using cmap")
-    except KeyError:
-        cmap = mpl.colormaps["viridis"]
+    # Grab the colorbar file from the recipe global metadata.
+    user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
+    colorbar = _load_colorbar_map(user_colorbar_file)
 
-    # Get the colorbar levels for this variable.
-    try:
-        levels = colorbar[varname]["levels"]
-        actual_cmap = mpl.cm.get_cmap(cmap)
-        norm = mpl.colors.BoundaryNorm(levels, ncolors=actual_cmap.N)
-        logging.debug("From colorbar dictionary: Using levels")
-    except KeyError:
+    # First try standard name, then long name, then varname.
+    varnames = list(filter(None, [cube.standard_name, cube.long_name, cube.var_name]))
+    for varname in varnames:
+        # Get the colormap for this variable.
         try:
-            # Get the range for this variable.
-            vmin, vmax = colorbar[varname]["min"], colorbar[varname]["max"]
-            logging.debug("From colorbar dictionary: Using min and max")
-            # Calculate levels from range.
-            levels = np.linspace(vmin, vmax, 20)
-            norm = None
+            cmap = mpl.colormaps[colorbar[varname]["cmap"]]
+            # Get the colorbar levels for this variable.
+            try:
+                levels = colorbar[varname]["levels"]
+                norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
+            except KeyError:
+                # Get the range for this variable.
+                vmin, vmax = colorbar[varname]["min"], colorbar[varname]["max"]
+                logging.debug("From colorbar dictionary: Using min and max")
+                # Calculate levels from range.
+                levels = np.linspace(vmin, vmax, 20)
+                norm = None
+                return cmap, levels, norm
         except KeyError:
-            levels = None
-            norm = None
+            continue
 
+    # Default if no varnames match.
+    cmap, levels, norm = mpl.colormaps["viridis"], None, None
     return cmap, levels, norm
 
 
@@ -236,7 +245,7 @@ def _plot_and_save_spatial_plot(
     fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
 
     # Specify the color bar
-    cmap, levels, norm = _colorbar_map_levels(cube.name())
+    cmap, levels, norm = _colorbar_map_levels(cube)
 
     if method == "contourf":
         # Filled contour plot of the field.
@@ -354,7 +363,7 @@ def _plot_and_save_postage_stamp_spatial_plot(
     fig = plt.figure(figsize=(10, 10))
 
     # Specify the color bar
-    cmap, levels, norm = _colorbar_map_levels(cube.name())
+    cmap, levels, norm = _colorbar_map_levels(cube)
 
     # Make a subplot for each member.
     for member, subplot in zip(
@@ -640,9 +649,9 @@ def _plot_and_save_histogram_series(
     title: str
         Plot title.
     vmin: float
-        minimum for colourbar
+        minimum for colorbar
     vmax: float
-        maximum for colourbar
+        maximum for colorbar
     histtype: str
         The type of histogram to plot. Options are "step" for a line
         histogram or "barstacked", "stepfilled". "Step" is the default option,

tests/operators/test_plots.py

@@ -14,9 +14,12 @@
 
 """Test plotting operators."""
 
+import json
 from pathlib import Path
 
 import iris.cube
+import matplotlib as mpl
+import numpy as np
 import pytest
 
 from CSET.operators import collapse, plot, read
@@ -36,6 +39,90 @@ def test_check_single_cube():
         plot._check_single_cube(non_cube)
 
 
+def test_load_colorbar_map():
+    """Colorbar is loaded correctly."""
+    colorbar = plot._load_colorbar_map()
+    assert isinstance(colorbar, dict)
+    # Check we can find an example definition.
+    assert colorbar["temperature_at_screen_level"] == {
+        "cmap": "RdYlBu_r",
+        "max": 323,
+        "min": 223,
+    }
+
+
+def test_load_colorbar_map_override(tmp_path):
+    """Colorbar is loaded correctly and overridden by the user definition."""
+    # Setup a user provided colorbar override.
+    user_definition = {"temperature_at_screen_level": {"max": 1000, "min": 0}}
+    user_colorbar_file = tmp_path / "colorbar.json"
+    with open(user_colorbar_file, "wt") as fp:
+        json.dump(user_definition, fp)
+
+    colorbar = plot._load_colorbar_map(user_colorbar_file)
+
+    assert isinstance(colorbar, dict)
+    # Check definition is updated.
+    assert colorbar["temperature_at_screen_level"] == {
+        "cmap": "RdYlBu_r",
+        "max": 1000,
+        "min": 0,
+    }
+    # Check we can still see unchanged definitions.
+    assert colorbar["temperature_at_screen_level_difference"] == {
+        "cmap": "bwr",
+        "max": 2,
+        "min": -2,
+    }
+
+
+def test_load_colorbar_map_override_file_not_found(tmp_path):
+    """Colorbar overridden by the user definition in non-existent file."""
+    user_colorbar_file = tmp_path / "colorbar.json"
+    colorbar = plot._load_colorbar_map(user_colorbar_file)
+    # Check it still returns the built-in one.
+    assert isinstance(colorbar, dict)
+
+
+def test_colorbar_map_levels(cube, tmp_working_dir):
+    """Colorbar definition is found for cube."""
+    cmap, levels, norm = plot._colorbar_map_levels(cube)
+    assert cmap == mpl.colormaps["RdYlBu_r"]
+    assert (levels == np.linspace(223, 323, 20)).all()
+    assert norm is None
+
+
+# Test will fail (but not cause an overall fail) as we can't test using levels
+# until there is at least one variable that uses them.
+@pytest.mark.xfail(reason="No colorbar currently uses levels")
+def test_colorbar_map_levels_def_on_levels(cube, tmp_working_dir):
+    """Colorbar definition that uses levels is found for cube."""
+    cmap, levels, norm = plot._colorbar_map_levels(cube)
+    assert cmap == mpl.colormaps[...]
+    assert levels == [1, 2, 3]
+    assert norm == ...
+
+
+def test_colorbar_map_levels_name_fallback(cube, tmp_working_dir):
+    """Colorbar definition is found for cube after checking its other names."""
+    cube.standard_name = None
+    cmap, levels, norm = plot._colorbar_map_levels(cube)
+    assert cmap == mpl.colormaps["RdYlBu_r"]
+    assert (levels == np.linspace(223, 323, 20)).all()
+    assert norm is None
+
+
+def test_colorbar_map_levels_unknown_variable_fallback(cube, tmp_working_dir):
+    """Colorbar definition doesn't exist for cube."""
+    cube.standard_name = None
+    cube.long_name = None
+    cube.var_name = "unknown"
+    cmap, levels, norm = plot._colorbar_map_levels(cube)
+    assert cmap == mpl.colormaps["viridis"]
+    assert levels is None
+    assert norm is None
+
+
 def test_spatial_contour_plot(cube, tmp_working_dir):
     """Plot spatial contour plot of instant air temp."""
     # Remove realization coord to increase coverage, and as its not needed.