Coverage for src/CSET/operators/_colormaps.py: 99%
213 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-19 11:18 +0000
1# © Crown copyright, Met Office (2022-2026) 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 with set color
95 if any("OBS" in name.upper() for name in model_names):
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):
184 cmap, levels, norm = custom_colormap_aviation_colour_state(cube)
185 return cmap, levels, norm
186 # If verification scores use custom colorbar
187 if any("RMSE_" in name for name in varnames):
188 cmap, levels, norm = custom_colormap_scores(cube)
189 return cmap, levels, norm
191 # If no valid colormap has been defined, use defaults and return.
192 if not cmap:
193 logging.warning("No colorbar definition exists for %s.", cube.name())
194 cmap, levels, norm = mpl.colormaps["viridis"], None, None
195 return cmap, levels, norm
197 # Test if pressure-level specific settings are provided for cube.
198 if pressure_level:
199 try:
200 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level]
201 except KeyError:
202 logging.debug(
203 "%s has no colorbar definition for pressure level %s.",
204 varname,
205 pressure_level,
206 )
208 # Check for availability of x-axis or y-axis user-specific overrides
209 # for setting level bounds for line plot types and return just levels.
210 # Line plots do not need a colormap, and just use the data range.
211 if axis:
212 if axis == "x":
213 try:
214 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"]
215 except KeyError:
216 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
217 if axis == "y":
218 try:
219 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"]
220 except KeyError:
221 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
222 # Check if user-specified auto-scaling for this variable
223 if vmin == "auto" or vmax == "auto":
224 levels = None
225 else:
226 levels = [vmin, vmax]
227 return None, levels, None
228 # Get and use the colorbar levels for this variable if spatial or histogram.
229 else:
230 try:
231 levels = var_colorbar["levels"]
232 # Use discrete bins when levels are specified, rather
233 # than a smooth range.
234 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
235 logging.debug("Using levels for %s colorbar.", varname)
236 logging.info("Using levels: %s", levels)
237 except KeyError:
238 # Get the range for this variable.
239 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
240 logging.debug("Using min and max for %s colorbar.", varname)
241 # Calculate levels from range.
242 if vmin == "auto" or vmax == "auto":
243 levels = None
244 else:
245 levels = np.linspace(vmin, vmax, 101)
246 norm = None
248 # Overwrite cmap, levels and norm for specific variables that
249 # require custom colorbar_map as these can not be defined in the
250 # JSON file.
251 cmap, levels, norm = custom_colormap_precipitation(cube, cmap, levels, norm)
252 cmap, levels, norm = custom_colormap_visibility_in_air(cube, cmap, levels, norm)
253 cmap, levels, norm = custom_colormap_celsius(cube, cmap, levels, norm)
254 return cmap, levels, norm
257def custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
258 """Get colormap for mask.
260 If "mask_for_" appears anywhere in the name of a cube this function will be called
261 regardless of the name of the variable to ensure a consistent plot.
263 Parameters
264 ----------
265 cube: Cube
266 Cube of variable for which the colorbar information is desired.
267 axis: "x", "y", optional
268 Select the levels for just this axis of a line plot. The min and max
269 can be set by xmin/xmax or ymin/ymax respectively. For variables where
270 setting a universal range is not desirable (e.g. temperature), users
271 can set ymin/ymax values to "auto" in the colorbar definitions file.
272 Where no additional xmin/xmax or ymin/ymax values are provided, the
273 axis bounds default to use the vmin/vmax values provided.
275 Returns
276 -------
277 cmap:
278 Matplotlib colormap.
279 levels:
280 List of levels to use for plotting. For continuous plots the min and max
281 should be taken as the range.
282 norm:
283 BoundaryNorm information.
284 """
285 if "difference" not in cube.long_name:
286 if axis:
287 levels = [0, 1]
288 # Complete settings based on levels.
289 return None, levels, None
290 else:
291 # Define the levels and colors.
292 levels = [0, 1, 2]
293 colors = ["white", "dodgerblue"]
294 # Create a custom color map.
295 cmap = mcolors.ListedColormap(colors)
296 # Normalize the levels.
297 norm = mcolors.BoundaryNorm(levels, cmap.N)
298 logging.debug("Colormap for %s.", cube.long_name)
299 return cmap, levels, norm
300 else:
301 if axis:
302 levels = [-1, 1]
303 return None, levels, None
304 else:
305 # Search for if mask difference, set to +/- 0.5 as values plotted <
306 # not <=.
307 levels = [-2, -0.5, 0.5, 2]
308 colors = ["goldenrod", "white", "teal"]
309 cmap = mcolors.ListedColormap(colors)
310 norm = mcolors.BoundaryNorm(levels, cmap.N)
311 logging.debug("Colormap for %s.", cube.long_name)
312 return cmap, levels, norm
315def custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
316 """Get a custom colorbar for a cube in the Beaufort Scale.
318 Specific variable ranges can be separately set in user-supplied definition
319 for x- or y-axis limits, or indicate where automated range preferred.
321 Parameters
322 ----------
323 cube: Cube
324 Cube of variable with Beaufort Scale in name.
325 axis: "x", "y", optional
326 Select the levels for just this axis of a line plot. The min and max
327 can be set by xmin/xmax or ymin/ymax respectively. For variables where
328 setting a universal range is not desirable (e.g. temperature), users
329 can set ymin/ymax values to "auto" in the colorbar definitions file.
330 Where no additional xmin/xmax or ymin/ymax values are provided, the
331 axis bounds default to use the vmin/vmax values provided.
333 Returns
334 -------
335 cmap:
336 Matplotlib colormap.
337 levels:
338 List of levels to use for plotting. For continuous plots the min and max
339 should be taken as the range.
340 norm:
341 BoundaryNorm information.
342 """
343 if "difference" not in cube.long_name:
344 if axis:
345 levels = [0, 12]
346 return None, levels, None
347 else:
348 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
349 colors = [
350 "black",
351 (0, 0, 0.6),
352 "blue",
353 "cyan",
354 "green",
355 "yellow",
356 (1, 0.5, 0),
357 "red",
358 "pink",
359 "magenta",
360 "purple",
361 "maroon",
362 "white",
363 ]
364 cmap = mcolors.ListedColormap(colors)
365 norm = mcolors.BoundaryNorm(levels, cmap.N)
366 logging.info("change colormap for Beaufort Scale colorbar.")
367 return cmap, levels, norm
368 else:
369 if axis:
370 levels = [-4, 4]
371 return None, levels, None
372 else:
373 levels = [
374 -3.5,
375 -2.5,
376 -1.5,
377 -0.5,
378 0.5,
379 1.5,
380 2.5,
381 3.5,
382 ]
383 cmap = plt.get_cmap("bwr", 8)
384 norm = mcolors.BoundaryNorm(levels, cmap.N)
385 return cmap, levels, norm
388def custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm):
389 """Return altered colormap for temperature with change in units to Celsius.
391 If "Celsius" appears anywhere in the name of a cube this function will be called.
393 Parameters
394 ----------
395 cube: Cube
396 Cube of variable for which the colorbar information is desired.
397 cmap: Matplotlib colormap.
398 levels: List
399 List of levels to use for plotting. For continuous plots the min and max
400 should be taken as the range.
401 norm: BoundaryNorm.
403 Returns
404 -------
405 cmap: Matplotlib colormap.
406 levels: List
407 List of levels to use for plotting. For continuous plots the min and max
408 should be taken as the range.
409 norm: BoundaryNorm.
410 """
411 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
412 if any("temperature" in name for name in varnames) and "Celsius" == cube.units:
413 levels = np.array(levels)
414 levels -= 273
415 levels = levels.tolist()
416 else:
417 # Do nothing keep the existing colourbar attributes
418 levels = levels
419 cmap = cmap
420 norm = norm
421 return cmap, levels, norm
424def custom_colormap_probability(
425 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None
426):
427 """Get a custom colorbar for a probability cube.
429 Specific variable ranges can be separately set in user-supplied definition
430 for x- or y-axis limits, or indicate where automated range preferred.
432 Parameters
433 ----------
434 cube: Cube
435 Cube of variable with probability in name.
436 axis: "x", "y", optional
437 Select the levels for just this axis of a line plot. The min and max
438 can be set by xmin/xmax or ymin/ymax respectively. For variables where
439 setting a universal range is not desirable (e.g. temperature), users
440 can set ymin/ymax values to "auto" in the colorbar definitions file.
441 Where no additional xmin/xmax or ymin/ymax values are provided, the
442 axis bounds default to use the vmin/vmax values provided.
444 Returns
445 -------
446 cmap:
447 Matplotlib colormap.
448 levels:
449 List of levels to use for plotting. For continuous plots the min and max
450 should be taken as the range.
451 norm:
452 BoundaryNorm information.
453 """
454 if axis:
455 levels = [0, 1]
456 return None, levels, None
457 else:
458 cmap = mcolors.ListedColormap(
459 [
460 "#FFFFFF",
461 "#636363",
462 "#e1dada",
463 "#B5CAFF",
464 "#8FB3FF",
465 "#7F97FF",
466 "#ABCF63",
467 "#E8F59E",
468 "#FFFA14",
469 "#FFD121",
470 "#FFA30A",
471 ]
472 )
473 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]
474 norm = mcolors.BoundaryNorm(levels, cmap.N)
475 return cmap, levels, norm
478def custom_colormap_precipitation(cube: iris.cube.Cube, cmap, levels, norm):
479 """Return a custom colormap for the current recipe."""
480 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
481 if (
482 any("surface_microphysical" in name for name in varnames)
483 and "difference" not in cube.long_name
484 and "mask" not in cube.long_name
485 ):
486 # Define the levels and colors
487 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256]
488 colors = [
489 "w",
490 (0, 0, 0.6),
491 "b",
492 "c",
493 "g",
494 "y",
495 (1, 0.5, 0),
496 "r",
497 "pink",
498 "m",
499 "purple",
500 "maroon",
501 "gray",
502 ]
503 # Create a custom colormap
504 cmap = mcolors.ListedColormap(colors)
505 # Normalize the levels
506 norm = mcolors.BoundaryNorm(levels, cmap.N)
507 logging.info("change colormap for surface_microphysical variable colorbar.")
508 else:
509 # do nothing and keep existing colorbar attributes
510 cmap = cmap
511 levels = levels
512 norm = norm
513 return cmap, levels, norm
516def custom_colormap_aviation_colour_state(cube: iris.cube.Cube):
517 """Return custom colormap for aviation colour state.
519 If "aviation_colour_state" appears anywhere in the name of a cube
520 this function will be called.
522 Parameters
523 ----------
524 cube: Cube
525 Cube of variable for which the colorbar information is desired.
527 Returns
528 -------
529 cmap: Matplotlib colormap.
530 levels: List
531 List of levels to use for plotting. For continuous plots the min and max
532 should be taken as the range.
533 norm: BoundaryNorm.
534 """
535 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5]
536 colors = [
537 "#87ceeb",
538 "#ffffff",
539 "#8ced69",
540 "#ffff00",
541 "#ffd700",
542 "#ffa500",
543 "#fe3620",
544 ]
545 # Create a custom colormap
546 cmap = mcolors.ListedColormap(colors)
547 # Normalise the levels
548 norm = mcolors.BoundaryNorm(levels, cmap.N)
549 return cmap, levels, norm
552def custom_colormap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm):
553 """Return a custom colormap for the current recipe."""
554 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
555 if (
556 any("visibility_in_air" in name for name in varnames)
557 and "difference" not in cube.long_name
558 and "mask" not in cube.long_name
559 ):
560 # Define the levels and colors (in km)
561 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]
562 norm = mcolors.BoundaryNorm(levels, cmap.N)
563 colours = [
564 "#8f00d6",
565 "#d10000",
566 "#ff9700",
567 "#ffff00",
568 "#00007f",
569 "#6c9ccd",
570 "#aae8ff",
571 "#37a648",
572 "#8edc64",
573 "#c5ffc5",
574 "#dcdcdc",
575 "#ffffff",
576 ]
577 # Create a custom colormap
578 cmap = mcolors.ListedColormap(colours)
579 # Normalize the levels
580 norm = mcolors.BoundaryNorm(levels, cmap.N)
581 logging.info("change colormap for visibility_in_air variable colorbar.")
582 else:
583 # do nothing and keep existing colorbar attributes
584 cmap = cmap
585 levels = levels
586 norm = norm
587 return cmap, levels, norm
590def custom_colormap_scores(cube: iris.cube.Cube):
591 """Return altered colormap for statistical metrics.
593 Parameters
594 ----------
595 cube: Cube
596 Cube of variable for which the colorbar information is desired.
598 Returns
599 -------
600 cmap: Matplotlib colormap.
601 levels: List
602 List of levels to use for plotting. For continuous plots the min and max
603 should be taken as the range.
604 norm: BoundaryNorm.
605 """
606 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
607 cmap, levels, norm = None, None, None
608 if any("RMSE_" in name for name in varnames): 608 ↛ 610line 608 didn't jump to line 610 because the condition on line 608 was always true
609 cmap = plt.get_cmap("PuRd", 51)
610 return cmap, levels, norm