Coverage for src / CSET / operators / _plot_colormaps.py: 95%
204 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-08 17:20 +0000
1# © Crown copyright, Met Office (2022-2025) and CSET contributors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
15"""Functions to support colormap settings for CSET plots."""
17import functools
18import importlib.resources
19import itertools
20import json
21import logging
22from typing import Literal
24import iris
25import matplotlib as mpl
26import matplotlib.colors as mcolors
27import matplotlib.pyplot as plt
28import numpy as np
30from CSET._common import (
31 combine_dicts,
32 get_recipe_metadata,
33 iter_maybe,
34)
36DEFAULT_DISCRETE_COLORS = mpl.colormaps["tab10"].colors + mpl.colormaps["Accent"].colors
39@functools.cache
40def _load_colorbar_map(user_colorbar_file: str = None) -> dict:
41 """Load the colorbar definitions from a file.
43 This is a separate function to make it cacheable.
44 """
45 colorbar_file = importlib.resources.files().joinpath("_colorbar_definition.json")
46 with open(colorbar_file, "rt", encoding="UTF-8") as fp:
47 colorbar = json.load(fp)
49 logging.debug("User colour bar file: %s", user_colorbar_file)
50 override_colorbar = {}
51 if user_colorbar_file:
52 try:
53 with open(user_colorbar_file, "rt", encoding="UTF-8") as fp:
54 override_colorbar = json.load(fp)
55 except FileNotFoundError:
56 logging.warning("Colorbar file does not exist. Using default values.")
58 # Overwrite values with the user supplied colorbar definition.
59 colorbar = combine_dicts(colorbar, override_colorbar)
60 return colorbar
63def _get_model_colors_map(cubes: iris.cube.CubeList | iris.cube.Cube) -> dict:
64 """Get an appropriate colors for model lines in line plots.
66 For each model in the list of cubes colors either from user provided
67 color definition file (so-called style file) or from default colors are mapped
68 to model_name attribute.
70 Parameters
71 ----------
72 cubes: CubeList or Cube
73 Cubes with model_name attribute
75 Returns
76 -------
77 model_colors_map:
78 Dictionary mapping model_name attribute to colors
79 """
80 user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
81 colorbar = _load_colorbar_map(user_colorbar_file)
82 model_names = sorted(
83 filter(
84 lambda x: x is not None,
85 (cube.attributes.get("model_name", None) for cube in iter_maybe(cubes)),
86 )
87 )
88 if not model_names:
89 return {}
90 use_user_colors = all(mname in colorbar.keys() for mname in model_names)
91 if use_user_colors: 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true
92 return {mname: colorbar[mname] for mname in model_names}
94 # Plot observations as first item
95 if "OBS" in [name.upper() for name in model_names]: 95 ↛ 96line 95 didn't jump to line 96 because the condition on line 95 was never true
96 colors = list(DEFAULT_DISCRETE_COLORS).copy()
97 colors.insert(0, mcolors.to_rgb("dimgray"))
98 ob_name = [name for name in model_names if "OBS" in name.upper()][0]
99 model_names.remove(ob_name)
100 model_names.insert(0, ob_name)
101 else:
102 colors = DEFAULT_DISCRETE_COLORS
104 color_list = itertools.cycle(colors)
105 return {mname: color for mname, color in zip(model_names, color_list, strict=False)}
108def _colorbar_map_levels(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
109 """Get an appropriate colorbar for the given cube.
111 For the given variable the appropriate colorbar is looked up from a
112 combination of the built-in CSET colorbar definitions, and any user supplied
113 definitions. As well as varying on variables, these definitions may also
114 exist for specific pressure levels to account for variables with
115 significantly different ranges at different heights. The colorbars also exist
116 for masks and mask differences for considering variable presence diagnostics.
117 Specific variable ranges can be separately set in user-supplied definition
118 for x- or y-axis limits, or indicate where automated range preferred.
120 Parameters
121 ----------
122 cube: Cube
123 Cube of variable for which the colorbar information is desired.
124 axis: "x", "y", optional
125 Select the levels for just this axis of a line plot. The min and max
126 can be set by xmin/xmax or ymin/ymax respectively. For variables where
127 setting a universal range is not desirable (e.g. temperature), users
128 can set ymin/ymax values to "auto" in the colorbar definitions file.
129 Where no additional xmin/xmax or ymin/ymax values are provided, the
130 axis bounds default to use the vmin/vmax values provided.
132 Returns
133 -------
134 cmap:
135 Matplotlib colormap.
136 levels:
137 List of levels to use for plotting. For continuous plots the min and max
138 should be taken as the range.
139 norm:
140 BoundaryNorm information.
141 """
142 # Grab the colorbar file from the recipe global metadata.
143 user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
144 colorbar = _load_colorbar_map(user_colorbar_file)
145 cmap = None
147 try:
148 # We assume that pressure is a scalar coordinate here.
149 pressure_level_raw = cube.coord("pressure").points[0]
150 # Ensure pressure_level is a string, as it is used as a JSON key.
151 pressure_level = str(int(pressure_level_raw))
152 except iris.exceptions.CoordinateNotFoundError:
153 pressure_level = None
155 # First try long name, then standard name, then var name. This order is used
156 # as long name is the one we correct between models, so it most likely to be
157 # consistent.
158 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name]))
159 varnames = [varname.replace("observed_", "") for varname in varnames]
160 for varname in varnames:
161 # Get the colormap for this variable.
162 try:
163 var_colorbar = colorbar[varname]
164 cmap = plt.get_cmap(colorbar[varname]["cmap"], 51)
165 varname_key = varname
166 break
167 except KeyError:
168 logging.debug("Cube name %s has no colorbar definition.", varname)
170 # Get colormap if it is a mask.
171 if any("mask_for_" in name for name in varnames):
172 cmap, levels, norm = _custom_colormap_mask(cube, axis=axis)
173 return cmap, levels, norm
174 # If winds on Beaufort Scale use custom colorbar and levels
175 if any("Beaufort_Scale" in name for name in varnames):
176 cmap, levels, norm = _custom_beaufort_scale(cube, axis=axis)
177 return cmap, levels, norm
178 # If probability is plotted use custom colorbar and levels
179 if any("probability_of_" in name for name in varnames):
180 cmap, levels, norm = _custom_colormap_probability(cube, axis=axis)
181 return cmap, levels, norm
182 # If aviation colour state use custom colorbar and levels
183 if any("aviation_colour_state" in name for name in varnames): 183 ↛ 184line 183 didn't jump to line 184 because the condition on line 183 was never true
184 cmap, levels, norm = _custom_colormap_aviation_colour_state(cube)
185 return cmap, levels, norm
187 # If no valid colormap has been defined, use defaults and return.
188 if not cmap:
189 logging.warning("No colorbar definition exists for %s.", cube.name())
190 cmap, levels, norm = mpl.colormaps["viridis"], None, None
191 return cmap, levels, norm
193 # Test if pressure-level specific settings are provided for cube.
194 if pressure_level:
195 try:
196 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level]
197 except KeyError:
198 logging.debug(
199 "%s has no colorbar definition for pressure level %s.",
200 varname,
201 pressure_level,
202 )
204 # Check for availability of x-axis or y-axis user-specific overrides
205 # for setting level bounds for line plot types and return just levels.
206 # Line plots do not need a colormap, and just use the data range.
207 if axis:
208 if axis == "x":
209 try:
210 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"]
211 except KeyError:
212 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
213 if axis == "y":
214 try:
215 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"]
216 except KeyError:
217 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
218 # Check if user-specified auto-scaling for this variable
219 if vmin == "auto" or vmax == "auto":
220 levels = None
221 else:
222 levels = [vmin, vmax]
223 return None, levels, None
224 # Get and use the colorbar levels for this variable if spatial or histogram.
225 else:
226 try:
227 levels = var_colorbar["levels"]
228 # Use discrete bins when levels are specified, rather
229 # than a smooth range.
230 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
231 logging.debug("Using levels for %s colorbar.", varname)
232 logging.info("Using levels: %s", levels)
233 except KeyError:
234 # Get the range for this variable.
235 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
236 logging.debug("Using min and max for %s colorbar.", varname)
237 # Calculate levels from range.
238 if vmin == "auto" or vmax == "auto": 238 ↛ 239line 238 didn't jump to line 239 because the condition on line 238 was never true
239 levels = None
240 else:
241 levels = np.linspace(vmin, vmax, 101)
242 norm = None
244 # Overwrite cmap, levels and norm for specific variables that
245 # require custom colorbar_map as these can not be defined in the
246 # JSON file.
247 cmap, levels, norm = _custom_colourmap_precipitation(cube, cmap, levels, norm)
248 cmap, levels, norm = _custom_colourmap_visibility_in_air(
249 cube, cmap, levels, norm
250 )
251 cmap, levels, norm = _custom_colormap_celsius(cube, cmap, levels, norm)
252 return cmap, levels, norm
255def _custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
256 """Get colourmap for mask.
258 If "mask_for_" appears anywhere in the name of a cube this function will be called
259 regardless of the name of the variable to ensure a consistent plot.
261 Parameters
262 ----------
263 cube: Cube
264 Cube of variable for which the colorbar information is desired.
265 axis: "x", "y", optional
266 Select the levels for just this axis of a line plot. The min and max
267 can be set by xmin/xmax or ymin/ymax respectively. For variables where
268 setting a universal range is not desirable (e.g. temperature), users
269 can set ymin/ymax values to "auto" in the colorbar definitions file.
270 Where no additional xmin/xmax or ymin/ymax values are provided, the
271 axis bounds default to use the vmin/vmax values provided.
273 Returns
274 -------
275 cmap:
276 Matplotlib colormap.
277 levels:
278 List of levels to use for plotting. For continuous plots the min and max
279 should be taken as the range.
280 norm:
281 BoundaryNorm information.
282 """
283 if "difference" not in cube.long_name:
284 if axis:
285 levels = [0, 1]
286 # Complete settings based on levels.
287 return None, levels, None
288 else:
289 # Define the levels and colors.
290 levels = [0, 1, 2]
291 colors = ["white", "dodgerblue"]
292 # Create a custom color map.
293 cmap = mcolors.ListedColormap(colors)
294 # Normalize the levels.
295 norm = mcolors.BoundaryNorm(levels, cmap.N)
296 logging.debug("Colourmap for %s.", cube.long_name)
297 return cmap, levels, norm
298 else:
299 if axis:
300 levels = [-1, 1]
301 return None, levels, None
302 else:
303 # Search for if mask difference, set to +/- 0.5 as values plotted <
304 # not <=.
305 levels = [-2, -0.5, 0.5, 2]
306 colors = ["goldenrod", "white", "teal"]
307 cmap = mcolors.ListedColormap(colors)
308 norm = mcolors.BoundaryNorm(levels, cmap.N)
309 logging.debug("Colourmap for %s.", cube.long_name)
310 return cmap, levels, norm
313def _custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
314 """Get a custom colorbar for a cube in the Beaufort Scale.
316 Specific variable ranges can be separately set in user-supplied definition
317 for x- or y-axis limits, or indicate where automated range preferred.
319 Parameters
320 ----------
321 cube: Cube
322 Cube of variable with Beaufort Scale in name.
323 axis: "x", "y", optional
324 Select the levels for just this axis of a line plot. The min and max
325 can be set by xmin/xmax or ymin/ymax respectively. For variables where
326 setting a universal range is not desirable (e.g. temperature), users
327 can set ymin/ymax values to "auto" in the colorbar definitions file.
328 Where no additional xmin/xmax or ymin/ymax values are provided, the
329 axis bounds default to use the vmin/vmax values provided.
331 Returns
332 -------
333 cmap:
334 Matplotlib colormap.
335 levels:
336 List of levels to use for plotting. For continuous plots the min and max
337 should be taken as the range.
338 norm:
339 BoundaryNorm information.
340 """
341 if "difference" not in cube.long_name:
342 if axis:
343 levels = [0, 12]
344 return None, levels, None
345 else:
346 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
347 colors = [
348 "black",
349 (0, 0, 0.6),
350 "blue",
351 "cyan",
352 "green",
353 "yellow",
354 (1, 0.5, 0),
355 "red",
356 "pink",
357 "magenta",
358 "purple",
359 "maroon",
360 "white",
361 ]
362 cmap = mcolors.ListedColormap(colors)
363 norm = mcolors.BoundaryNorm(levels, cmap.N)
364 logging.info("change colormap for Beaufort Scale colorbar.")
365 return cmap, levels, norm
366 else:
367 if axis:
368 levels = [-4, 4]
369 return None, levels, None
370 else:
371 levels = [
372 -3.5,
373 -2.5,
374 -1.5,
375 -0.5,
376 0.5,
377 1.5,
378 2.5,
379 3.5,
380 ]
381 cmap = plt.get_cmap("bwr", 8)
382 norm = mcolors.BoundaryNorm(levels, cmap.N)
383 return cmap, levels, norm
386def _custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm):
387 """Return altered colourmap for temperature with change in units to Celsius.
389 If "Celsius" appears anywhere in the name of a cube this function will be called.
391 Parameters
392 ----------
393 cube: Cube
394 Cube of variable for which the colorbar information is desired.
395 cmap: Matplotlib colormap.
396 levels: List
397 List of levels to use for plotting. For continuous plots the min and max
398 should be taken as the range.
399 norm: BoundaryNorm.
401 Returns
402 -------
403 cmap: Matplotlib colormap.
404 levels: List
405 List of levels to use for plotting. For continuous plots the min and max
406 should be taken as the range.
407 norm: BoundaryNorm.
408 """
409 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
410 if any("temperature" in name for name in varnames) and "Celsius" == cube.units:
411 levels = np.array(levels)
412 levels -= 273
413 levels = levels.tolist()
414 else:
415 # Do nothing keep the existing colourbar attributes
416 levels = levels
417 cmap = cmap
418 norm = norm
419 return cmap, levels, norm
422def _custom_colormap_probability(
423 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None
424):
425 """Get a custom colorbar for a probability cube.
427 Specific variable ranges can be separately set in user-supplied definition
428 for x- or y-axis limits, or indicate where automated range preferred.
430 Parameters
431 ----------
432 cube: Cube
433 Cube of variable with probability in name.
434 axis: "x", "y", optional
435 Select the levels for just this axis of a line plot. The min and max
436 can be set by xmin/xmax or ymin/ymax respectively. For variables where
437 setting a universal range is not desirable (e.g. temperature), users
438 can set ymin/ymax values to "auto" in the colorbar definitions file.
439 Where no additional xmin/xmax or ymin/ymax values are provided, the
440 axis bounds default to use the vmin/vmax values provided.
442 Returns
443 -------
444 cmap:
445 Matplotlib colormap.
446 levels:
447 List of levels to use for plotting. For continuous plots the min and max
448 should be taken as the range.
449 norm:
450 BoundaryNorm information.
451 """
452 if axis:
453 levels = [0, 1]
454 return None, levels, None
455 else:
456 cmap = mcolors.ListedColormap(
457 [
458 "#FFFFFF",
459 "#636363",
460 "#e1dada",
461 "#B5CAFF",
462 "#8FB3FF",
463 "#7F97FF",
464 "#ABCF63",
465 "#E8F59E",
466 "#FFFA14",
467 "#FFD121",
468 "#FFA30A",
469 ]
470 )
471 levels = [0.0, 0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
472 norm = mcolors.BoundaryNorm(levels, cmap.N)
473 return cmap, levels, norm
476def _custom_colourmap_precipitation(cube: iris.cube.Cube, cmap, levels, norm):
477 """Return a custom colourmap for the current recipe."""
478 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
479 if (
480 any("surface_microphysical" in name for name in varnames)
481 and "difference" not in cube.long_name
482 and "mask" not in cube.long_name
483 ):
484 # Define the levels and colors
485 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256]
486 colors = [
487 "w",
488 (0, 0, 0.6),
489 "b",
490 "c",
491 "g",
492 "y",
493 (1, 0.5, 0),
494 "r",
495 "pink",
496 "m",
497 "purple",
498 "maroon",
499 "gray",
500 ]
501 # Create a custom colormap
502 cmap = mcolors.ListedColormap(colors)
503 # Normalize the levels
504 norm = mcolors.BoundaryNorm(levels, cmap.N)
505 logging.info("change colormap for surface_microphysical variable colorbar.")
506 else:
507 # do nothing and keep existing colorbar attributes
508 cmap = cmap
509 levels = levels
510 norm = norm
511 return cmap, levels, norm
514def _custom_colormap_aviation_colour_state(cube: iris.cube.Cube):
515 """Return custom colourmap for aviation colour state.
517 If "aviation_colour_state" appears anywhere in the name of a cube
518 this function will be called.
520 Parameters
521 ----------
522 cube: Cube
523 Cube of variable for which the colorbar information is desired.
525 Returns
526 -------
527 cmap: Matplotlib colormap.
528 levels: List
529 List of levels to use for plotting. For continuous plots the min and max
530 should be taken as the range.
531 norm: BoundaryNorm.
532 """
533 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5]
534 colors = [
535 "#87ceeb",
536 "#ffffff",
537 "#8ced69",
538 "#ffff00",
539 "#ffd700",
540 "#ffa500",
541 "#fe3620",
542 ]
543 # Create a custom colormap
544 cmap = mcolors.ListedColormap(colors)
545 # Normalise the levels
546 norm = mcolors.BoundaryNorm(levels, cmap.N)
547 return cmap, levels, norm
550def _custom_colourmap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm):
551 """Return a custom colourmap for the current recipe."""
552 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
553 if (
554 any("visibility_in_air" in name for name in varnames)
555 and "difference" not in cube.long_name
556 and "mask" not in cube.long_name
557 ):
558 # Define the levels and colors (in km)
559 levels = [0, 0.05, 0.1, 0.2, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 50.0, 70.0, 100.0]
560 norm = mcolors.BoundaryNorm(levels, cmap.N)
561 colours = [
562 "#8f00d6",
563 "#d10000",
564 "#ff9700",
565 "#ffff00",
566 "#00007f",
567 "#6c9ccd",
568 "#aae8ff",
569 "#37a648",
570 "#8edc64",
571 "#c5ffc5",
572 "#dcdcdc",
573 "#ffffff",
574 ]
575 # Create a custom colormap
576 cmap = mcolors.ListedColormap(colours)
577 # Normalize the levels
578 norm = mcolors.BoundaryNorm(levels, cmap.N)
579 logging.info("change colormap for visibility_in_air variable colorbar.")
580 else:
581 # do nothing and keep existing colorbar attributes
582 cmap = cmap
583 levels = levels
584 norm = norm
585 return cmap, levels, norm