Coverage for src/CSET/operators/_colormaps.py: 92%
217 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-18 10:49 +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 for varname in varnames:
150 # Get the colormap for this variable.
151 try:
152 var_colorbar = colorbar[varname]
153 cmap = plt.get_cmap(colorbar[varname]["cmap"], 51)
154 varname_key = varname
155 break
156 except KeyError:
157 logging.debug("Cube name %s has no colorbar definition.", varname)
159 # Get colormap if it is a mask.
160 if any("mask_for_" in name for name in varnames):
161 cmap, levels, norm = custom_colormap_mask(cube, axis=axis)
162 return cmap, levels, norm
163 # If winds on Beaufort Scale use custom colorbar and levels
164 if any("Beaufort_Scale" in name for name in varnames):
165 cmap, levels, norm = custom_beaufort_scale(cube, axis=axis)
166 return cmap, levels, norm
167 # If probability is plotted use custom colorbar and levels
168 if any("probability_of_" in name for name in varnames):
169 cmap, levels, norm = custom_colormap_probability(cube, axis=axis)
170 return cmap, levels, norm
171 # If aviation colour state use custom colorbar and levels
172 if any("aviation_colour_state" in name for name in varnames): 172 ↛ 173line 172 didn't jump to line 173 because the condition on line 172 was never true
173 cmap, levels, norm = custom_colormap_aviation_colour_state(cube)
174 return cmap, levels, norm
175 # If verification scores use custom colorbar
176 if any("RMSE_" in name for name in varnames):
177 cmap, levels, norm = custom_colormap_scores(cube)
178 return cmap, levels, norm
179 if any("CURV_" in name for name in varnames): 179 ↛ 180line 179 didn't jump to line 180 because the condition on line 179 was never true
180 cmap, levels, norm = custom_colormap_curv(cube)
181 return cmap, levels, norm
183 # If no valid colormap has been defined, use defaults and return.
184 if not cmap:
185 logging.warning("No colorbar definition exists for %s.", cube.name())
186 cmap, levels, norm = mpl.colormaps["viridis"], None, None
187 return cmap, levels, norm
189 # Test if pressure-level specific settings are provided for cube.
190 if pressure_level:
191 try:
192 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level]
193 except KeyError:
194 logging.debug(
195 "%s has no colorbar definition for pressure level %s.",
196 varname,
197 pressure_level,
198 )
200 # Check for availability of x-axis or y-axis user-specific overrides
201 # for setting level bounds for line plot types and return just levels.
202 # Line plots do not need a colormap, and just use the data range.
203 if axis:
204 if axis == "x":
205 try:
206 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"]
207 except KeyError:
208 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
209 if axis == "y":
210 try:
211 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"]
212 except KeyError:
213 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
214 # Check if user-specified auto-scaling for this variable
215 if vmin == "auto" or vmax == "auto":
216 levels = None
217 else:
218 levels = [vmin, vmax]
219 return None, levels, None
220 # Get and use the colorbar levels for this variable if spatial or histogram.
221 else:
222 try:
223 levels = var_colorbar["levels"]
224 # Use discrete bins when levels are specified, rather
225 # than a smooth range.
226 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
227 logging.debug("Using levels for %s colorbar.", varname)
228 logging.info("Using levels: %s", levels)
229 except KeyError:
230 # Get the range for this variable.
231 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
232 logging.debug("Using min and max for %s colorbar.", varname)
233 # Calculate levels from range.
234 if vmin == "auto" or vmax == "auto": 234 ↛ 235line 234 didn't jump to line 235 because the condition on line 234 was never true
235 levels = None
236 else:
237 levels = np.linspace(vmin, vmax, 101)
238 norm = None
240 # Overwrite cmap, levels and norm for specific variables that
241 # require custom colorbar_map as these can not be defined in the
242 # JSON file.
243 cmap, levels, norm = custom_colormap_precipitation(cube, cmap, levels, norm)
244 cmap, levels, norm = custom_colormap_visibility_in_air(cube, cmap, levels, norm)
245 cmap, levels, norm = custom_colormap_celsius(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_curv(cube: iris.cube.Cube):
606 """Return custom colourmap for curv.
608 If "CURV_" appears anywhere in the name of a cube
609 this function will be called.
611 Parameters
612 ----------
613 cube: Cube
614 Cube of variable for which the colorbar information is desired.
616 Returns
617 -------
618 cmap: Matplotlib colormap.
619 levels: List
620 List of levels to use for plotting. For continuous plots the min and max
621 should be taken as the range.
622 norm: BoundaryNorm.
623 """
624 if "16" in cube.long_name:
625 levels = [-17, -15, -13, -11, -9, -7, -5, -3, -1, 1, 3, 5, 7, 9, 11, 13, 15, 17]
626 colors = [
627 "#01153e",
628 "#030764",
629 "#00008b",
630 "#0000ff",
631 "#0323df",
632 "#069af3",
633 "#00ffff",
634 "#7fffd4",
635 "#ffffff",
636 "#ffd700",
637 "#fac205",
638 "#ffa500",
639 "#f97306",
640 "#ff4500",
641 "#ff0000",
642 "#dc143c",
643 "#a52a2a",
644 ]
645 else:
646 levels = [-9, -7, -5, -3, -1, 1, 3, 5, 7, 9]
647 colors = [
648 "#01153e",
649 "#00008b",
650 "#0323df",
651 "#00ffff",
652 "#ffffff",
653 "#fac205",
654 "#f97306",
655 "#ff0000",
656 "#a52a2a",
657 ]
658 # Create a custom colormap
659 cmap = mcolors.ListedColormap(colors)
660 # Normalise the levels
661 norm = mcolors.BoundaryNorm(levels, cmap.N)
662 return cmap, levels, norm