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