Coverage for src/CSET/operators/_colormaps.py: 99%
230 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-16 13:33 +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 color_list = itertools.cycle(DEFAULT_DISCRETE_COLORS)
95 return {mname: color for mname, color in zip(model_names, color_list, strict=False)}
98def colorbar_map_levels(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
99 """Get an appropriate colorbar for the given cube.
101 For the given variable the appropriate colorbar is looked up from a
102 combination of the built-in CSET colorbar definitions, and any user supplied
103 definitions. As well as varying on variables, these definitions may also
104 exist for specific pressure levels to account for variables with
105 significantly different ranges at different heights. The colorbars also exist
106 for masks and mask differences for considering variable presence diagnostics.
107 Specific variable ranges can be separately set in user-supplied definition
108 for x- or y-axis limits, or indicate where automated range preferred.
110 Parameters
111 ----------
112 cube: Cube
113 Cube of variable for which the colorbar information is desired.
114 axis: "x", "y", optional
115 Select the levels for just this axis of a line plot. The min and max
116 can be set by xmin/xmax or ymin/ymax respectively. For variables where
117 setting a universal range is not desirable (e.g. temperature), users
118 can set ymin/ymax values to "auto" in the colorbar definitions file.
119 Where no additional xmin/xmax or ymin/ymax values are provided, the
120 axis bounds default to use the vmin/vmax values provided.
122 Returns
123 -------
124 cmap:
125 Matplotlib colormap.
126 levels:
127 List of levels to use for plotting. For continuous plots the min and max
128 should be taken as the range.
129 norm:
130 BoundaryNorm information.
131 """
132 # Grab the colorbar file from the recipe global metadata.
133 user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
134 colorbar = load_colorbar_map(user_colorbar_file)
135 cmap = None
137 try:
138 # We assume that pressure is a scalar coordinate here.
139 pressure_level_raw = cube.coord("pressure").points[0]
140 # Ensure pressure_level is a string, as it is used as a JSON key.
141 pressure_level = str(int(pressure_level_raw))
142 except iris.exceptions.CoordinateNotFoundError:
143 pressure_level = None
145 # First try long name, then standard name, then var name. This order is used
146 # as long name is the one we correct between models, so it most likely to be
147 # consistent.
148 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name]))
149 # Treat observation-labelled var names consistently with model var names.
150 varnames = [varname.replace("observed_", "") for varname in varnames]
151 for varname in varnames:
152 # Get the colormap for this variable.
153 try:
154 var_colorbar = colorbar[varname]
155 cmap = plt.get_cmap(colorbar[varname]["cmap"], 51)
156 varname_key = varname
157 break
158 except KeyError:
159 logging.debug("Cube name %s has no colorbar definition.", varname)
161 # Get colormap if it is a mask.
162 if any("mask_for_" in name for name in varnames):
163 cmap, levels, norm = custom_colormap_mask(cube, axis=axis)
164 return cmap, levels, norm
165 # If winds on Beaufort Scale use custom colorbar and levels
166 if any("Beaufort_Scale" in name for name in varnames):
167 cmap, levels, norm = custom_beaufort_scale(cube, axis=axis)
168 return cmap, levels, norm
169 # If probability is plotted use custom colorbar and levels
170 if any("probability_of_" in name for name in varnames):
171 cmap, levels, norm = custom_colormap_probability(cube, axis=axis)
172 return cmap, levels, norm
173 # If aviation colour state use custom colorbar and levels
174 if any("aviation_colour_state" in name for name in varnames):
175 cmap, levels, norm = custom_colormap_aviation_colour_state(cube)
176 return cmap, levels, norm
177 # If verification scores use custom colorbar
178 if any("RMSE_" in name for name in varnames):
179 cmap, levels, norm = custom_colormap_scores(cube)
180 return cmap, levels, norm
182 # If no valid colormap has been defined, use defaults and return.
183 if not cmap:
184 logging.warning("No colorbar definition exists for %s.", cube.name())
185 cmap, levels, norm = mpl.colormaps["viridis"], None, None
186 return cmap, levels, norm
188 # Test if pressure-level specific settings are provided for cube.
189 if pressure_level:
190 try:
191 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level]
192 except KeyError:
193 logging.debug(
194 "%s has no colorbar definition for pressure level %s.",
195 varname,
196 pressure_level,
197 )
199 # Check for availability of x-axis or y-axis user-specific overrides
200 # for setting level bounds for line plot types and return just levels.
201 # Line plots do not need a colormap, and just use the data range.
202 if axis:
203 if axis == "x":
204 try:
205 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"]
206 except KeyError:
207 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
208 if axis == "y":
209 try:
210 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"]
211 except KeyError:
212 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
213 # Check if user-specified auto-scaling for this variable
214 if vmin == "auto" or vmax == "auto":
215 levels = None
216 else:
217 levels = [vmin, vmax]
218 return None, levels, None
219 # Get and use the colorbar levels for this variable if spatial or histogram.
220 else:
221 try:
222 levels = var_colorbar["levels"]
223 # Use discrete bins when levels are specified, rather
224 # than a smooth range.
225 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
226 logging.debug("Using levels for %s colorbar.", varname)
227 logging.info("Using levels: %s", levels)
228 except KeyError:
229 # Get the range for this variable.
230 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
231 logging.debug("Using min and max for %s colorbar.", varname)
232 # Calculate levels from range.
233 if vmin == "auto" or vmax == "auto":
234 levels = None
235 else:
236 levels = np.linspace(vmin, vmax, 101)
237 norm = None
239 # Overwrite cmap, levels and norm for specific variables that
240 # require custom colorbar_map as these can not be defined in the
241 # JSON file.
242 cmap, levels, norm = custom_colormap_precipitation(cube, cmap, levels, norm)
243 cmap, levels, norm = custom_colormap_visibility_in_air(cube, cmap, levels, norm)
244 cmap, levels, norm = custom_colormap_celsius(cube, cmap, levels, norm)
245 cmap, levels, norm = custom_colormap_feature_tracking(cube, cmap, levels, norm)
246 return cmap, levels, norm
249def custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
250 """Get colormap for mask.
252 If "mask_for_" appears anywhere in the name of a cube this function will be called
253 regardless of the name of the variable to ensure a consistent plot.
255 Parameters
256 ----------
257 cube: Cube
258 Cube of variable for which the colorbar information is desired.
259 axis: "x", "y", optional
260 Select the levels for just this axis of a line plot. The min and max
261 can be set by xmin/xmax or ymin/ymax respectively. For variables where
262 setting a universal range is not desirable (e.g. temperature), users
263 can set ymin/ymax values to "auto" in the colorbar definitions file.
264 Where no additional xmin/xmax or ymin/ymax values are provided, the
265 axis bounds default to use the vmin/vmax values provided.
267 Returns
268 -------
269 cmap:
270 Matplotlib colormap.
271 levels:
272 List of levels to use for plotting. For continuous plots the min and max
273 should be taken as the range.
274 norm:
275 BoundaryNorm information.
276 """
277 if "difference" not in cube.long_name:
278 if axis:
279 levels = [0, 1]
280 # Complete settings based on levels.
281 return None, levels, None
282 else:
283 # Define the levels and colors.
284 levels = [0, 1, 2]
285 colors = ["white", "dodgerblue"]
286 # Create a custom color map.
287 cmap = mcolors.ListedColormap(colors)
288 # Normalize the levels.
289 norm = mcolors.BoundaryNorm(levels, cmap.N)
290 logging.debug("Colormap for %s.", cube.long_name)
291 return cmap, levels, norm
292 else:
293 if axis:
294 levels = [-1, 1]
295 return None, levels, None
296 else:
297 # Search for if mask difference, set to +/- 0.5 as values plotted <
298 # not <=.
299 levels = [-2, -0.5, 0.5, 2]
300 colors = ["goldenrod", "white", "teal"]
301 cmap = mcolors.ListedColormap(colors)
302 norm = mcolors.BoundaryNorm(levels, cmap.N)
303 logging.debug("Colormap for %s.", cube.long_name)
304 return cmap, levels, norm
307def custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
308 """Get a custom colorbar for a cube in the Beaufort Scale.
310 Specific variable ranges can be separately set in user-supplied definition
311 for x- or y-axis limits, or indicate where automated range preferred.
313 Parameters
314 ----------
315 cube: Cube
316 Cube of variable with Beaufort Scale in name.
317 axis: "x", "y", optional
318 Select the levels for just this axis of a line plot. The min and max
319 can be set by xmin/xmax or ymin/ymax respectively. For variables where
320 setting a universal range is not desirable (e.g. temperature), users
321 can set ymin/ymax values to "auto" in the colorbar definitions file.
322 Where no additional xmin/xmax or ymin/ymax values are provided, the
323 axis bounds default to use the vmin/vmax values provided.
325 Returns
326 -------
327 cmap:
328 Matplotlib colormap.
329 levels:
330 List of levels to use for plotting. For continuous plots the min and max
331 should be taken as the range.
332 norm:
333 BoundaryNorm information.
334 """
335 if "difference" not in cube.long_name:
336 if axis:
337 levels = [0, 12]
338 return None, levels, None
339 else:
340 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
341 colors = [
342 "black",
343 (0, 0, 0.6),
344 "blue",
345 "cyan",
346 "green",
347 "yellow",
348 (1, 0.5, 0),
349 "red",
350 "pink",
351 "magenta",
352 "purple",
353 "maroon",
354 "white",
355 ]
356 cmap = mcolors.ListedColormap(colors)
357 norm = mcolors.BoundaryNorm(levels, cmap.N)
358 logging.info("change colormap for Beaufort Scale colorbar.")
359 return cmap, levels, norm
360 else:
361 if axis:
362 levels = [-4, 4]
363 return None, levels, None
364 else:
365 levels = [
366 -3.5,
367 -2.5,
368 -1.5,
369 -0.5,
370 0.5,
371 1.5,
372 2.5,
373 3.5,
374 ]
375 cmap = plt.get_cmap("bwr", 8)
376 norm = mcolors.BoundaryNorm(levels, cmap.N)
377 return cmap, levels, norm
380def custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm):
381 """Return altered colormap for temperature with change in units to Celsius.
383 If "Celsius" appears anywhere in the name of a cube this function will be called.
385 Parameters
386 ----------
387 cube: Cube
388 Cube of variable for which the colorbar information is desired.
389 cmap: Matplotlib colormap.
390 levels: List
391 List of levels to use for plotting. For continuous plots the min and max
392 should be taken as the range.
393 norm: BoundaryNorm.
395 Returns
396 -------
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.
402 """
403 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
404 if any("temperature" in name for name in varnames) and "Celsius" == cube.units:
405 levels = np.array(levels)
406 levels -= 273
407 levels = levels.tolist()
408 else:
409 # Do nothing keep the existing colourbar attributes
410 levels = levels
411 cmap = cmap
412 norm = norm
413 return cmap, levels, norm
416def custom_colormap_probability(
417 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None
418):
419 """Get a custom colorbar for a probability cube.
421 Specific variable ranges can be separately set in user-supplied definition
422 for x- or y-axis limits, or indicate where automated range preferred.
424 Parameters
425 ----------
426 cube: Cube
427 Cube of variable with probability in name.
428 axis: "x", "y", optional
429 Select the levels for just this axis of a line plot. The min and max
430 can be set by xmin/xmax or ymin/ymax respectively. For variables where
431 setting a universal range is not desirable (e.g. temperature), users
432 can set ymin/ymax values to "auto" in the colorbar definitions file.
433 Where no additional xmin/xmax or ymin/ymax values are provided, the
434 axis bounds default to use the vmin/vmax values provided.
436 Returns
437 -------
438 cmap:
439 Matplotlib colormap.
440 levels:
441 List of levels to use for plotting. For continuous plots the min and max
442 should be taken as the range.
443 norm:
444 BoundaryNorm information.
445 """
446 if axis:
447 levels = [0, 1]
448 return None, levels, None
449 else:
450 cmap = mcolors.ListedColormap(
451 [
452 "#FFFFFF",
453 "#636363",
454 "#e1dada",
455 "#B5CAFF",
456 "#8FB3FF",
457 "#7F97FF",
458 "#ABCF63",
459 "#E8F59E",
460 "#FFFA14",
461 "#FFD121",
462 "#FFA30A",
463 ]
464 )
465 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]
466 norm = mcolors.BoundaryNorm(levels, cmap.N)
467 return cmap, levels, norm
470def custom_colormap_precipitation(cube: iris.cube.Cube, cmap, levels, norm):
471 """Return a custom colormap for the current recipe."""
472 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
473 if (
474 any("surface_microphysical" in name for name in varnames)
475 and "difference" not in cube.long_name
476 and "mask" not in cube.long_name
477 ):
478 # Define the levels and colors
479 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256]
480 colors = [
481 "w",
482 (0, 0, 0.6),
483 "b",
484 "c",
485 "g",
486 "y",
487 (1, 0.5, 0),
488 "r",
489 "pink",
490 "m",
491 "purple",
492 "maroon",
493 "gray",
494 ]
495 # Create a custom colormap
496 cmap = mcolors.ListedColormap(colors)
497 # Normalize the levels
498 norm = mcolors.BoundaryNorm(levels, cmap.N)
499 logging.info("change colormap for surface_microphysical variable colorbar.")
500 else:
501 # do nothing and keep existing colorbar attributes
502 cmap = cmap
503 levels = levels
504 norm = norm
505 return cmap, levels, norm
508def custom_colormap_aviation_colour_state(cube: iris.cube.Cube):
509 """Return custom colormap for aviation colour state.
511 If "aviation_colour_state" appears anywhere in the name of a cube
512 this function will be called.
514 Parameters
515 ----------
516 cube: Cube
517 Cube of variable for which the colorbar information is desired.
519 Returns
520 -------
521 cmap: Matplotlib colormap.
522 levels: List
523 List of levels to use for plotting. For continuous plots the min and max
524 should be taken as the range.
525 norm: BoundaryNorm.
526 """
527 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5]
528 colors = [
529 "#87ceeb",
530 "#ffffff",
531 "#8ced69",
532 "#ffff00",
533 "#ffd700",
534 "#ffa500",
535 "#fe3620",
536 ]
537 # Create a custom colormap
538 cmap = mcolors.ListedColormap(colors)
539 # Normalise the levels
540 norm = mcolors.BoundaryNorm(levels, cmap.N)
541 return cmap, levels, norm
544def custom_colormap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm):
545 """Return a custom colormap for the current recipe."""
546 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
547 if (
548 any("visibility_in_air" in name for name in varnames)
549 and "difference" not in cube.long_name
550 and "mask" not in cube.long_name
551 ):
552 # Define the levels and colors (in km)
553 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]
554 norm = mcolors.BoundaryNorm(levels, cmap.N)
555 colours = [
556 "#8f00d6",
557 "#d10000",
558 "#ff9700",
559 "#ffff00",
560 "#00007f",
561 "#6c9ccd",
562 "#aae8ff",
563 "#37a648",
564 "#8edc64",
565 "#c5ffc5",
566 "#dcdcdc",
567 "#ffffff",
568 ]
569 # Create a custom colormap
570 cmap = mcolors.ListedColormap(colours)
571 # Normalize the levels
572 norm = mcolors.BoundaryNorm(levels, cmap.N)
573 logging.info("change colormap for visibility_in_air variable colorbar.")
574 else:
575 # do nothing and keep existing colorbar attributes
576 cmap = cmap
577 levels = levels
578 norm = norm
579 return cmap, levels, norm
582def custom_colormap_scores(cube: iris.cube.Cube):
583 """Return altered colormap for statistical metrics.
585 Parameters
586 ----------
587 cube: Cube
588 Cube of variable for which the colorbar information is desired.
590 Returns
591 -------
592 cmap: Matplotlib colormap.
593 levels: List
594 List of levels to use for plotting. For continuous plots the min and max
595 should be taken as the range.
596 norm: BoundaryNorm.
597 """
598 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
599 cmap, levels, norm = None, None, None
600 if any("RMSE_" in name for name in varnames): 600 ↛ 602line 600 didn't jump to line 602 because the condition on line 600 was always true
601 cmap = plt.get_cmap("PuRd", 51)
602 return cmap, levels, norm
605def custom_colormap_feature_tracking(cube: iris.cube.Cube, cmap, levels, norm):
606 """Return altered colormap for feature tracking.
608 Parameters
609 ----------
610 cube: Cube
611 Cube of variable for which the colorbar information is desired.
613 Returns
614 -------
615 cmap: Matplotlib colormap.
616 levels: List
617 List of levels to use for plotting. For continuous plots the min and max
618 should be taken as the range.
619 norm: BoundaryNorm.
620 """
621 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name]))
622 if (
623 any("feature_id" in name for name in varnames)
624 and "difference" not in cube.long_name
625 and "mask" not in cube.long_name
626 ):
627 # Define the levels and colors
628 levels = np.linspace(1, np.ma.max(cube.data), 10)
629 cmap = plt.get_cmap("viridis")
630 # Normalize the levels
631 norm = mcolors.BoundaryNorm(levels, cmap.N)
632 logging.info("change colormap for feature id variable colorbar.")
633 elif (
634 any("feature_lifetime" in name for name in varnames)
635 and "difference" not in cube.long_name
636 and "mask" not in cube.long_name
637 ):
638 # Define the levels and colors
639 levels = np.linspace(1, np.ma.max(cube.data), 10)
640 cmap = plt.get_cmap("YlGnBu")
641 # Normalize the levels
642 norm = mcolors.BoundaryNorm(levels, cmap.N)
643 logging.info("change colormap for feature lifetime variable colorbar.")
644 elif (
645 any("feature_init" in name for name in varnames)
646 and "difference" not in cube.long_name
647 and "mask" not in cube.long_name
648 ):
649 # Define the levels and colors
650 levels = np.array([0.5, 1])
651 cmap = plt.get_cmap("Blues")
652 # Normalize the levels
653 norm = mcolors.BoundaryNorm(levels, cmap.N)
654 logging.info("change colormap for feature init variable colorbar.")
656 else:
657 # do nothing and keep existing colorbar attributes
658 cmap = cmap
659 levels = levels
660 norm = norm
662 # Set all non-feature data to white
663 if any("feature" in name for name in varnames):
664 cmap.with_extremes(under="white")
666 return cmap, levels, norm