Merge pull request #1140 from MetOffice/949_populate_colorbar_json_fi…

…le_with_environment_variable

Support flexible pressure_level dependent colourbars

src/CSET/operators/_colorbar_definition.json

@@ -633,6 +633,49 @@
     "max": 1e-06,
     "min": -1e-06
   },
+  "temperature_at_pressure_levels": {
+    "cmap": "RdYlBu_r",
+    "max": 330,
+    "min": 200,
+    "pressure_levels": {
+      "100": {
+        "max": 240,
+        "min": 200
+      },
+      "1000": {
+        "max": 320,
+        "min": 240
+      },
+      "200": {
+        "max": 240,
+        "min": 200
+      },
+      "250": {
+        "max": 240,
+        "min": 200
+      },
+      "300": {
+        "max": 250,
+        "min": 210
+      },
+      "500": {
+        "max": 280,
+        "min": 240
+      },
+      "700": {
+        "max": 300,
+        "min": 240
+      },
+      "850": {
+        "max": 310,
+        "min": 250
+      },
+      "950": {
+        "max": 310,
+        "min": 250
+      }
+    }
+  },
   "temperature_at_screen_level": {
     "cmap": "RdYlBu_r",
     "max": 323,

src/CSET/operators/plot.py

@@ -167,9 +167,13 @@ def _load_colorbar_map(user_colorbar_file: str = None) -> dict:
 
 
 def _colorbar_map_levels(cube: iris.cube.Cube):
-    """Specify the color map and levels.
+    """Get an appropriate colorbar for the given cube.
 
-    For the given variable name, from a colorbar dictionary file.
+    For the given variable the appropriate colorbar is looked up from a
+    combination of the built-in CSET colorbar definitions, and any user supplied
+    definitions. As well as varying on variables, these definitions may also
+    exist for specific pressure levels to account for variables with
+    significantly different ranges at different heights.
 
     Parameters
     ----------
@@ -190,28 +194,60 @@ def _colorbar_map_levels(cube: iris.cube.Cube):
     user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
     colorbar = _load_colorbar_map(user_colorbar_file)
 
-    # First try standard name, then long name, then varname.
-    varnames = list(filter(None, [cube.standard_name, cube.long_name, cube.var_name]))
+    try:
+        # We assume that pressure is a scalar coordinate here.
+        pressure_level_raw = cube.coord("pressure").points[0]
+        # Ensure pressure_level is a string, as it is used as a JSON key.
+        pressure_level = str(int(pressure_level_raw))
+    except iris.exceptions.CoordinateNotFoundError:
+        pressure_level = None
+
+    # First try long name, then standard name, then var name. This order is used
+    # as long name is the one we correct between models, so it most likely to be
+    # consistent.
+    varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
     for varname in varnames:
         # Get the colormap for this variable.
         try:
             cmap = mpl.colormaps[colorbar[varname]["cmap"]]
+            if pressure_level is None:
+                var_colorbar = colorbar[varname]
+            else:
+                # If pressure level is specified for cube use a pressure-level
+                # specific colorbar, if one exists.
+                try:
+                    var_colorbar = colorbar[varname]["pressure_levels"][pressure_level]
+                except KeyError:
+                    logging.warning(
+                        "%s has no colorbar definition for pressure level %s.",
+                        varname,
+                        pressure_level,
+                    )
+                    # Fallback to variable default.
+                    var_colorbar = colorbar[varname]
+
             # Get the colorbar levels for this variable.
             try:
-                levels = colorbar[varname]["levels"]
+                levels = var_colorbar["levels"]
+                # Use discrete bins when levels are specified, rather than a
+                # smooth range.
                 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
+                logging.debug("Using levels for %s colorbar.", varname)
             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")
+                vmin, vmax = var_colorbar["min"], var_colorbar["max"]
+                logging.debug("Using min and max for %s colorbar.", varname)
                 # Calculate levels from range.
-                levels = np.linspace(vmin, vmax, 20)
+                levels = np.linspace(vmin, vmax, 51)
                 norm = None
-                return cmap, levels, norm
+            return cmap, levels, norm
         except KeyError:
+            logging.debug("Cube name %s has no colorbar definition.", varname)
+            # Retry with next name.
             continue
 
     # Default if no varnames match.
+    logging.warning("No colorbar definition exists for %s.", cube.name())
     cmap, levels, norm = mpl.colormaps["viridis"], None, None
     return cmap, levels, norm
 
@@ -676,7 +712,7 @@ def _plot_and_save_histogram_series(
             vmin = 0
             vmax = 400  # Manually set vmin/vmax to override json derived value.
         else:
-            bins = np.linspace(vmin, vmax, 50)
+            bins = np.linspace(vmin, vmax, 51)
 
         # Reshape cube data into a single array to allow for a single histogram.
         # Otherwise we plot xdim histograms stacked.

tests/operators/test_plots.py

@@ -88,7 +88,7 @@ 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 (levels == np.linspace(223, 323, 51)).all()
     assert norm is None
 
 
@@ -108,7 +108,7 @@ def test_colorbar_map_levels_name_fallback(cube, tmp_working_dir):
     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 (levels == np.linspace(223, 323, 51)).all()
     assert norm is None
 
 
@@ -123,6 +123,15 @@ def test_colorbar_map_levels_unknown_variable_fallback(cube, tmp_working_dir):
     assert norm is None
 
 
+def test_colorbar_map_levels_pressure_level(transect_source_cube, tmp_working_dir):
+    """Pressure level specific colorbar definition is picked up."""
+    cube_250hPa = transect_source_cube.extract(iris.Constraint(pressure=250))
+    cmap, levels, norm = plot._colorbar_map_levels(cube_250hPa)
+    assert cmap == mpl.colormaps["RdYlBu_r"]
+    assert (levels == np.linspace(200, 240, 51)).all()
+    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.