Coverage for src/CSET/operators/_colormaps.py: 97%
205 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-17 15:44 +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 return cmap, levels, norm
246def custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
247 """Get colormap for mask.
249 If "mask_for_" appears anywhere in the name of a cube this function will be called
250 regardless of the name of the variable to ensure a consistent plot.
252 Parameters
253 ----------
254 cube: Cube
255 Cube of variable for which the colorbar information is desired.
256 axis: "x", "y", optional
257 Select the levels for just this axis of a line plot. The min and max
258 can be set by xmin/xmax or ymin/ymax respectively. For variables where
259 setting a universal range is not desirable (e.g. temperature), users
260 can set ymin/ymax values to "auto" in the colorbar definitions file.
261 Where no additional xmin/xmax or ymin/ymax values are provided, the
262 axis bounds default to use the vmin/vmax values provided.
264 Returns
265 -------
266 cmap:
267 Matplotlib colormap.
268 levels:
269 List of levels to use for plotting. For continuous plots the min and max
270 should be taken as the range.
271 norm:
272 BoundaryNorm information.
273 """
274 if "difference" not in cube.long_name:
275 if axis:
276 levels = [0, 1]
277 # Complete settings based on levels.
278 return None, levels, None
279 else:
280 # Define the levels and colors.
281 levels = [0, 1, 2]
282 colors = ["white", "dodgerblue"]
283 # Create a custom color map.
284 cmap = mcolors.ListedColormap(colors)
285 # Normalize the levels.
286 norm = mcolors.BoundaryNorm(levels, cmap.N)
287 logging.debug("Colormap for %s.", cube.long_name)
288 return cmap, levels, norm
289 else:
290 if axis:
291 levels = [-1, 1]
292 return None, levels, None
293 else:
294 # Search for if mask difference, set to +/- 0.5 as values plotted <
295 # not <=.
296 levels = [-2, -0.5, 0.5, 2]
297 colors = ["goldenrod", "white", "teal"]
298 cmap = mcolors.ListedColormap(colors)
299 norm = mcolors.BoundaryNorm(levels, cmap.N)
300 logging.debug("Colormap for %s.", cube.long_name)
301 return cmap, levels, norm
304def custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
305 """Get a custom colorbar for a cube in the Beaufort Scale.
307 Specific variable ranges can be separately set in user-supplied definition
308 for x- or y-axis limits, or indicate where automated range preferred.
310 Parameters
311 ----------
312 cube: Cube
313 Cube of variable with Beaufort Scale in name.
314 axis: "x", "y", optional
315 Select the levels for just this axis of a line plot. The min and max
316 can be set by xmin/xmax or ymin/ymax respectively. For variables where
317 setting a universal range is not desirable (e.g. temperature), users
318 can set ymin/ymax values to "auto" in the colorbar definitions file.
319 Where no additional xmin/xmax or ymin/ymax values are provided, the
320 axis bounds default to use the vmin/vmax values provided.
322 Returns
323 -------
324 cmap:
325 Matplotlib colormap.
326 levels:
327 List of levels to use for plotting. For continuous plots the min and max
328 should be taken as the range.
329 norm:
330 BoundaryNorm information.
331 """
332 if "difference" not in cube.long_name:
333 if axis:
334 levels = [0, 12]
335 return None, levels, None
336 else:
337 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
338 colors = [
339 "black",
340 (0, 0, 0.6),
341 "blue",
342 "cyan",
343 "green",
344 "yellow",
345 (1, 0.5, 0),
346 "red",
347 "pink",
348 "magenta",
349 "purple",
350 "maroon",
351 "white",
352 ]
353 cmap = mcolors.ListedColormap(colors)
354 norm = mcolors.BoundaryNorm(levels, cmap.N)
355 logging.info("change colormap for Beaufort Scale colorbar.")
356 return cmap, levels, norm
357 else:
358 if axis:
359 levels = [-4, 4]
360 return None, levels, None
361 else:
362 levels = [
363 -3.5,
364 -2.5,
365 -1.5,
366 -0.5,
367 0.5,
368 1.5,
369 2.5,
370 3.5,
371 ]
372 cmap = plt.get_cmap("bwr", 8)
373 norm = mcolors.BoundaryNorm(levels, cmap.N)
374 return cmap, levels, norm
377def custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm):
378 """Return altered colormap for temperature with change in units to Celsius.
380 If "Celsius" appears anywhere in the name of a cube this function will be called.
382 Parameters
383 ----------
384 cube: Cube
385 Cube of variable for which the colorbar information is desired.
386 cmap: Matplotlib colormap.
387 levels: List
388 List of levels to use for plotting. For continuous plots the min and max
389 should be taken as the range.
390 norm: BoundaryNorm.
392 Returns
393 -------
394 cmap: Matplotlib colormap.
395 levels: List
396 List of levels to use for plotting. For continuous plots the min and max
397 should be taken as the range.
398 norm: BoundaryNorm.
399 """
400 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
401 if any("temperature" in name for name in varnames) and "Celsius" == cube.units:
402 levels = np.array(levels)
403 levels -= 273
404 levels = levels.tolist()
405 else:
406 # Do nothing keep the existing colourbar attributes
407 levels = levels
408 cmap = cmap
409 norm = norm
410 return cmap, levels, norm
413def custom_colormap_probability(
414 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None
415):
416 """Get a custom colorbar for a probability cube.
418 Specific variable ranges can be separately set in user-supplied definition
419 for x- or y-axis limits, or indicate where automated range preferred.
421 Parameters
422 ----------
423 cube: Cube
424 Cube of variable with probability in name.
425 axis: "x", "y", optional
426 Select the levels for just this axis of a line plot. The min and max
427 can be set by xmin/xmax or ymin/ymax respectively. For variables where
428 setting a universal range is not desirable (e.g. temperature), users
429 can set ymin/ymax values to "auto" in the colorbar definitions file.
430 Where no additional xmin/xmax or ymin/ymax values are provided, the
431 axis bounds default to use the vmin/vmax values provided.
433 Returns
434 -------
435 cmap:
436 Matplotlib colormap.
437 levels:
438 List of levels to use for plotting. For continuous plots the min and max
439 should be taken as the range.
440 norm:
441 BoundaryNorm information.
442 """
443 if axis:
444 levels = [0, 1]
445 return None, levels, None
446 else:
447 cmap = mcolors.ListedColormap(
448 [
449 "#FFFFFF",
450 "#636363",
451 "#e1dada",
452 "#B5CAFF",
453 "#8FB3FF",
454 "#7F97FF",
455 "#ABCF63",
456 "#E8F59E",
457 "#FFFA14",
458 "#FFD121",
459 "#FFA30A",
460 ]
461 )
462 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]
463 norm = mcolors.BoundaryNorm(levels, cmap.N)
464 return cmap, levels, norm
467def custom_colormap_precipitation(cube: iris.cube.Cube, cmap, levels, norm):
468 """Return a custom colormap for the current recipe."""
469 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
470 if (
471 any("surface_microphysical" in name for name in varnames)
472 and "difference" not in cube.long_name
473 and "mask" not in cube.long_name
474 ):
475 # Define the levels and colors
476 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256]
477 colors = [
478 "w",
479 (0, 0, 0.6),
480 "b",
481 "c",
482 "g",
483 "y",
484 (1, 0.5, 0),
485 "r",
486 "pink",
487 "m",
488 "purple",
489 "maroon",
490 "gray",
491 ]
492 # Create a custom colormap
493 cmap = mcolors.ListedColormap(colors)
494 # Normalize the levels
495 norm = mcolors.BoundaryNorm(levels, cmap.N)
496 logging.info("change colormap for surface_microphysical variable colorbar.")
497 else:
498 # do nothing and keep existing colorbar attributes
499 cmap = cmap
500 levels = levels
501 norm = norm
502 return cmap, levels, norm
505def custom_colormap_aviation_colour_state(cube: iris.cube.Cube):
506 """Return custom colormap for aviation colour state.
508 If "aviation_colour_state" appears anywhere in the name of a cube
509 this function will be called.
511 Parameters
512 ----------
513 cube: Cube
514 Cube of variable for which the colorbar information is desired.
516 Returns
517 -------
518 cmap: Matplotlib colormap.
519 levels: List
520 List of levels to use for plotting. For continuous plots the min and max
521 should be taken as the range.
522 norm: BoundaryNorm.
523 """
524 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5]
525 colors = [
526 "#87ceeb",
527 "#ffffff",
528 "#8ced69",
529 "#ffff00",
530 "#ffd700",
531 "#ffa500",
532 "#fe3620",
533 ]
534 # Create a custom colormap
535 cmap = mcolors.ListedColormap(colors)
536 # Normalise the levels
537 norm = mcolors.BoundaryNorm(levels, cmap.N)
538 return cmap, levels, norm
541def custom_colormap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm):
542 """Return a custom colormap for the current recipe."""
543 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
544 if (
545 any("visibility_in_air" in name for name in varnames)
546 and "difference" not in cube.long_name
547 and "mask" not in cube.long_name
548 ):
549 # Define the levels and colors (in km)
550 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]
551 norm = mcolors.BoundaryNorm(levels, cmap.N)
552 colours = [
553 "#8f00d6",
554 "#d10000",
555 "#ff9700",
556 "#ffff00",
557 "#00007f",
558 "#6c9ccd",
559 "#aae8ff",
560 "#37a648",
561 "#8edc64",
562 "#c5ffc5",
563 "#dcdcdc",
564 "#ffffff",
565 ]
566 # Create a custom colormap
567 cmap = mcolors.ListedColormap(colours)
568 # Normalize the levels
569 norm = mcolors.BoundaryNorm(levels, cmap.N)
570 logging.info("change colormap for visibility_in_air variable colorbar.")
571 else:
572 # do nothing and keep existing colorbar attributes
573 cmap = cmap
574 levels = levels
575 norm = norm
576 return cmap, levels, norm
579def custom_colormap_scores(cube: iris.cube.Cube):
580 """Return altered colormap for statistical metrics.
582 Parameters
583 ----------
584 cube: Cube
585 Cube of variable for which the colorbar information is desired.
587 Returns
588 -------
589 cmap: Matplotlib colormap.
590 levels: List
591 List of levels to use for plotting. For continuous plots the min and max
592 should be taken as the range.
593 norm: BoundaryNorm.
594 """
595 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
596 cmap, levels, norm = None, None, None
597 if any("RMSE_" in name for name in varnames): 597 ↛ 599line 597 didn't jump to line 599 because the condition on line 597 was always true
598 cmap = plt.get_cmap("PuRd", 51)
599 return cmap, levels, norm