Coverage for src / CSET / operators / plot.py: 89%
929 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-11-27 11:58 +0000
1# © Crown copyright, Met Office (2022-2025) 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"""Operators to produce various kinds of plots."""
17import fcntl
18import functools
19import importlib.resources
20import itertools
21import json
22import logging
23import math
24import os
25import sys
26from typing import Literal
28import cartopy.crs as ccrs
29import iris
30import iris.coords
31import iris.cube
32import iris.exceptions
33import iris.plot as iplt
34import matplotlib as mpl
35import matplotlib.colors as mcolors
36import matplotlib.pyplot as plt
37import numpy as np
38import scipy.fft as fft
39from iris.cube import Cube
40from markdown_it import MarkdownIt
42from CSET._common import (
43 combine_dicts,
44 get_recipe_metadata,
45 iter_maybe,
46 render_file,
47 slugify,
48)
49from CSET.operators._utils import (
50 fully_equalise_attributes,
51 get_cube_yxcoordname,
52 is_transect,
53)
54from CSET.operators.collapse import collapse
55from CSET.operators.misc import _extract_common_time_points
56from CSET.operators.regrid import regrid_onto_cube
58# Use a non-interactive plotting backend.
59mpl.use("agg")
61DEFAULT_DISCRETE_COLORS = mpl.colormaps["tab10"].colors + mpl.colormaps["Accent"].colors
63############################
64# Private helper functions #
65############################
68def _append_to_plot_index(plot_index: list) -> list:
69 """Add plots into the plot index, returning the complete plot index."""
70 with open("meta.json", "r+t", encoding="UTF-8") as fp:
71 fcntl.flock(fp, fcntl.LOCK_EX)
72 fp.seek(0)
73 meta = json.load(fp)
74 complete_plot_index = meta.get("plots", [])
75 complete_plot_index = complete_plot_index + plot_index
76 meta["plots"] = complete_plot_index
77 if os.getenv("CYLC_TASK_CYCLE_POINT") and not bool(
78 os.getenv("DO_CASE_AGGREGATION")
79 ):
80 meta["case_date"] = os.getenv("CYLC_TASK_CYCLE_POINT", "")
81 fp.seek(0)
82 fp.truncate()
83 json.dump(meta, fp, indent=2)
84 return complete_plot_index
87def _check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube:
88 """Ensure a single cube is given.
90 If a CubeList of length one is given that the contained cube is returned,
91 otherwise an error is raised.
93 Parameters
94 ----------
95 cube: Cube | CubeList
96 The cube to check.
98 Returns
99 -------
100 cube: Cube
101 The checked cube.
103 Raises
104 ------
105 TypeError
106 If the input cube is not a Cube or CubeList of a single Cube.
107 """
108 if isinstance(cube, iris.cube.Cube):
109 return cube
110 if isinstance(cube, iris.cube.CubeList):
111 if len(cube) == 1:
112 return cube[0]
113 raise TypeError("Must have a single cube", cube)
116def _py312_importlib_resources_files_shim():
117 """Importlib behaviour changed in 3.12 to avoid circular dependencies.
119 This shim is needed until python 3.12 is our oldest supported version, after
120 which it can just be replaced by directly using importlib.resources.files.
121 """
122 if sys.version_info.minor >= 12:
123 files = importlib.resources.files()
124 else:
125 import CSET.operators
127 files = importlib.resources.files(CSET.operators)
128 return files
131def _make_plot_html_page(plots: list):
132 """Create a HTML page to display a plot image."""
133 # Debug check that plots actually contains some strings.
134 assert isinstance(plots[0], str)
136 # Load HTML template file.
137 operator_files = _py312_importlib_resources_files_shim()
138 template_file = operator_files.joinpath("_plot_page_template.html")
140 # Get some metadata.
141 meta = get_recipe_metadata()
142 title = meta.get("title", "Untitled")
143 description = MarkdownIt().render(meta.get("description", "*No description.*"))
145 # Prepare template variables.
146 variables = {
147 "title": title,
148 "description": description,
149 "initial_plot": plots[0],
150 "plots": plots,
151 "title_slug": slugify(title),
152 }
154 # Render template.
155 html = render_file(template_file, **variables)
157 # Save completed HTML.
158 with open("index.html", "wt", encoding="UTF-8") as fp:
159 fp.write(html)
162@functools.cache
163def _load_colorbar_map(user_colorbar_file: str = None) -> dict:
164 """Load the colorbar definitions from a file.
166 This is a separate function to make it cacheable.
167 """
168 colorbar_file = _py312_importlib_resources_files_shim().joinpath(
169 "_colorbar_definition.json"
170 )
171 with open(colorbar_file, "rt", encoding="UTF-8") as fp:
172 colorbar = json.load(fp)
174 logging.debug("User colour bar file: %s", user_colorbar_file)
175 override_colorbar = {}
176 if user_colorbar_file:
177 try:
178 with open(user_colorbar_file, "rt", encoding="UTF-8") as fp:
179 override_colorbar = json.load(fp)
180 except FileNotFoundError:
181 logging.warning("Colorbar file does not exist. Using default values.")
183 # Overwrite values with the user supplied colorbar definition.
184 colorbar = combine_dicts(colorbar, override_colorbar)
185 return colorbar
188def _get_model_colors_map(cubes: iris.cube.CubeList | iris.cube.Cube) -> dict:
189 """Get an appropriate colors for model lines in line plots.
191 For each model in the list of cubes colors either from user provided
192 color definition file (so-called style file) or from default colors are mapped
193 to model_name attribute.
195 Parameters
196 ----------
197 cubes: CubeList or Cube
198 Cubes with model_name attribute
200 Returns
201 -------
202 model_colors_map:
203 Dictionary mapping model_name attribute to colors
204 """
205 user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
206 colorbar = _load_colorbar_map(user_colorbar_file)
207 model_names = sorted(
208 filter(
209 lambda x: x is not None,
210 (cube.attributes.get("model_name", None) for cube in iter_maybe(cubes)),
211 )
212 )
213 if not model_names:
214 return {}
215 use_user_colors = all(mname in colorbar.keys() for mname in model_names)
216 if use_user_colors: 216 ↛ 217line 216 didn't jump to line 217 because the condition on line 216 was never true
217 return {mname: colorbar[mname] for mname in model_names}
219 color_list = itertools.cycle(DEFAULT_DISCRETE_COLORS)
220 return {mname: color for mname, color in zip(model_names, color_list, strict=False)}
223def _colorbar_map_levels(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
224 """Get an appropriate colorbar for the given cube.
226 For the given variable the appropriate colorbar is looked up from a
227 combination of the built-in CSET colorbar definitions, and any user supplied
228 definitions. As well as varying on variables, these definitions may also
229 exist for specific pressure levels to account for variables with
230 significantly different ranges at different heights. The colorbars also exist
231 for masks and mask differences for considering variable presence diagnostics.
232 Specific variable ranges can be separately set in user-supplied definition
233 for x- or y-axis limits, or indicate where automated range preferred.
235 Parameters
236 ----------
237 cube: Cube
238 Cube of variable for which the colorbar information is desired.
239 axis: "x", "y", optional
240 Select the levels for just this axis of a line plot. The min and max
241 can be set by xmin/xmax or ymin/ymax respectively. For variables where
242 setting a universal range is not desirable (e.g. temperature), users
243 can set ymin/ymax values to "auto" in the colorbar definitions file.
244 Where no additional xmin/xmax or ymin/ymax values are provided, the
245 axis bounds default to use the vmin/vmax values provided.
247 Returns
248 -------
249 cmap:
250 Matplotlib colormap.
251 levels:
252 List of levels to use for plotting. For continuous plots the min and max
253 should be taken as the range.
254 norm:
255 BoundaryNorm information.
256 """
257 # Grab the colorbar file from the recipe global metadata.
258 user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
259 colorbar = _load_colorbar_map(user_colorbar_file)
260 cmap = None
262 try:
263 # We assume that pressure is a scalar coordinate here.
264 pressure_level_raw = cube.coord("pressure").points[0]
265 # Ensure pressure_level is a string, as it is used as a JSON key.
266 pressure_level = str(int(pressure_level_raw))
267 except iris.exceptions.CoordinateNotFoundError:
268 pressure_level = None
270 # First try long name, then standard name, then var name. This order is used
271 # as long name is the one we correct between models, so it most likely to be
272 # consistent.
273 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name]))
274 for varname in varnames:
275 # Get the colormap for this variable.
276 try:
277 var_colorbar = colorbar[varname]
278 cmap = plt.get_cmap(colorbar[varname]["cmap"], 51)
279 varname_key = varname
280 break
281 except KeyError:
282 logging.debug("Cube name %s has no colorbar definition.", varname)
284 # Get colormap if it is a mask.
285 if any("mask_for_" in name for name in varnames):
286 cmap, levels, norm = _custom_colormap_mask(cube, axis=axis)
287 return cmap, levels, norm
288 # If winds on Beaufort Scale use custom colorbar and levels
289 if any("Beaufort_Scale" in name for name in varnames):
290 cmap, levels, norm = _custom_beaufort_scale(cube, axis=axis)
291 return cmap, levels, norm
292 # If probability is plotted use custom colorbar and levels
293 if any("probability_of_" in name for name in varnames):
294 cmap, levels, norm = _custom_colormap_probability(cube, axis=axis)
295 return cmap, levels, norm
296 # If aviation colour state use custom colorbar and levels
297 if any("aviation_colour_state" in name for name in varnames): 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true
298 cmap, levels, norm = _custom_colormap_aviation_colour_state(cube)
299 return cmap, levels, norm
301 # If no valid colormap has been defined, use defaults and return.
302 if not cmap:
303 logging.warning("No colorbar definition exists for %s.", cube.name())
304 cmap, levels, norm = mpl.colormaps["viridis"], None, None
305 return cmap, levels, norm
307 # Test if pressure-level specific settings are provided for cube.
308 if pressure_level:
309 try:
310 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level]
311 except KeyError:
312 logging.debug(
313 "%s has no colorbar definition for pressure level %s.",
314 varname,
315 pressure_level,
316 )
318 # Check for availability of x-axis or y-axis user-specific overrides
319 # for setting level bounds for line plot types and return just levels.
320 # Line plots do not need a colormap, and just use the data range.
321 if axis:
322 if axis == "x":
323 try:
324 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"]
325 except KeyError:
326 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
327 if axis == "y":
328 try:
329 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"]
330 except KeyError:
331 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
332 # Check if user-specified auto-scaling for this variable
333 if vmin == "auto" or vmax == "auto":
334 levels = None
335 else:
336 levels = [vmin, vmax]
337 return None, levels, None
338 # Get and use the colorbar levels for this variable if spatial or histogram.
339 else:
340 try:
341 levels = var_colorbar["levels"]
342 # Use discrete bins when levels are specified, rather
343 # than a smooth range.
344 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
345 logging.debug("Using levels for %s colorbar.", varname)
346 logging.info("Using levels: %s", levels)
347 except KeyError:
348 # Get the range for this variable.
349 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
350 logging.debug("Using min and max for %s colorbar.", varname)
351 # Calculate levels from range.
352 levels = np.linspace(vmin, vmax, 101)
353 norm = None
355 # Overwrite cmap, levels and norm for specific variables that
356 # require custom colorbar_map as these can not be defined in the
357 # JSON file.
358 cmap, levels, norm = _custom_colourmap_precipitation(cube, cmap, levels, norm)
359 cmap, levels, norm = _custom_colourmap_visibility_in_air(
360 cube, cmap, levels, norm
361 )
362 cmap, levels, norm = _custom_colormap_celsius(cube, cmap, levels, norm)
363 return cmap, levels, norm
366def _setup_spatial_map(
367 cube: iris.cube.Cube,
368 figure,
369 cmap,
370 grid_size: int | None = None,
371 subplot: int | None = None,
372):
373 """Define map projections, extent and add coastlines for spatial plots.
375 For spatial map plots, a relevant map projection for rotated or non-rotated inputs
376 is specified, and map extent defined based on the input data.
378 Parameters
379 ----------
380 cube: Cube
381 2 dimensional (lat and lon) Cube of the data to plot.
382 figure:
383 Matplotlib Figure object holding all plot elements.
384 cmap:
385 Matplotlib colormap.
386 grid_size: int, optional
387 Size of grid for subplots if multiple spatial subplots in figure.
388 subplot: int, optional
389 Subplot index if multiple spatial subplots in figure.
391 Returns
392 -------
393 axes:
394 Matplotlib GeoAxes definition.
395 """
396 # Identify min/max plot bounds.
397 try:
398 lat_axis, lon_axis = get_cube_yxcoordname(cube)
399 x1 = np.min(cube.coord(lon_axis).points)
400 x2 = np.max(cube.coord(lon_axis).points)
401 y1 = np.min(cube.coord(lat_axis).points)
402 y2 = np.max(cube.coord(lat_axis).points)
404 # Adjust bounds within +/- 180.0 if x dimension extends beyond half-globe.
405 if np.abs(x2 - x1) > 180.0:
406 x1 = x1 - 180.0
407 x2 = x2 - 180.0
408 logging.debug("Adjusting plot bounds to fit global extent.")
410 # Consider map projection orientation.
411 # Adapting orientation enables plotting across international dateline.
412 # Users can adapt the default central_longitude if alternative projections views.
413 if x2 > 180.0:
414 central_longitude = 180.0
415 else:
416 central_longitude = 0.0
418 # Define spatial map projection.
419 coord_system = cube.coord(lat_axis).coord_system
420 if isinstance(coord_system, iris.coord_systems.RotatedGeogCS):
421 # Define rotated pole map projection for rotated pole inputs.
422 projection = ccrs.RotatedPole(
423 pole_longitude=coord_system.grid_north_pole_longitude,
424 pole_latitude=coord_system.grid_north_pole_latitude,
425 central_rotated_longitude=0.0,
426 )
427 crs = projection
428 else:
429 # Define regular map projection for non-rotated pole inputs.
430 # Alternatives might include e.g. for global model outputs:
431 # projection=ccrs.Robinson(central_longitude=X.y, globe=None)
432 # See also https://scitools.org.uk/cartopy/docs/v0.15/crs/projections.html.
433 projection = ccrs.PlateCarree(central_longitude=central_longitude)
434 crs = ccrs.PlateCarree()
436 # Define axes for plot (or subplot) with required map projection.
437 if subplot is not None:
438 axes = figure.add_subplot(
439 grid_size, grid_size, subplot, projection=projection
440 )
441 else:
442 axes = figure.add_subplot(projection=projection)
444 # Add coastlines if cube contains x and y map coordinates.
445 if cmap.name in ["viridis", "Greys"]:
446 coastcol = "magenta"
447 else:
448 coastcol = "black"
449 logging.debug("Plotting coastlines in colour %s.", coastcol)
450 axes.coastlines(resolution="10m", color=coastcol)
452 # If is lat/lon spatial map, fix extent to keep plot tight.
453 # Specifying crs within set_extent helps ensure only data region is shown.
454 if isinstance(coord_system, iris.coord_systems.GeogCS):
455 axes.set_extent([x1, x2, y1, y2], crs=crs)
457 except ValueError:
458 # Skip if not both x and y map coordinates.
459 axes = figure.gca()
460 pass
462 return axes
465def _get_plot_resolution() -> int:
466 """Get resolution of rasterised plots in pixels per inch."""
467 return get_recipe_metadata().get("plot_resolution", 100)
470def _plot_and_save_spatial_plot(
471 cube: iris.cube.Cube,
472 filename: str,
473 title: str,
474 method: Literal["contourf", "pcolormesh"],
475 **kwargs,
476):
477 """Plot and save a spatial plot.
479 Parameters
480 ----------
481 cube: Cube
482 2 dimensional (lat and lon) Cube of the data to plot.
483 filename: str
484 Filename of the plot to write.
485 title: str
486 Plot title.
487 method: "contourf" | "pcolormesh"
488 The plotting method to use.
489 """
490 # Setup plot details, size, resolution, etc.
491 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
493 # Specify the color bar
494 cmap, levels, norm = _colorbar_map_levels(cube)
496 # Setup plot map projection, extent and coastlines.
497 axes = _setup_spatial_map(cube, fig, cmap)
499 # Plot the field.
500 if method == "contourf":
501 # Filled contour plot of the field.
502 plot = iplt.contourf(cube, cmap=cmap, levels=levels, norm=norm)
503 elif method == "pcolormesh":
504 try:
505 vmin = min(levels)
506 vmax = max(levels)
507 except TypeError:
508 vmin, vmax = None, None
509 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax
510 # if levels are defined.
511 if norm is not None:
512 vmin = None
513 vmax = None
514 logging.debug("Plotting using defined levels.")
515 plot = iplt.pcolormesh(cube, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax)
516 else:
517 raise ValueError(f"Unknown plotting method: {method}")
519 # Check to see if transect, and if so, adjust y axis.
520 if is_transect(cube):
521 if "pressure" in [coord.name() for coord in cube.coords()]:
522 axes.invert_yaxis()
523 axes.set_yscale("log")
524 axes.set_ylim(1100, 100)
525 # If both model_level_number and level_height exists, iplt can construct
526 # plot as a function of height above orography (NOT sea level).
527 elif {"model_level_number", "level_height"}.issubset( 527 ↛ 532line 527 didn't jump to line 532 because the condition on line 527 was always true
528 {coord.name() for coord in cube.coords()}
529 ):
530 axes.set_yscale("log")
532 axes.set_title(
533 f"{title}\n"
534 f"Start Lat: {cube.attributes['transect_coords'].split('_')[0]}"
535 f" Start Lon: {cube.attributes['transect_coords'].split('_')[1]}"
536 f" End Lat: {cube.attributes['transect_coords'].split('_')[2]}"
537 f" End Lon: {cube.attributes['transect_coords'].split('_')[3]}",
538 fontsize=16,
539 )
541 else:
542 # Add title.
543 axes.set_title(title, fontsize=16)
545 # Add watermark with min/max/mean. Currently not user togglable.
546 # In the bbox dictionary, fc and ec are hex colour codes for grey shade.
547 axes.annotate(
548 f"Min: {np.min(cube.data):.3g} Max: {np.max(cube.data):.3g} Mean: {np.mean(cube.data):.3g}",
549 xy=(1, -0.05),
550 xycoords="axes fraction",
551 xytext=(-5, 5),
552 textcoords="offset points",
553 ha="right",
554 va="bottom",
555 size=11,
556 bbox=dict(boxstyle="round", fc="#cccccc", ec="#808080", alpha=0.9),
557 )
559 # Add colour bar.
560 cbar = fig.colorbar(plot, orientation="horizontal", pad=0.042, shrink=0.7)
561 cbar.set_label(label=f"{cube.name()} ({cube.units})", size=14)
562 # add ticks and tick_labels for every levels if less than 20 levels exist
563 if levels is not None and len(levels) < 20:
564 cbar.set_ticks(levels)
565 cbar.set_ticklabels([f"{level:.2f}" for level in levels])
566 if "visibility" in cube.name(): 566 ↛ 567line 566 didn't jump to line 567 because the condition on line 566 was never true
567 cbar.set_ticklabels([f"{level:.3g}" for level in levels])
568 logging.debug("Set colorbar ticks and labels.")
570 # Save plot.
571 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
572 logging.info("Saved spatial plot to %s", filename)
573 plt.close(fig)
576def _plot_and_save_postage_stamp_spatial_plot(
577 cube: iris.cube.Cube,
578 filename: str,
579 stamp_coordinate: str,
580 title: str,
581 method: Literal["contourf", "pcolormesh"],
582 **kwargs,
583):
584 """Plot postage stamp spatial plots from an ensemble.
586 Parameters
587 ----------
588 cube: Cube
589 Iris cube of data to be plotted. It must have the stamp coordinate.
590 filename: str
591 Filename of the plot to write.
592 stamp_coordinate: str
593 Coordinate that becomes different plots.
594 method: "contourf" | "pcolormesh"
595 The plotting method to use.
597 Raises
598 ------
599 ValueError
600 If the cube doesn't have the right dimensions.
601 """
602 # Use the smallest square grid that will fit the members.
603 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
605 fig = plt.figure(figsize=(10, 10))
607 # Specify the color bar
608 cmap, levels, norm = _colorbar_map_levels(cube)
610 # Make a subplot for each member.
611 for member, subplot in zip(
612 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
613 ):
614 # Setup subplot map projection, extent and coastlines.
615 axes = _setup_spatial_map(
616 member, fig, cmap, grid_size=grid_size, subplot=subplot
617 )
618 if method == "contourf":
619 # Filled contour plot of the field.
620 plot = iplt.contourf(member, cmap=cmap, levels=levels, norm=norm)
621 elif method == "pcolormesh":
622 if levels is not None:
623 vmin = min(levels)
624 vmax = max(levels)
625 else:
626 raise TypeError("Unknown vmin and vmax range.")
627 vmin, vmax = None, None
628 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax
629 # if levels are defined.
630 if norm is not None: 630 ↛ 631line 630 didn't jump to line 631 because the condition on line 630 was never true
631 vmin = None
632 vmax = None
633 # pcolormesh plot of the field.
634 plot = iplt.pcolormesh(member, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax)
635 else:
636 raise ValueError(f"Unknown plotting method: {method}")
637 axes.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
638 axes.set_axis_off()
640 # Put the shared colorbar in its own axes.
641 colorbar_axes = fig.add_axes([0.15, 0.07, 0.7, 0.03])
642 colorbar = fig.colorbar(
643 plot, colorbar_axes, orientation="horizontal", pad=0.042, shrink=0.7
644 )
645 colorbar.set_label(f"{cube.name()} ({cube.units})", size=14)
647 # Overall figure title.
648 fig.suptitle(title, fontsize=16)
650 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
651 logging.info("Saved contour postage stamp plot to %s", filename)
652 plt.close(fig)
655def _plot_and_save_line_series(
656 cubes: iris.cube.CubeList,
657 coords: list[iris.coords.Coord],
658 ensemble_coord: str,
659 filename: str,
660 title: str,
661 **kwargs,
662):
663 """Plot and save a 1D line series.
665 Parameters
666 ----------
667 cubes: Cube or CubeList
668 Cube or CubeList containing the cubes to plot on the y-axis.
669 coords: list[Coord]
670 Coordinates to plot on the x-axis, one per cube.
671 ensemble_coord: str
672 Ensemble coordinate in the cube.
673 filename: str
674 Filename of the plot to write.
675 title: str
676 Plot title.
677 """
678 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
680 model_colors_map = _get_model_colors_map(cubes)
682 # Store min/max ranges.
683 y_levels = []
685 # Check match-up across sequence coords gives consistent sizes
686 _validate_cubes_coords(cubes, coords)
688 for cube, coord in zip(cubes, coords, strict=True):
689 label = None
690 color = "black"
691 if model_colors_map:
692 label = cube.attributes.get("model_name")
693 color = model_colors_map.get(label)
694 for cube_slice in cube.slices_over(ensemble_coord):
695 # Label with (control) if part of an ensemble or not otherwise.
696 if cube_slice.coord(ensemble_coord).points == [0]:
697 iplt.plot(
698 coord,
699 cube_slice,
700 color=color,
701 marker="o",
702 ls="-",
703 lw=3,
704 label=f"{label} (control)"
705 if len(cube.coord(ensemble_coord).points) > 1
706 else label,
707 )
708 # Label with (perturbed) if part of an ensemble and not the control.
709 else:
710 iplt.plot(
711 coord,
712 cube_slice,
713 color=color,
714 ls="-",
715 lw=1.5,
716 alpha=0.75,
717 label=f"{label} (member)",
718 )
720 # Calculate the global min/max if multiple cubes are given.
721 _, levels, _ = _colorbar_map_levels(cube, axis="y")
722 if levels is not None: 722 ↛ 723line 722 didn't jump to line 723 because the condition on line 722 was never true
723 y_levels.append(min(levels))
724 y_levels.append(max(levels))
726 # Get the current axes.
727 ax = plt.gca()
729 # Add some labels and tweak the style.
730 # check if cubes[0] works for single cube if not CubeList
731 ax.set_xlabel(f"{coords[0].name()} / {coords[0].units}", fontsize=14)
732 ax.set_ylabel(f"{cubes[0].name()} / {cubes[0].units}", fontsize=14)
733 ax.set_title(title, fontsize=16)
735 ax.ticklabel_format(axis="y", useOffset=False)
736 ax.tick_params(axis="x", labelrotation=15)
737 ax.tick_params(axis="both", labelsize=12)
739 # Set y limits to global min and max, autoscale if colorbar doesn't exist.
740 if y_levels: 740 ↛ 741line 740 didn't jump to line 741 because the condition on line 740 was never true
741 ax.set_ylim(min(y_levels), max(y_levels))
742 # Add zero line.
743 if min(y_levels) < 0.0 and max(y_levels) > 0.0:
744 ax.axhline(y=0, xmin=0, xmax=1, ls="-", color="grey", lw=2)
745 logging.debug(
746 "Line plot with y-axis limits %s-%s", min(y_levels), max(y_levels)
747 )
748 else:
749 ax.autoscale()
751 # Add gridlines
752 ax.grid(linestyle="--", color="grey", linewidth=1)
753 # Ientify unique labels for legend
754 handles = list(
755 {
756 label: handle
757 for (handle, label) in zip(*ax.get_legend_handles_labels(), strict=True)
758 }.values()
759 )
760 ax.legend(handles=handles, loc="best", ncol=1, frameon=False, fontsize=16)
762 # Save plot.
763 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
764 logging.info("Saved line plot to %s", filename)
765 plt.close(fig)
768def _plot_and_save_vertical_line_series(
769 cubes: iris.cube.CubeList,
770 coords: list[iris.coords.Coord],
771 ensemble_coord: str,
772 filename: str,
773 series_coordinate: str,
774 title: str,
775 vmin: float,
776 vmax: float,
777 **kwargs,
778):
779 """Plot and save a 1D line series in vertical.
781 Parameters
782 ----------
783 cubes: CubeList
784 1 dimensional Cube or CubeList of the data to plot on x-axis.
785 coord: list[Coord]
786 Coordinates to plot on the y-axis, one per cube.
787 ensemble_coord: str
788 Ensemble coordinate in the cube.
789 filename: str
790 Filename of the plot to write.
791 series_coordinate: str
792 Coordinate to use as vertical axis.
793 title: str
794 Plot title.
795 vmin: float
796 Minimum value for the x-axis.
797 vmax: float
798 Maximum value for the x-axis.
799 """
800 # plot the vertical pressure axis using log scale
801 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
803 model_colors_map = _get_model_colors_map(cubes)
805 # Check match-up across sequence coords gives consistent sizes
806 _validate_cubes_coords(cubes, coords)
808 for cube, coord in zip(cubes, coords, strict=True):
809 label = None
810 color = "black"
811 if model_colors_map: 811 ↛ 812line 811 didn't jump to line 812 because the condition on line 811 was never true
812 label = cube.attributes.get("model_name")
813 color = model_colors_map.get(label)
815 for cube_slice in cube.slices_over(ensemble_coord):
816 # If ensemble data given plot control member with (control)
817 # unless single forecast.
818 if cube_slice.coord(ensemble_coord).points == [0]:
819 iplt.plot(
820 cube_slice,
821 coord,
822 color=color,
823 marker="o",
824 ls="-",
825 lw=3,
826 label=f"{label} (control)"
827 if len(cube.coord(ensemble_coord).points) > 1
828 else label,
829 )
830 # If ensemble data given plot perturbed members with (perturbed).
831 else:
832 iplt.plot(
833 cube_slice,
834 coord,
835 color=color,
836 ls="-",
837 lw=1.5,
838 alpha=0.75,
839 label=f"{label} (member)",
840 )
842 # Get the current axis
843 ax = plt.gca()
845 # Special handling for pressure level data.
846 if series_coordinate == "pressure": 846 ↛ 868line 846 didn't jump to line 868 because the condition on line 846 was always true
847 # Invert y-axis and set to log scale.
848 ax.invert_yaxis()
849 ax.set_yscale("log")
851 # Define y-ticks and labels for pressure log axis.
852 y_tick_labels = [
853 "1000",
854 "850",
855 "700",
856 "500",
857 "300",
858 "200",
859 "100",
860 ]
861 y_ticks = [1000, 850, 700, 500, 300, 200, 100]
863 # Set y-axis limits and ticks.
864 ax.set_ylim(1100, 100)
866 # Test if series_coordinate is model level data. The UM data uses
867 # model_level_number and lfric uses full_levels as coordinate.
868 elif series_coordinate in ("model_level_number", "full_levels", "half_levels"):
869 # Define y-ticks and labels for vertical axis.
870 y_ticks = iter_maybe(cubes)[0].coord(series_coordinate).points
871 y_tick_labels = [str(int(i)) for i in y_ticks]
872 ax.set_ylim(min(y_ticks), max(y_ticks))
874 ax.set_yticks(y_ticks)
875 ax.set_yticklabels(y_tick_labels)
877 # Set x-axis limits.
878 ax.set_xlim(vmin, vmax)
879 # Mark y=0 if present in plot.
880 if vmin < 0.0 and vmax > 0.0: 880 ↛ 881line 880 didn't jump to line 881 because the condition on line 880 was never true
881 ax.axvline(x=0, ymin=0, ymax=1, ls="-", color="grey", lw=2)
883 # Add some labels and tweak the style.
884 ax.set_ylabel(f"{coord.name()} / {coord.units}", fontsize=14)
885 ax.set_xlabel(
886 f"{iter_maybe(cubes)[0].name()} / {iter_maybe(cubes)[0].units}", fontsize=14
887 )
888 ax.set_title(title, fontsize=16)
889 ax.ticklabel_format(axis="x")
890 ax.tick_params(axis="y")
891 ax.tick_params(axis="both", labelsize=12)
893 # Add gridlines
894 ax.grid(linestyle="--", color="grey", linewidth=1)
895 # Ientify unique labels for legend
896 handles = list(
897 {
898 label: handle
899 for (handle, label) in zip(*ax.get_legend_handles_labels(), strict=True)
900 }.values()
901 )
902 ax.legend(handles=handles, loc="best", ncol=1, frameon=False, fontsize=16)
904 # Save plot.
905 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
906 logging.info("Saved line plot to %s", filename)
907 plt.close(fig)
910def _plot_and_save_scatter_plot(
911 cube_x: iris.cube.Cube | iris.cube.CubeList,
912 cube_y: iris.cube.Cube | iris.cube.CubeList,
913 filename: str,
914 title: str,
915 one_to_one: bool,
916 model_names: list[str] = None,
917 **kwargs,
918):
919 """Plot and save a 2D scatter plot.
921 Parameters
922 ----------
923 cube_x: Cube | CubeList
924 1 dimensional Cube or CubeList of the data to plot on x-axis.
925 cube_y: Cube | CubeList
926 1 dimensional Cube or CubeList of the data to plot on y-axis.
927 filename: str
928 Filename of the plot to write.
929 title: str
930 Plot title.
931 one_to_one: bool
932 Whether a 1:1 line is plotted.
933 """
934 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
935 # plot the cube_x and cube_y 1D fields as a scatter plot. If they are CubeLists this ensures
936 # to pair each cube from cube_x with the corresponding cube from cube_y, allowing to iterate
937 # over the pairs simultaneously.
939 # Ensure cube_x and cube_y are iterable
940 cube_x_iterable = iter_maybe(cube_x)
941 cube_y_iterable = iter_maybe(cube_y)
943 for cube_x_iter, cube_y_iter in zip(cube_x_iterable, cube_y_iterable, strict=True):
944 iplt.scatter(cube_x_iter, cube_y_iter)
945 if one_to_one is True:
946 plt.plot(
947 [
948 np.nanmin([np.nanmin(cube_y.data), np.nanmin(cube_x.data)]),
949 np.nanmax([np.nanmax(cube_y.data), np.nanmax(cube_x.data)]),
950 ],
951 [
952 np.nanmin([np.nanmin(cube_y.data), np.nanmin(cube_x.data)]),
953 np.nanmax([np.nanmax(cube_y.data), np.nanmax(cube_x.data)]),
954 ],
955 "k",
956 linestyle="--",
957 )
958 ax = plt.gca()
960 # Add some labels and tweak the style.
961 if model_names is None:
962 ax.set_xlabel(f"{cube_x[0].name()} / {cube_x[0].units}", fontsize=14)
963 ax.set_ylabel(f"{cube_y[0].name()} / {cube_y[0].units}", fontsize=14)
964 else:
965 # Add the model names, these should be order of base (x) and other (y).
966 ax.set_xlabel(
967 f"{model_names[0]}_{cube_x[0].name()} / {cube_x[0].units}", fontsize=14
968 )
969 ax.set_ylabel(
970 f"{model_names[1]}_{cube_y[0].name()} / {cube_y[0].units}", fontsize=14
971 )
972 ax.set_title(title, fontsize=16)
973 ax.ticklabel_format(axis="y", useOffset=False)
974 ax.tick_params(axis="x", labelrotation=15)
975 ax.tick_params(axis="both", labelsize=12)
976 ax.autoscale()
978 # Save plot.
979 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
980 logging.info("Saved scatter plot to %s", filename)
981 plt.close(fig)
984def _plot_and_save_vector_plot(
985 cube_u: iris.cube.Cube,
986 cube_v: iris.cube.Cube,
987 filename: str,
988 title: str,
989 method: Literal["contourf", "pcolormesh"],
990 **kwargs,
991):
992 """Plot and save a 2D vector plot.
994 Parameters
995 ----------
996 cube_u: Cube
997 2 dimensional Cube of u component of the data.
998 cube_v: Cube
999 2 dimensional Cube of v component of the data.
1000 filename: str
1001 Filename of the plot to write.
1002 title: str
1003 Plot title.
1004 """
1005 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1007 # Create a cube containing the magnitude of the vector field.
1008 cube_vec_mag = (cube_u**2 + cube_v**2) ** 0.5
1009 cube_vec_mag.rename(f"{cube_u.name()}_{cube_v.name()}_magnitude")
1011 # Specify the color bar
1012 cmap, levels, norm = _colorbar_map_levels(cube_vec_mag)
1014 # Setup plot map projection, extent and coastlines.
1015 axes = _setup_spatial_map(cube_vec_mag, fig, cmap)
1017 if method == "contourf": 1017 ↛ 1020line 1017 didn't jump to line 1020 because the condition on line 1017 was always true
1018 # Filled contour plot of the field.
1019 plot = iplt.contourf(cube_vec_mag, cmap=cmap, levels=levels, norm=norm)
1020 elif method == "pcolormesh":
1021 try:
1022 vmin = min(levels)
1023 vmax = max(levels)
1024 except TypeError:
1025 vmin, vmax = None, None
1026 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax
1027 # if levels are defined.
1028 if norm is not None:
1029 vmin = None
1030 vmax = None
1031 plot = iplt.pcolormesh(cube_vec_mag, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax)
1032 else:
1033 raise ValueError(f"Unknown plotting method: {method}")
1035 # Check to see if transect, and if so, adjust y axis.
1036 if is_transect(cube_vec_mag): 1036 ↛ 1037line 1036 didn't jump to line 1037 because the condition on line 1036 was never true
1037 if "pressure" in [coord.name() for coord in cube_vec_mag.coords()]:
1038 axes.invert_yaxis()
1039 axes.set_yscale("log")
1040 axes.set_ylim(1100, 100)
1041 # If both model_level_number and level_height exists, iplt can construct
1042 # plot as a function of height above orography (NOT sea level).
1043 elif {"model_level_number", "level_height"}.issubset(
1044 {coord.name() for coord in cube_vec_mag.coords()}
1045 ):
1046 axes.set_yscale("log")
1048 axes.set_title(
1049 f"{title}\n"
1050 f"Start Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[0]}"
1051 f" Start Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[1]}"
1052 f" End Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[2]}"
1053 f" End Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[3]}",
1054 fontsize=16,
1055 )
1057 else:
1058 # Add title.
1059 axes.set_title(title, fontsize=16)
1061 # Add watermark with min/max/mean. Currently not user togglable.
1062 # In the bbox dictionary, fc and ec are hex colour codes for grey shade.
1063 axes.annotate(
1064 f"Min: {np.min(cube_vec_mag.data):.3g} Max: {np.max(cube_vec_mag.data):.3g} Mean: {np.mean(cube_vec_mag.data):.3g}",
1065 xy=(1, -0.05),
1066 xycoords="axes fraction",
1067 xytext=(-5, 5),
1068 textcoords="offset points",
1069 ha="right",
1070 va="bottom",
1071 size=11,
1072 bbox=dict(boxstyle="round", fc="#cccccc", ec="#808080", alpha=0.9),
1073 )
1075 # Add colour bar.
1076 cbar = fig.colorbar(plot, orientation="horizontal", pad=0.042, shrink=0.7)
1077 cbar.set_label(label=f"{cube_vec_mag.name()} ({cube_vec_mag.units})", size=14)
1078 # add ticks and tick_labels for every levels if less than 20 levels exist
1079 if levels is not None and len(levels) < 20: 1079 ↛ 1080line 1079 didn't jump to line 1080 because the condition on line 1079 was never true
1080 cbar.set_ticks(levels)
1081 cbar.set_ticklabels([f"{level:.1f}" for level in levels])
1083 # 30 barbs along the longest axis of the plot, or a barb per point for data
1084 # with less than 30 points.
1085 step = max(max(cube_u.shape) // 30, 1)
1086 iplt.quiver(cube_u[::step, ::step], cube_v[::step, ::step], pivot="middle")
1088 # Save plot.
1089 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1090 logging.info("Saved vector plot to %s", filename)
1091 plt.close(fig)
1094def _plot_and_save_histogram_series(
1095 cubes: iris.cube.Cube | iris.cube.CubeList,
1096 filename: str,
1097 title: str,
1098 vmin: float,
1099 vmax: float,
1100 **kwargs,
1101):
1102 """Plot and save a histogram series.
1104 Parameters
1105 ----------
1106 cubes: Cube or CubeList
1107 2 dimensional Cube or CubeList of the data to plot as histogram.
1108 filename: str
1109 Filename of the plot to write.
1110 title: str
1111 Plot title.
1112 vmin: float
1113 minimum for colorbar
1114 vmax: float
1115 maximum for colorbar
1116 """
1117 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1118 ax = plt.gca()
1120 model_colors_map = _get_model_colors_map(cubes)
1122 # Set default that histograms will produce probability density function
1123 # at each bin (integral over range sums to 1).
1124 density = True
1126 for cube in iter_maybe(cubes):
1127 # Easier to check title (where var name originates)
1128 # than seeing if long names exist etc.
1129 # Exception case, where distribution better fits log scales/bins.
1130 if "surface_microphysical" in title:
1131 if "amount" in title: 1131 ↛ 1133line 1131 didn't jump to line 1133 because the condition on line 1131 was never true
1132 # Compute histogram following Klingaman et al. (2017): ASoP
1133 bin2 = np.exp(np.log(0.02) + 0.1 * np.linspace(0, 99, 100))
1134 bins = np.pad(bin2, (1, 0), "constant", constant_values=0)
1135 density = False
1136 else:
1137 bins = 10.0 ** (
1138 np.arange(-10, 27, 1) / 10.0
1139 ) # Suggestion from RMED toolbox.
1140 bins = np.insert(bins, 0, 0)
1141 ax.set_yscale("log")
1142 vmin = bins[1]
1143 vmax = bins[-1] # Manually set vmin/vmax to override json derived value.
1144 ax.set_xscale("log")
1145 elif "lightning" in title:
1146 bins = [0, 1, 2, 3, 4, 5]
1147 else:
1148 bins = np.linspace(vmin, vmax, 51)
1149 logging.debug(
1150 "Plotting histogram with %s bins %s - %s.",
1151 np.size(bins),
1152 np.min(bins),
1153 np.max(bins),
1154 )
1156 # Reshape cube data into a single array to allow for a single histogram.
1157 # Otherwise we plot xdim histograms stacked.
1158 cube_data_1d = (cube.data).flatten()
1160 label = None
1161 color = "black"
1162 if model_colors_map: 1162 ↛ 1163line 1162 didn't jump to line 1163 because the condition on line 1162 was never true
1163 label = cube.attributes.get("model_name")
1164 color = model_colors_map[label]
1165 x, y = np.histogram(cube_data_1d, bins=bins, density=density)
1167 # Compute area under curve.
1168 if "surface_microphysical" in title and "amount" in title: 1168 ↛ 1169line 1168 didn't jump to line 1169 because the condition on line 1168 was never true
1169 bin_mean = (bins[:-1] + bins[1:]) / 2.0
1170 x = x * bin_mean / x.sum()
1171 x = x[1:]
1172 y = y[1:]
1174 ax.plot(
1175 y[:-1], x, color=color, linewidth=3, marker="o", markersize=6, label=label
1176 )
1178 # Add some labels and tweak the style.
1179 ax.set_title(title, fontsize=16)
1180 ax.set_xlabel(
1181 f"{iter_maybe(cubes)[0].name()} / {iter_maybe(cubes)[0].units}", fontsize=14
1182 )
1183 ax.set_ylabel("Normalised probability density", fontsize=14)
1184 if "surface_microphysical" in title and "amount" in title: 1184 ↛ 1185line 1184 didn't jump to line 1185 because the condition on line 1184 was never true
1185 ax.set_ylabel(
1186 f"Contribution to mean ({iter_maybe(cubes)[0].units})", fontsize=14
1187 )
1188 ax.set_xlim(vmin, vmax)
1189 ax.tick_params(axis="both", labelsize=12)
1191 # Overlay grid-lines onto histogram plot.
1192 ax.grid(linestyle="--", color="grey", linewidth=1)
1193 if model_colors_map: 1193 ↛ 1194line 1193 didn't jump to line 1194 because the condition on line 1193 was never true
1194 ax.legend(loc="best", ncol=1, frameon=False, fontsize=16)
1196 # Save plot.
1197 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1198 logging.info("Saved line plot to %s", filename)
1199 plt.close(fig)
1202def _plot_and_save_postage_stamp_histogram_series(
1203 cube: iris.cube.Cube,
1204 filename: str,
1205 title: str,
1206 stamp_coordinate: str,
1207 vmin: float,
1208 vmax: float,
1209 **kwargs,
1210):
1211 """Plot and save postage (ensemble members) stamps for a histogram series.
1213 Parameters
1214 ----------
1215 cube: Cube
1216 2 dimensional Cube of the data to plot as histogram.
1217 filename: str
1218 Filename of the plot to write.
1219 title: str
1220 Plot title.
1221 stamp_coordinate: str
1222 Coordinate that becomes different plots.
1223 vmin: float
1224 minimum for pdf x-axis
1225 vmax: float
1226 maximum for pdf x-axis
1227 """
1228 # Use the smallest square grid that will fit the members.
1229 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
1231 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1232 # Make a subplot for each member.
1233 for member, subplot in zip(
1234 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
1235 ):
1236 # Implicit interface is much easier here, due to needing to have the
1237 # cartopy GeoAxes generated.
1238 plt.subplot(grid_size, grid_size, subplot)
1239 # Reshape cube data into a single array to allow for a single histogram.
1240 # Otherwise we plot xdim histograms stacked.
1241 member_data_1d = (member.data).flatten()
1242 plt.hist(member_data_1d, density=True, stacked=True)
1243 ax = plt.gca()
1244 ax.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
1245 ax.set_xlim(vmin, vmax)
1247 # Overall figure title.
1248 fig.suptitle(title, fontsize=16)
1250 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1251 logging.info("Saved histogram postage stamp plot to %s", filename)
1252 plt.close(fig)
1255def _plot_and_save_postage_stamps_in_single_plot_histogram_series(
1256 cube: iris.cube.Cube,
1257 filename: str,
1258 title: str,
1259 stamp_coordinate: str,
1260 vmin: float,
1261 vmax: float,
1262 **kwargs,
1263):
1264 fig, ax = plt.subplots(figsize=(10, 10), facecolor="w", edgecolor="k")
1265 ax.set_title(title, fontsize=16)
1266 ax.set_xlim(vmin, vmax)
1267 ax.set_xlabel(f"{cube.name()} / {cube.units}", fontsize=14)
1268 ax.set_ylabel("normalised probability density", fontsize=14)
1269 # Loop over all slices along the stamp_coordinate
1270 for member in cube.slices_over(stamp_coordinate):
1271 # Flatten the member data to 1D
1272 member_data_1d = member.data.flatten()
1273 # Plot the histogram using plt.hist
1274 plt.hist(
1275 member_data_1d,
1276 density=True,
1277 stacked=True,
1278 label=f"Member #{member.coord(stamp_coordinate).points[0]}",
1279 )
1281 # Add a legend
1282 ax.legend(fontsize=16)
1284 # Save the figure to a file
1285 plt.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1287 # Close the figure
1288 plt.close(fig)
1291def _plot_and_save_power_spectrum_series(
1292 cubes: iris.cube.Cube | iris.cube.CubeList,
1293 filename: str,
1294 title: str,
1295 **kwargs,
1296):
1297 """Plot and save a power spectrum series.
1299 Parameters
1300 ----------
1301 cubes: Cube or CubeList
1302 2 dimensional Cube or CubeList of the data to plot as power spectrum.
1303 filename: str
1304 Filename of the plot to write.
1305 title: str
1306 Plot title.
1307 """
1308 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1309 ax = plt.gca()
1311 model_colors_map = _get_model_colors_map(cubes)
1313 for cube in iter_maybe(cubes):
1314 # Calculate power spectrum
1316 # Extract time coordinate and convert to datetime
1317 time_coord = cube.coord("time")
1318 time_points = time_coord.units.num2date(time_coord.points)
1320 # Choose one time point (e.g., the first one)
1321 target_time = time_points[0]
1323 # Bind target_time inside the lambda using a default argument
1324 time_constraint = iris.Constraint(
1325 time=lambda cell, target_time=target_time: cell.point == target_time
1326 )
1328 cube = cube.extract(time_constraint)
1330 if cube.ndim == 2:
1331 cube_3d = cube.data[np.newaxis, :, :]
1332 logging.debug("Adding in new axis for a 2 dimensional cube.")
1333 elif cube.ndim == 3: 1333 ↛ 1334line 1333 didn't jump to line 1334 because the condition on line 1333 was never true
1334 cube_3d = cube.data
1335 else:
1336 raise ValueError("Cube dimensions unsuitable for power spectra code")
1337 raise ValueError(
1338 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional."
1339 )
1341 # Calculate spectra
1342 ps_array = _DCT_ps(cube_3d)
1344 ps_cube = iris.cube.Cube(
1345 ps_array,
1346 long_name="power_spectra",
1347 )
1349 ps_cube.attributes["model_name"] = cube.attributes.get("model_name")
1351 # Create a frequency/wavelength array for coordinate
1352 ps_len = ps_cube.data.shape[1]
1353 freqs = np.arange(1, ps_len + 1)
1354 freq_coord = iris.coords.DimCoord(freqs, long_name="frequency", units="1")
1356 # Convert datetime to numeric time using original units
1357 numeric_time = time_coord.units.date2num(time_points)
1358 # Create a new DimCoord with numeric time
1359 new_time_coord = iris.coords.DimCoord(
1360 numeric_time, standard_name="time", units=time_coord.units
1361 )
1363 # Add time and frequency coordinate to spectra cube.
1364 ps_cube.add_dim_coord(new_time_coord.copy(), 0)
1365 ps_cube.add_dim_coord(freq_coord.copy(), 1)
1367 # Extract data from the cube
1368 frequency = ps_cube.coord("frequency").points
1369 power_spectrum = ps_cube.data
1371 label = None
1372 color = "black"
1373 if model_colors_map: 1373 ↛ 1374line 1373 didn't jump to line 1374 because the condition on line 1373 was never true
1374 label = ps_cube.attributes.get("model_name")
1375 color = model_colors_map[label]
1376 ax.plot(frequency, power_spectrum[0], color=color, label=label)
1378 # Add some labels and tweak the style.
1379 ax.set_title(title, fontsize=16)
1380 ax.set_xlabel("Wavenumber", fontsize=14)
1381 ax.set_ylabel("Power", fontsize=14)
1382 ax.tick_params(axis="both", labelsize=12)
1384 # Set log-log scale
1385 ax.set_xscale("log")
1386 ax.set_yscale("log")
1388 # Overlay grid-lines onto power spectrum plot.
1389 ax.grid(linestyle="--", color="grey", linewidth=1)
1390 if model_colors_map: 1390 ↛ 1391line 1390 didn't jump to line 1391 because the condition on line 1390 was never true
1391 ax.legend(loc="best", ncol=1, frameon=False, fontsize=16)
1393 # Save plot.
1394 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1395 logging.info("Saved line plot to %s", filename)
1396 plt.close(fig)
1399def _plot_and_save_postage_stamp_power_spectrum_series(
1400 cube: iris.cube.Cube,
1401 filename: str,
1402 title: str,
1403 stamp_coordinate: str,
1404 **kwargs,
1405):
1406 """Plot and save postage (ensemble members) stamps for a power spectrum series.
1408 Parameters
1409 ----------
1410 cube: Cube
1411 2 dimensional Cube of the data to plot as power spectrum.
1412 filename: str
1413 Filename of the plot to write.
1414 title: str
1415 Plot title.
1416 stamp_coordinate: str
1417 Coordinate that becomes different plots.
1418 """
1419 # Use the smallest square grid that will fit the members.
1420 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
1422 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1423 # Make a subplot for each member.
1424 for member, subplot in zip(
1425 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
1426 ):
1427 # Implicit interface is much easier here, due to needing to have the
1428 # cartopy GeoAxes generated.
1429 plt.subplot(grid_size, grid_size, subplot)
1431 frequency = member.coord("frequency").points
1432 power_spectrum = member.data
1434 ax = plt.gca()
1435 ax.plot(frequency, power_spectrum[0])
1436 ax.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
1438 # Overall figure title.
1439 fig.suptitle(title, fontsize=16)
1441 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1442 logging.info("Saved power spectra postage stamp plot to %s", filename)
1443 plt.close(fig)
1446def _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series(
1447 cube: iris.cube.Cube,
1448 filename: str,
1449 title: str,
1450 stamp_coordinate: str,
1451 **kwargs,
1452):
1453 fig, ax = plt.subplots(figsize=(10, 10), facecolor="w", edgecolor="k")
1454 ax.set_title(title, fontsize=16)
1455 ax.set_xlabel(f"{cube.name()} / {cube.units}", fontsize=14)
1456 ax.set_ylabel("Power", fontsize=14)
1457 # Loop over all slices along the stamp_coordinate
1458 for member in cube.slices_over(stamp_coordinate):
1459 frequency = member.coord("frequency").points
1460 power_spectrum = member.data
1461 ax.plot(
1462 frequency,
1463 power_spectrum[0],
1464 label=f"Member #{member.coord(stamp_coordinate).points[0]}",
1465 )
1467 # Add a legend
1468 ax.legend(fontsize=16)
1470 # Save the figure to a file
1471 plt.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1473 # Close the figure
1474 plt.close(fig)
1477def _spatial_plot(
1478 method: Literal["contourf", "pcolormesh"],
1479 cube: iris.cube.Cube,
1480 filename: str | None,
1481 sequence_coordinate: str,
1482 stamp_coordinate: str,
1483):
1484 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
1486 A 2D spatial field can be plotted, but if the sequence_coordinate is present
1487 then a sequence of plots will be produced. Similarly if the stamp_coordinate
1488 is present then postage stamp plots will be produced.
1490 Parameters
1491 ----------
1492 method: "contourf" | "pcolormesh"
1493 The plotting method to use.
1494 cube: Cube
1495 Iris cube of the data to plot. It should have two spatial dimensions,
1496 such as lat and lon, and may also have a another two dimension to be
1497 plotted sequentially and/or as postage stamp plots.
1498 filename: str | None
1499 Name of the plot to write, used as a prefix for plot sequences. If None
1500 uses the recipe name.
1501 sequence_coordinate: str
1502 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
1503 This coordinate must exist in the cube.
1504 stamp_coordinate: str
1505 Coordinate about which to plot postage stamp plots. Defaults to
1506 ``"realization"``.
1508 Raises
1509 ------
1510 ValueError
1511 If the cube doesn't have the right dimensions.
1512 TypeError
1513 If the cube isn't a single cube.
1514 """
1515 recipe_title = get_recipe_metadata().get("title", "Untitled")
1517 # Ensure we have a name for the plot file.
1518 if filename is None:
1519 filename = slugify(recipe_title)
1521 # Ensure we've got a single cube.
1522 cube = _check_single_cube(cube)
1524 # Make postage stamp plots if stamp_coordinate exists and has more than a
1525 # single point.
1526 plotting_func = _plot_and_save_spatial_plot
1527 try:
1528 if cube.coord(stamp_coordinate).shape[0] > 1:
1529 plotting_func = _plot_and_save_postage_stamp_spatial_plot
1530 except iris.exceptions.CoordinateNotFoundError:
1531 pass
1533 # Must have a sequence coordinate.
1534 try:
1535 cube.coord(sequence_coordinate)
1536 except iris.exceptions.CoordinateNotFoundError as err:
1537 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
1539 # Create a plot for each value of the sequence coordinate.
1540 plot_index = []
1541 nplot = np.size(cube.coord(sequence_coordinate).points)
1542 for cube_slice in cube.slices_over(sequence_coordinate):
1543 # Use sequence value so multiple sequences can merge.
1544 sequence_value = cube_slice.coord(sequence_coordinate).points[0]
1545 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
1546 coord = cube_slice.coord(sequence_coordinate)
1547 # Format the coordinate value in a unit appropriate way.
1548 title = f"{recipe_title}\n [{coord.units.title(coord.points[0])}]"
1549 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs
1550 if nplot == 1 and coord.has_bounds:
1551 if np.size(coord.bounds) > 1:
1552 title = f"{recipe_title}\n [{coord.units.title(coord.bounds[0][0])} to {coord.units.title(coord.bounds[0][1])}]"
1553 # Do the actual plotting.
1554 plotting_func(
1555 cube_slice,
1556 filename=plot_filename,
1557 stamp_coordinate=stamp_coordinate,
1558 title=title,
1559 method=method,
1560 )
1561 plot_index.append(plot_filename)
1563 # Add list of plots to plot metadata.
1564 complete_plot_index = _append_to_plot_index(plot_index)
1566 # Make a page to display the plots.
1567 _make_plot_html_page(complete_plot_index)
1570def _custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
1571 """Get colourmap for mask.
1573 If "mask_for_" appears anywhere in the name of a cube this function will be called
1574 regardless of the name of the variable to ensure a consistent plot.
1576 Parameters
1577 ----------
1578 cube: Cube
1579 Cube of variable for which the colorbar information is desired.
1580 axis: "x", "y", optional
1581 Select the levels for just this axis of a line plot. The min and max
1582 can be set by xmin/xmax or ymin/ymax respectively. For variables where
1583 setting a universal range is not desirable (e.g. temperature), users
1584 can set ymin/ymax values to "auto" in the colorbar definitions file.
1585 Where no additional xmin/xmax or ymin/ymax values are provided, the
1586 axis bounds default to use the vmin/vmax values provided.
1588 Returns
1589 -------
1590 cmap:
1591 Matplotlib colormap.
1592 levels:
1593 List of levels to use for plotting. For continuous plots the min and max
1594 should be taken as the range.
1595 norm:
1596 BoundaryNorm information.
1597 """
1598 if "difference" not in cube.long_name:
1599 if axis:
1600 levels = [0, 1]
1601 # Complete settings based on levels.
1602 return None, levels, None
1603 else:
1604 # Define the levels and colors.
1605 levels = [0, 1, 2]
1606 colors = ["white", "dodgerblue"]
1607 # Create a custom color map.
1608 cmap = mcolors.ListedColormap(colors)
1609 # Normalize the levels.
1610 norm = mcolors.BoundaryNorm(levels, cmap.N)
1611 logging.debug("Colourmap for %s.", cube.long_name)
1612 return cmap, levels, norm
1613 else:
1614 if axis:
1615 levels = [-1, 1]
1616 return None, levels, None
1617 else:
1618 # Search for if mask difference, set to +/- 0.5 as values plotted <
1619 # not <=.
1620 levels = [-2, -0.5, 0.5, 2]
1621 colors = ["goldenrod", "white", "teal"]
1622 cmap = mcolors.ListedColormap(colors)
1623 norm = mcolors.BoundaryNorm(levels, cmap.N)
1624 logging.debug("Colourmap for %s.", cube.long_name)
1625 return cmap, levels, norm
1628def _custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
1629 """Get a custom colorbar for a cube in the Beaufort Scale.
1631 Specific variable ranges can be separately set in user-supplied definition
1632 for x- or y-axis limits, or indicate where automated range preferred.
1634 Parameters
1635 ----------
1636 cube: Cube
1637 Cube of variable with Beaufort Scale in name.
1638 axis: "x", "y", optional
1639 Select the levels for just this axis of a line plot. The min and max
1640 can be set by xmin/xmax or ymin/ymax respectively. For variables where
1641 setting a universal range is not desirable (e.g. temperature), users
1642 can set ymin/ymax values to "auto" in the colorbar definitions file.
1643 Where no additional xmin/xmax or ymin/ymax values are provided, the
1644 axis bounds default to use the vmin/vmax values provided.
1646 Returns
1647 -------
1648 cmap:
1649 Matplotlib colormap.
1650 levels:
1651 List of levels to use for plotting. For continuous plots the min and max
1652 should be taken as the range.
1653 norm:
1654 BoundaryNorm information.
1655 """
1656 if "difference" not in cube.long_name:
1657 if axis:
1658 levels = [0, 12]
1659 return None, levels, None
1660 else:
1661 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
1662 colors = [
1663 "black",
1664 (0, 0, 0.6),
1665 "blue",
1666 "cyan",
1667 "green",
1668 "yellow",
1669 (1, 0.5, 0),
1670 "red",
1671 "pink",
1672 "magenta",
1673 "purple",
1674 "maroon",
1675 "white",
1676 ]
1677 cmap = mcolors.ListedColormap(colors)
1678 norm = mcolors.BoundaryNorm(levels, cmap.N)
1679 logging.info("change colormap for Beaufort Scale colorbar.")
1680 return cmap, levels, norm
1681 else:
1682 if axis:
1683 levels = [-4, 4]
1684 return None, levels, None
1685 else:
1686 levels = [
1687 -3.5,
1688 -2.5,
1689 -1.5,
1690 -0.5,
1691 0.5,
1692 1.5,
1693 2.5,
1694 3.5,
1695 ]
1696 cmap = plt.get_cmap("bwr", 8)
1697 norm = mcolors.BoundaryNorm(levels, cmap.N)
1698 return cmap, levels, norm
1701def _custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm):
1702 """Return altered colourmap for temperature with change in units to Celsius.
1704 If "Celsius" appears anywhere in the name of a cube this function will be called.
1706 Parameters
1707 ----------
1708 cube: Cube
1709 Cube of variable for which the colorbar information is desired.
1710 cmap: Matplotlib colormap.
1711 levels: List
1712 List of levels to use for plotting. For continuous plots the min and max
1713 should be taken as the range.
1714 norm: BoundaryNorm.
1716 Returns
1717 -------
1718 cmap: Matplotlib colormap.
1719 levels: List
1720 List of levels to use for plotting. For continuous plots the min and max
1721 should be taken as the range.
1722 norm: BoundaryNorm.
1723 """
1724 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
1725 if any("temperature" in name for name in varnames) and "Celsius" == cube.units:
1726 levels = np.array(levels)
1727 levels -= 273
1728 levels = levels.tolist()
1729 else:
1730 # Do nothing keep the existing colourbar attributes
1731 levels = levels
1732 cmap = cmap
1733 norm = norm
1734 return cmap, levels, norm
1737def _custom_colormap_probability(
1738 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None
1739):
1740 """Get a custom colorbar for a probability cube.
1742 Specific variable ranges can be separately set in user-supplied definition
1743 for x- or y-axis limits, or indicate where automated range preferred.
1745 Parameters
1746 ----------
1747 cube: Cube
1748 Cube of variable with probability in name.
1749 axis: "x", "y", optional
1750 Select the levels for just this axis of a line plot. The min and max
1751 can be set by xmin/xmax or ymin/ymax respectively. For variables where
1752 setting a universal range is not desirable (e.g. temperature), users
1753 can set ymin/ymax values to "auto" in the colorbar definitions file.
1754 Where no additional xmin/xmax or ymin/ymax values are provided, the
1755 axis bounds default to use the vmin/vmax values provided.
1757 Returns
1758 -------
1759 cmap:
1760 Matplotlib colormap.
1761 levels:
1762 List of levels to use for plotting. For continuous plots the min and max
1763 should be taken as the range.
1764 norm:
1765 BoundaryNorm information.
1766 """
1767 if axis:
1768 levels = [0, 1]
1769 return None, levels, None
1770 else:
1771 cmap = mcolors.ListedColormap(
1772 [
1773 "#FFFFFF",
1774 "#636363",
1775 "#e1dada",
1776 "#B5CAFF",
1777 "#8FB3FF",
1778 "#7F97FF",
1779 "#ABCF63",
1780 "#E8F59E",
1781 "#FFFA14",
1782 "#FFD121",
1783 "#FFA30A",
1784 ]
1785 )
1786 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]
1787 norm = mcolors.BoundaryNorm(levels, cmap.N)
1788 return cmap, levels, norm
1791def _custom_colourmap_precipitation(cube: iris.cube.Cube, cmap, levels, norm):
1792 """Return a custom colourmap for the current recipe."""
1793 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
1794 if (
1795 any("surface_microphysical" in name for name in varnames)
1796 and "difference" not in cube.long_name
1797 and "mask" not in cube.long_name
1798 ):
1799 # Define the levels and colors
1800 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256]
1801 colors = [
1802 "w",
1803 (0, 0, 0.6),
1804 "b",
1805 "c",
1806 "g",
1807 "y",
1808 (1, 0.5, 0),
1809 "r",
1810 "pink",
1811 "m",
1812 "purple",
1813 "maroon",
1814 "gray",
1815 ]
1816 # Create a custom colormap
1817 cmap = mcolors.ListedColormap(colors)
1818 # Normalize the levels
1819 norm = mcolors.BoundaryNorm(levels, cmap.N)
1820 logging.info("change colormap for surface_microphysical variable colorbar.")
1821 else:
1822 # do nothing and keep existing colorbar attributes
1823 cmap = cmap
1824 levels = levels
1825 norm = norm
1826 return cmap, levels, norm
1829def _custom_colormap_aviation_colour_state(cube: iris.cube.Cube):
1830 """Return custom colourmap for aviation colour state.
1832 If "aviation_colour_state" appears anywhere in the name of a cube
1833 this function will be called.
1835 Parameters
1836 ----------
1837 cube: Cube
1838 Cube of variable for which the colorbar information is desired.
1840 Returns
1841 -------
1842 cmap: Matplotlib colormap.
1843 levels: List
1844 List of levels to use for plotting. For continuous plots the min and max
1845 should be taken as the range.
1846 norm: BoundaryNorm.
1847 """
1848 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5]
1849 colors = [
1850 "#87ceeb",
1851 "#ffffff",
1852 "#8ced69",
1853 "#ffff00",
1854 "#ffd700",
1855 "#ffa500",
1856 "#fe3620",
1857 ]
1858 # Create a custom colormap
1859 cmap = mcolors.ListedColormap(colors)
1860 # Normalise the levels
1861 norm = mcolors.BoundaryNorm(levels, cmap.N)
1862 return cmap, levels, norm
1865def _custom_colourmap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm):
1866 """Return a custom colourmap for the current recipe."""
1867 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
1868 if (
1869 any("visibility_in_air" in name for name in varnames)
1870 and "difference" not in cube.long_name
1871 and "mask" not in cube.long_name
1872 ):
1873 # Define the levels and colors (in km)
1874 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]
1875 norm = mcolors.BoundaryNorm(levels, cmap.N)
1876 colours = [
1877 "#8f00d6",
1878 "#d10000",
1879 "#ff9700",
1880 "#ffff00",
1881 "#00007f",
1882 "#6c9ccd",
1883 "#aae8ff",
1884 "#37a648",
1885 "#8edc64",
1886 "#c5ffc5",
1887 "#dcdcdc",
1888 "#ffffff",
1889 ]
1890 # Create a custom colormap
1891 cmap = mcolors.ListedColormap(colours)
1892 # Normalize the levels
1893 norm = mcolors.BoundaryNorm(levels, cmap.N)
1894 logging.info("change colormap for visibility_in_air variable colorbar.")
1895 else:
1896 # do nothing and keep existing colorbar attributes
1897 cmap = cmap
1898 levels = levels
1899 norm = norm
1900 return cmap, levels, norm
1903def _get_num_models(cube: iris.cube.Cube | iris.cube.CubeList) -> int:
1904 """Return number of models based on cube attributes."""
1905 model_names = list(
1906 filter(
1907 lambda x: x is not None,
1908 {cb.attributes.get("model_name", None) for cb in iter_maybe(cube)},
1909 )
1910 )
1911 if not model_names:
1912 logging.debug("Missing model names. Will assume single model.")
1913 return 1
1914 else:
1915 return len(model_names)
1918def _validate_cube_shape(
1919 cube: iris.cube.Cube | iris.cube.CubeList, num_models: int
1920) -> None:
1921 """Check all cubes have a model name."""
1922 if isinstance(cube, iris.cube.CubeList) and len(cube) != num_models: 1922 ↛ 1923line 1922 didn't jump to line 1923 because the condition on line 1922 was never true
1923 raise ValueError(
1924 f"The number of model names ({num_models}) should equal the number "
1925 f"of cubes ({len(cube)})."
1926 )
1929def _validate_cubes_coords(
1930 cubes: iris.cube.CubeList, coords: list[iris.coords.Coord]
1931) -> None:
1932 """Check same number of cubes as sequence coordinate for zip functions."""
1933 if len(cubes) != len(coords): 1933 ↛ 1934line 1933 didn't jump to line 1934 because the condition on line 1933 was never true
1934 raise ValueError(
1935 f"The number of CubeList entries ({len(cubes)}) should equal the number "
1936 f"of sequence coordinates ({len(coords)})."
1937 f"Check that number of time entries in input data are consistent if "
1938 f"performing time-averaging steps prior to plotting outputs."
1939 )
1942####################
1943# Public functions #
1944####################
1947def spatial_contour_plot(
1948 cube: iris.cube.Cube,
1949 filename: str = None,
1950 sequence_coordinate: str = "time",
1951 stamp_coordinate: str = "realization",
1952 **kwargs,
1953) -> iris.cube.Cube:
1954 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
1956 A 2D spatial field can be plotted, but if the sequence_coordinate is present
1957 then a sequence of plots will be produced. Similarly if the stamp_coordinate
1958 is present then postage stamp plots will be produced.
1960 Parameters
1961 ----------
1962 cube: Cube
1963 Iris cube of the data to plot. It should have two spatial dimensions,
1964 such as lat and lon, and may also have a another two dimension to be
1965 plotted sequentially and/or as postage stamp plots.
1966 filename: str, optional
1967 Name of the plot to write, used as a prefix for plot sequences. Defaults
1968 to the recipe name.
1969 sequence_coordinate: str, optional
1970 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
1971 This coordinate must exist in the cube.
1972 stamp_coordinate: str, optional
1973 Coordinate about which to plot postage stamp plots. Defaults to
1974 ``"realization"``.
1976 Returns
1977 -------
1978 Cube
1979 The original cube (so further operations can be applied).
1981 Raises
1982 ------
1983 ValueError
1984 If the cube doesn't have the right dimensions.
1985 TypeError
1986 If the cube isn't a single cube.
1987 """
1988 _spatial_plot("contourf", cube, filename, sequence_coordinate, stamp_coordinate)
1989 return cube
1992def spatial_pcolormesh_plot(
1993 cube: iris.cube.Cube,
1994 filename: str = None,
1995 sequence_coordinate: str = "time",
1996 stamp_coordinate: str = "realization",
1997 **kwargs,
1998) -> iris.cube.Cube:
1999 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
2001 A 2D spatial field can be plotted, but if the sequence_coordinate is present
2002 then a sequence of plots will be produced. Similarly if the stamp_coordinate
2003 is present then postage stamp plots will be produced.
2005 This function is significantly faster than ``spatial_contour_plot``,
2006 especially at high resolutions, and should be preferred unless contiguous
2007 contour areas are important.
2009 Parameters
2010 ----------
2011 cube: Cube
2012 Iris cube of the data to plot. It should have two spatial dimensions,
2013 such as lat and lon, and may also have a another two dimension to be
2014 plotted sequentially and/or as postage stamp plots.
2015 filename: str, optional
2016 Name of the plot to write, used as a prefix for plot sequences. Defaults
2017 to the recipe name.
2018 sequence_coordinate: str, optional
2019 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2020 This coordinate must exist in the cube.
2021 stamp_coordinate: str, optional
2022 Coordinate about which to plot postage stamp plots. Defaults to
2023 ``"realization"``.
2025 Returns
2026 -------
2027 Cube
2028 The original cube (so further operations can be applied).
2030 Raises
2031 ------
2032 ValueError
2033 If the cube doesn't have the right dimensions.
2034 TypeError
2035 If the cube isn't a single cube.
2036 """
2037 _spatial_plot("pcolormesh", cube, filename, sequence_coordinate, stamp_coordinate)
2038 return cube
2041# TODO: Expand function to handle ensemble data.
2042# line_coordinate: str, optional
2043# Coordinate about which to plot multiple lines. Defaults to
2044# ``"realization"``.
2045def plot_line_series(
2046 cube: iris.cube.Cube | iris.cube.CubeList,
2047 filename: str = None,
2048 series_coordinate: str = "time",
2049 # line_coordinate: str = "realization",
2050 **kwargs,
2051) -> iris.cube.Cube | iris.cube.CubeList:
2052 """Plot a line plot for the specified coordinate.
2054 The Cube or CubeList must be 1D.
2056 Parameters
2057 ----------
2058 iris.cube | iris.cube.CubeList
2059 Cube or CubeList of the data to plot. The individual cubes should have a single dimension.
2060 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
2061 We do not support different data such as temperature and humidity in the same CubeList for plotting.
2062 filename: str, optional
2063 Name of the plot to write, used as a prefix for plot sequences. Defaults
2064 to the recipe name.
2065 series_coordinate: str, optional
2066 Coordinate about which to make a series. Defaults to ``"time"``. This
2067 coordinate must exist in the cube.
2069 Returns
2070 -------
2071 iris.cube.Cube | iris.cube.CubeList
2072 The original Cube or CubeList (so further operations can be applied).
2073 plotted data.
2075 Raises
2076 ------
2077 ValueError
2078 If the cubes don't have the right dimensions.
2079 TypeError
2080 If the cube isn't a Cube or CubeList.
2081 """
2082 # Ensure we have a name for the plot file.
2083 title = get_recipe_metadata().get("title", "Untitled")
2085 if filename is None:
2086 filename = slugify(title)
2088 # Add file extension.
2089 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
2091 num_models = _get_num_models(cube)
2093 _validate_cube_shape(cube, num_models)
2095 # Iterate over all cubes and extract coordinate to plot.
2096 cubes = iter_maybe(cube)
2097 coords = []
2098 for cube in cubes:
2099 try:
2100 coords.append(cube.coord(series_coordinate))
2101 except iris.exceptions.CoordinateNotFoundError as err:
2102 raise ValueError(
2103 f"Cube must have a {series_coordinate} coordinate."
2104 ) from err
2105 if cube.ndim > 2 or not cube.coords("realization"):
2106 raise ValueError("Cube must be 1D or 2D with a realization coordinate.")
2108 # Do the actual plotting.
2109 _plot_and_save_line_series(cubes, coords, "realization", plot_filename, title)
2111 # Add list of plots to plot metadata.
2112 plot_index = _append_to_plot_index([plot_filename])
2114 # Make a page to display the plots.
2115 _make_plot_html_page(plot_index)
2117 return cube
2120def plot_vertical_line_series(
2121 cubes: iris.cube.Cube | iris.cube.CubeList,
2122 filename: str = None,
2123 series_coordinate: str = "model_level_number",
2124 sequence_coordinate: str = "time",
2125 # line_coordinate: str = "realization",
2126 **kwargs,
2127) -> iris.cube.Cube | iris.cube.CubeList:
2128 """Plot a line plot against a type of vertical coordinate.
2130 The Cube or CubeList must be 1D.
2132 A 1D line plot with y-axis as pressure coordinate can be plotted, but if the sequence_coordinate is present
2133 then a sequence of plots will be produced.
2135 Parameters
2136 ----------
2137 iris.cube | iris.cube.CubeList
2138 Cube or CubeList of the data to plot. The individual cubes should have a single dimension.
2139 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
2140 We do not support different data such as temperature and humidity in the same CubeList for plotting.
2141 filename: str, optional
2142 Name of the plot to write, used as a prefix for plot sequences. Defaults
2143 to the recipe name.
2144 series_coordinate: str, optional
2145 Coordinate to plot on the y-axis. Can be ``pressure`` or
2146 ``model_level_number`` for UM, or ``full_levels`` or ``half_levels``
2147 for LFRic. Defaults to ``model_level_number``.
2148 This coordinate must exist in the cube.
2149 sequence_coordinate: str, optional
2150 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2151 This coordinate must exist in the cube.
2153 Returns
2154 -------
2155 iris.cube.Cube | iris.cube.CubeList
2156 The original Cube or CubeList (so further operations can be applied).
2157 Plotted data.
2159 Raises
2160 ------
2161 ValueError
2162 If the cubes doesn't have the right dimensions.
2163 TypeError
2164 If the cube isn't a Cube or CubeList.
2165 """
2166 # Ensure we have a name for the plot file.
2167 recipe_title = get_recipe_metadata().get("title", "Untitled")
2169 if filename is None:
2170 filename = slugify(recipe_title)
2172 cubes = iter_maybe(cubes)
2173 # Initialise empty list to hold all data from all cubes in a CubeList
2174 all_data = []
2176 # Store min/max ranges for x range.
2177 x_levels = []
2179 num_models = _get_num_models(cubes)
2181 _validate_cube_shape(cubes, num_models)
2183 # Iterate over all cubes in cube or CubeList and plot.
2184 coords = []
2185 for cube in cubes:
2186 # Test if series coordinate i.e. pressure level exist for any cube with cube.ndim >=1.
2187 try:
2188 coords.append(cube.coord(series_coordinate))
2189 except iris.exceptions.CoordinateNotFoundError as err:
2190 raise ValueError(
2191 f"Cube must have a {series_coordinate} coordinate."
2192 ) from err
2194 try:
2195 if cube.ndim > 1 or not cube.coords("realization"): 2195 ↛ 2203line 2195 didn't jump to line 2203 because the condition on line 2195 was always true
2196 cube.coord(sequence_coordinate)
2197 except iris.exceptions.CoordinateNotFoundError as err:
2198 raise ValueError(
2199 f"Cube must have a {sequence_coordinate} coordinate or be 1D, or 2D with a realization coordinate."
2200 ) from err
2202 # Get minimum and maximum from levels information.
2203 _, levels, _ = _colorbar_map_levels(cube, axis="x")
2204 if levels is not None: 2204 ↛ 2208line 2204 didn't jump to line 2208 because the condition on line 2204 was always true
2205 x_levels.append(min(levels))
2206 x_levels.append(max(levels))
2207 else:
2208 all_data.append(cube.data)
2210 if len(x_levels) == 0: 2210 ↛ 2212line 2210 didn't jump to line 2212 because the condition on line 2210 was never true
2211 # Combine all data into a single NumPy array
2212 combined_data = np.concatenate(all_data)
2214 # Set the lower and upper limit for the x-axis to ensure all plots have
2215 # same range. This needs to read the whole cube over the range of the
2216 # sequence and if applicable postage stamp coordinate.
2217 vmin = np.floor(combined_data.min())
2218 vmax = np.ceil(combined_data.max())
2219 else:
2220 vmin = min(x_levels)
2221 vmax = max(x_levels)
2223 # Matching the slices (matching by seq coord point; it may happen that
2224 # evaluated models do not cover the same seq coord range, hence matching
2225 # necessary)
2226 def filter_cube_iterables(cube_iterables) -> bool:
2227 return len(cube_iterables) == len(coords)
2229 cube_iterables = filter(
2230 filter_cube_iterables,
2231 (
2232 iris.cube.CubeList(
2233 s
2234 for s in itertools.chain.from_iterable(
2235 cb.slices_over(sequence_coordinate) for cb in cubes
2236 )
2237 if s.coord(sequence_coordinate).points[0] == point
2238 )
2239 for point in sorted(
2240 set(
2241 itertools.chain.from_iterable(
2242 cb.coord(sequence_coordinate).points for cb in cubes
2243 )
2244 )
2245 )
2246 ),
2247 )
2249 # Create a plot for each value of the sequence coordinate.
2250 # Allowing for multiple cubes in a CubeList to be plotted in the same plot for
2251 # similar sequence values. Passing a CubeList into the internal plotting function
2252 # for similar values of the sequence coordinate. cube_slice can be an iris.cube.Cube
2253 # or an iris.cube.CubeList.
2254 plot_index = []
2255 nplot = np.size(cubes[0].coord(sequence_coordinate).points)
2256 for cubes_slice in cube_iterables:
2257 # Use sequence value so multiple sequences can merge.
2258 seq_coord = cubes_slice[0].coord(sequence_coordinate)
2259 sequence_value = seq_coord.points[0]
2260 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
2261 # Format the coordinate value in a unit appropriate way.
2262 title = f"{recipe_title}\n [{seq_coord.units.title(sequence_value)}]"
2263 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs
2264 if nplot == 1 and seq_coord.has_bounds: 2264 ↛ 2265line 2264 didn't jump to line 2265 because the condition on line 2264 was never true
2265 if np.size(seq_coord.bounds) > 1:
2266 title = f"{recipe_title}\n [{seq_coord.units.title(seq_coord.bounds[0][0])} to {seq_coord.units.title(seq_coord.bounds[0][1])}]"
2267 # Do the actual plotting.
2268 _plot_and_save_vertical_line_series(
2269 cubes_slice,
2270 coords,
2271 "realization",
2272 plot_filename,
2273 series_coordinate,
2274 title=title,
2275 vmin=vmin,
2276 vmax=vmax,
2277 )
2278 plot_index.append(plot_filename)
2280 # Add list of plots to plot metadata.
2281 complete_plot_index = _append_to_plot_index(plot_index)
2283 # Make a page to display the plots.
2284 _make_plot_html_page(complete_plot_index)
2286 return cubes
2289def qq_plot(
2290 cubes: iris.cube.CubeList,
2291 coordinates: list[str],
2292 percentiles: list[float],
2293 model_names: list[str],
2294 filename: str = None,
2295 one_to_one: bool = True,
2296 **kwargs,
2297) -> iris.cube.CubeList:
2298 """Plot a Quantile-Quantile plot between two models for common time points.
2300 The cubes will be normalised by collapsing each cube to its percentiles. Cubes are
2301 collapsed within the operator over all specified coordinates such as
2302 grid_latitude, grid_longitude, vertical levels, but also realisation representing
2303 ensemble members to ensure a 1D cube (array).
2305 Parameters
2306 ----------
2307 cubes: iris.cube.CubeList
2308 Two cubes of the same variable with different models.
2309 coordinate: list[str]
2310 The list of coordinates to collapse over. This list should be
2311 every coordinate within the cube to result in a 1D cube around
2312 the percentile coordinate.
2313 percent: list[float]
2314 A list of percentiles to appear in the plot.
2315 model_names: list[str]
2316 A list of model names to appear on the axis of the plot.
2317 filename: str, optional
2318 Filename of the plot to write.
2319 one_to_one: bool, optional
2320 If True a 1:1 line is plotted; if False it is not. Default is True.
2322 Raises
2323 ------
2324 ValueError
2325 When the cubes are not compatible.
2327 Notes
2328 -----
2329 The quantile-quantile plot is a variant on the scatter plot representing
2330 two datasets by their quantiles (percentiles) for common time points.
2331 This plot does not use a theoretical distribution to compare against, but
2332 compares percentiles of two datasets. This plot does
2333 not use all raw data points, but plots the selected percentiles (quantiles) of
2334 each variable instead for the two datasets, thereby normalising the data for a
2335 direct comparison between the selected percentiles of the two dataset distributions.
2337 Quantile-quantile plots are valuable for comparing against
2338 observations and other models. Identical percentiles between the variables
2339 will lie on the one-to-one line implying the values correspond well to each
2340 other. Where there is a deviation from the one-to-one line a range of
2341 possibilities exist depending on how and where the data is shifted (e.g.,
2342 Wilks 2011 [Wilks2011]_).
2344 For distributions above the one-to-one line the distribution is left-skewed;
2345 below is right-skewed. A distinct break implies a bimodal distribution, and
2346 closer values/values further apart at the tails imply poor representation of
2347 the extremes.
2349 References
2350 ----------
2351 .. [Wilks2011] Wilks, D.S., (2011) "Statistical Methods in the Atmospheric
2352 Sciences" Third Edition, vol. 100, Academic Press, Oxford, UK, 676 pp.
2353 """
2354 # Check cubes using same functionality as the difference operator.
2355 if len(cubes) != 2:
2356 raise ValueError("cubes should contain exactly 2 cubes.")
2357 base: Cube = cubes.extract_cube(iris.AttributeConstraint(cset_comparison_base=1))
2358 other: Cube = cubes.extract_cube(
2359 iris.Constraint(
2360 cube_func=lambda cube: "cset_comparison_base" not in cube.attributes
2361 )
2362 )
2364 # Get spatial coord names.
2365 base_lat_name, base_lon_name = get_cube_yxcoordname(base)
2366 other_lat_name, other_lon_name = get_cube_yxcoordname(other)
2368 # Ensure cubes to compare are on common differencing grid.
2369 # This is triggered if either
2370 # i) latitude and longitude shapes are not the same. Note grid points
2371 # are not compared directly as these can differ through rounding
2372 # errors.
2373 # ii) or variables are known to often sit on different grid staggering
2374 # in different models (e.g. cell center vs cell edge), as is the case
2375 # for UM and LFRic comparisons.
2376 # In future greater choice of regridding method might be applied depending
2377 # on variable type. Linear regridding can in general be appropriate for smooth
2378 # variables. Care should be taken with interpretation of differences
2379 # given this dependency on regridding.
2380 if (
2381 base.coord(base_lat_name).shape != other.coord(other_lat_name).shape
2382 or base.coord(base_lon_name).shape != other.coord(other_lon_name).shape
2383 ) or (
2384 base.long_name
2385 in [
2386 "eastward_wind_at_10m",
2387 "northward_wind_at_10m",
2388 "northward_wind_at_cell_centres",
2389 "eastward_wind_at_cell_centres",
2390 "zonal_wind_at_pressure_levels",
2391 "meridional_wind_at_pressure_levels",
2392 "potential_vorticity_at_pressure_levels",
2393 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",
2394 ]
2395 ):
2396 logging.debug(
2397 "Linear regridding base cube to other grid to compute differences"
2398 )
2399 base = regrid_onto_cube(base, other, method="Linear")
2401 # Extract just common time points.
2402 base, other = _extract_common_time_points(base, other)
2404 # Equalise attributes so we can merge.
2405 fully_equalise_attributes([base, other])
2406 logging.debug("Base: %s\nOther: %s", base, other)
2408 # Collapse cubes.
2409 base = collapse(
2410 base,
2411 coordinate=coordinates,
2412 method="PERCENTILE",
2413 additional_percent=percentiles,
2414 )
2415 other = collapse(
2416 other,
2417 coordinate=coordinates,
2418 method="PERCENTILE",
2419 additional_percent=percentiles,
2420 )
2422 # Ensure we have a name for the plot file.
2423 title = get_recipe_metadata().get("title", "Untitled")
2425 if filename is None:
2426 filename = slugify(title)
2428 # Add file extension.
2429 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
2431 # Do the actual plotting on a scatter plot
2432 _plot_and_save_scatter_plot(
2433 base, other, plot_filename, title, one_to_one, model_names
2434 )
2436 # Add list of plots to plot metadata.
2437 plot_index = _append_to_plot_index([plot_filename])
2439 # Make a page to display the plots.
2440 _make_plot_html_page(plot_index)
2442 return iris.cube.CubeList([base, other])
2445def scatter_plot(
2446 cube_x: iris.cube.Cube | iris.cube.CubeList,
2447 cube_y: iris.cube.Cube | iris.cube.CubeList,
2448 filename: str = None,
2449 one_to_one: bool = True,
2450 **kwargs,
2451) -> iris.cube.CubeList:
2452 """Plot a scatter plot between two variables.
2454 Both cubes must be 1D.
2456 Parameters
2457 ----------
2458 cube_x: Cube | CubeList
2459 1 dimensional Cube of the data to plot on y-axis.
2460 cube_y: Cube | CubeList
2461 1 dimensional Cube of the data to plot on x-axis.
2462 filename: str, optional
2463 Filename of the plot to write.
2464 one_to_one: bool, optional
2465 If True a 1:1 line is plotted; if False it is not. Default is True.
2467 Returns
2468 -------
2469 cubes: CubeList
2470 CubeList of the original x and y cubes for further processing.
2472 Raises
2473 ------
2474 ValueError
2475 If the cube doesn't have the right dimensions and cubes not the same
2476 size.
2477 TypeError
2478 If the cube isn't a single cube.
2480 Notes
2481 -----
2482 Scatter plots are used for determining if there is a relationship between
2483 two variables. Positive relations have a slope going from bottom left to top
2484 right; Negative relations have a slope going from top left to bottom right.
2485 """
2486 # Iterate over all cubes in cube or CubeList and plot.
2487 for cube_iter in iter_maybe(cube_x):
2488 # Check cubes are correct shape.
2489 cube_iter = _check_single_cube(cube_iter)
2490 if cube_iter.ndim > 1:
2491 raise ValueError("cube_x must be 1D.")
2493 # Iterate over all cubes in cube or CubeList and plot.
2494 for cube_iter in iter_maybe(cube_y):
2495 # Check cubes are correct shape.
2496 cube_iter = _check_single_cube(cube_iter)
2497 if cube_iter.ndim > 1:
2498 raise ValueError("cube_y must be 1D.")
2500 # Ensure we have a name for the plot file.
2501 title = get_recipe_metadata().get("title", "Untitled")
2503 if filename is None:
2504 filename = slugify(title)
2506 # Add file extension.
2507 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
2509 # Do the actual plotting.
2510 _plot_and_save_scatter_plot(cube_x, cube_y, plot_filename, title, one_to_one)
2512 # Add list of plots to plot metadata.
2513 plot_index = _append_to_plot_index([plot_filename])
2515 # Make a page to display the plots.
2516 _make_plot_html_page(plot_index)
2518 return iris.cube.CubeList([cube_x, cube_y])
2521def vector_plot(
2522 cube_u: iris.cube.Cube,
2523 cube_v: iris.cube.Cube,
2524 filename: str = None,
2525 sequence_coordinate: str = "time",
2526 **kwargs,
2527) -> iris.cube.CubeList:
2528 """Plot a vector plot based on the input u and v components."""
2529 recipe_title = get_recipe_metadata().get("title", "Untitled")
2531 # Ensure we have a name for the plot file.
2532 if filename is None: 2532 ↛ 2533line 2532 didn't jump to line 2533 because the condition on line 2532 was never true
2533 filename = slugify(recipe_title)
2535 # Cubes must have a matching sequence coordinate.
2536 try:
2537 # Check that the u and v cubes have the same sequence coordinate.
2538 if cube_u.coord(sequence_coordinate) != cube_v.coord(sequence_coordinate): 2538 ↛ 2539line 2538 didn't jump to line 2539 because the condition on line 2538 was never true
2539 raise ValueError("Coordinates do not match.")
2540 except (iris.exceptions.CoordinateNotFoundError, ValueError) as err:
2541 raise ValueError(
2542 f"Cubes should have matching {sequence_coordinate} coordinate:\n{cube_u}\n{cube_v}"
2543 ) from err
2545 # Create a plot for each value of the sequence coordinate.
2546 plot_index = []
2547 for cube_u_slice, cube_v_slice in zip(
2548 cube_u.slices_over(sequence_coordinate),
2549 cube_v.slices_over(sequence_coordinate),
2550 strict=True,
2551 ):
2552 # Use sequence value so multiple sequences can merge.
2553 sequence_value = cube_u_slice.coord(sequence_coordinate).points[0]
2554 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
2555 coord = cube_u_slice.coord(sequence_coordinate)
2556 # Format the coordinate value in a unit appropriate way.
2557 title = f"{recipe_title}\n{coord.units.title(coord.points[0])}"
2558 # Do the actual plotting.
2559 _plot_and_save_vector_plot(
2560 cube_u_slice,
2561 cube_v_slice,
2562 filename=plot_filename,
2563 title=title,
2564 method="contourf",
2565 )
2566 plot_index.append(plot_filename)
2568 # Add list of plots to plot metadata.
2569 complete_plot_index = _append_to_plot_index(plot_index)
2571 # Make a page to display the plots.
2572 _make_plot_html_page(complete_plot_index)
2574 return iris.cube.CubeList([cube_u, cube_v])
2577def plot_histogram_series(
2578 cubes: iris.cube.Cube | iris.cube.CubeList,
2579 filename: str = None,
2580 sequence_coordinate: str = "time",
2581 stamp_coordinate: str = "realization",
2582 single_plot: bool = False,
2583 **kwargs,
2584) -> iris.cube.Cube | iris.cube.CubeList:
2585 """Plot a histogram plot for each vertical level provided.
2587 A histogram plot can be plotted, but if the sequence_coordinate (i.e. time)
2588 is present then a sequence of plots will be produced using the time slider
2589 functionality to scroll through histograms against time. If a
2590 stamp_coordinate is present then postage stamp plots will be produced. If
2591 stamp_coordinate and single_plot is True, all postage stamp plots will be
2592 plotted in a single plot instead of separate postage stamp plots.
2594 Parameters
2595 ----------
2596 cubes: Cube | iris.cube.CubeList
2597 Iris cube or CubeList of the data to plot. It should have a single dimension other
2598 than the stamp coordinate.
2599 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
2600 We do not support different data such as temperature and humidity in the same CubeList for plotting.
2601 filename: str, optional
2602 Name of the plot to write, used as a prefix for plot sequences. Defaults
2603 to the recipe name.
2604 sequence_coordinate: str, optional
2605 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2606 This coordinate must exist in the cube and will be used for the time
2607 slider.
2608 stamp_coordinate: str, optional
2609 Coordinate about which to plot postage stamp plots. Defaults to
2610 ``"realization"``.
2611 single_plot: bool, optional
2612 If True, all postage stamp plots will be plotted in a single plot. If
2613 False, each postage stamp plot will be plotted separately. Is only valid
2614 if stamp_coordinate exists and has more than a single point.
2616 Returns
2617 -------
2618 iris.cube.Cube | iris.cube.CubeList
2619 The original Cube or CubeList (so further operations can be applied).
2620 Plotted data.
2622 Raises
2623 ------
2624 ValueError
2625 If the cube doesn't have the right dimensions.
2626 TypeError
2627 If the cube isn't a Cube or CubeList.
2628 """
2629 recipe_title = get_recipe_metadata().get("title", "Untitled")
2631 cubes = iter_maybe(cubes)
2633 # Ensure we have a name for the plot file.
2634 if filename is None:
2635 filename = slugify(recipe_title)
2637 # Internal plotting function.
2638 plotting_func = _plot_and_save_histogram_series
2640 num_models = _get_num_models(cubes)
2642 _validate_cube_shape(cubes, num_models)
2644 # If several histograms are plotted with time as sequence_coordinate for the
2645 # time slider option.
2646 for cube in cubes:
2647 try:
2648 cube.coord(sequence_coordinate)
2649 except iris.exceptions.CoordinateNotFoundError as err:
2650 raise ValueError(
2651 f"Cube must have a {sequence_coordinate} coordinate."
2652 ) from err
2654 # Get minimum and maximum from levels information.
2655 levels = None
2656 for cube in cubes: 2656 ↛ 2672line 2656 didn't jump to line 2672 because the loop on line 2656 didn't complete
2657 # First check if user-specified "auto" range variable.
2658 # This maintains the value of levels as None, so proceed.
2659 _, levels, _ = _colorbar_map_levels(cube, axis="y")
2660 if levels is None:
2661 break
2662 # If levels is changed, recheck to use the vmin,vmax or
2663 # levels-based ranges for histogram plots.
2664 _, levels, _ = _colorbar_map_levels(cube)
2665 logging.debug("levels: %s", levels)
2666 if levels is not None: 2666 ↛ 2656line 2666 didn't jump to line 2656 because the condition on line 2666 was always true
2667 vmin = min(levels)
2668 vmax = max(levels)
2669 logging.debug("Updated vmin, vmax: %s, %s", vmin, vmax)
2670 break
2672 if levels is None:
2673 vmin = min(cb.data.min() for cb in cubes)
2674 vmax = max(cb.data.max() for cb in cubes)
2676 # Make postage stamp plots if stamp_coordinate exists and has more than a
2677 # single point. If single_plot is True:
2678 # -- all postage stamp plots will be plotted in a single plot instead of
2679 # separate postage stamp plots.
2680 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are
2681 # produced per single model only
2682 if num_models == 1: 2682 ↛ 2695line 2682 didn't jump to line 2695 because the condition on line 2682 was always true
2683 if ( 2683 ↛ 2687line 2683 didn't jump to line 2687 because the condition on line 2683 was never true
2684 stamp_coordinate in [c.name() for c in cubes[0].coords()]
2685 and cubes[0].coord(stamp_coordinate).shape[0] > 1
2686 ):
2687 if single_plot:
2688 plotting_func = (
2689 _plot_and_save_postage_stamps_in_single_plot_histogram_series
2690 )
2691 else:
2692 plotting_func = _plot_and_save_postage_stamp_histogram_series
2693 cube_iterables = cubes[0].slices_over(sequence_coordinate)
2694 else:
2695 all_points = sorted(
2696 set(
2697 itertools.chain.from_iterable(
2698 cb.coord(sequence_coordinate).points for cb in cubes
2699 )
2700 )
2701 )
2702 all_slices = list(
2703 itertools.chain.from_iterable(
2704 cb.slices_over(sequence_coordinate) for cb in cubes
2705 )
2706 )
2707 # Matched slices (matched by seq coord point; it may happen that
2708 # evaluated models do not cover the same seq coord range, hence matching
2709 # necessary)
2710 cube_iterables = [
2711 iris.cube.CubeList(
2712 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point
2713 )
2714 for point in all_points
2715 ]
2717 plot_index = []
2718 nplot = np.size(cube.coord(sequence_coordinate).points)
2719 # Create a plot for each value of the sequence coordinate. Allowing for
2720 # multiple cubes in a CubeList to be plotted in the same plot for similar
2721 # sequence values. Passing a CubeList into the internal plotting function
2722 # for similar values of the sequence coordinate. cube_slice can be an
2723 # iris.cube.Cube or an iris.cube.CubeList.
2724 for cube_slice in cube_iterables:
2725 single_cube = cube_slice
2726 if isinstance(cube_slice, iris.cube.CubeList): 2726 ↛ 2727line 2726 didn't jump to line 2727 because the condition on line 2726 was never true
2727 single_cube = cube_slice[0]
2729 # Use sequence value so multiple sequences can merge.
2730 sequence_value = single_cube.coord(sequence_coordinate).points[0]
2731 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
2732 coord = single_cube.coord(sequence_coordinate)
2733 # Format the coordinate value in a unit appropriate way.
2734 title = f"{recipe_title}\n [{coord.units.title(coord.points[0])}]"
2735 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs
2736 if nplot == 1 and coord.has_bounds: 2736 ↛ 2737line 2736 didn't jump to line 2737 because the condition on line 2736 was never true
2737 if np.size(coord.bounds) > 1:
2738 title = f"{recipe_title}\n [{coord.units.title(coord.bounds[0][0])} to {coord.units.title(coord.bounds[0][1])}]"
2739 # Do the actual plotting.
2740 plotting_func(
2741 cube_slice,
2742 filename=plot_filename,
2743 stamp_coordinate=stamp_coordinate,
2744 title=title,
2745 vmin=vmin,
2746 vmax=vmax,
2747 )
2748 plot_index.append(plot_filename)
2750 # Add list of plots to plot metadata.
2751 complete_plot_index = _append_to_plot_index(plot_index)
2753 # Make a page to display the plots.
2754 _make_plot_html_page(complete_plot_index)
2756 return cubes
2759def plot_power_spectrum_series(
2760 cubes: iris.cube.Cube | iris.cube.CubeList,
2761 filename: str = None,
2762 sequence_coordinate: str = "time",
2763 stamp_coordinate: str = "realization",
2764 single_plot: bool = False,
2765 **kwargs,
2766) -> iris.cube.Cube | iris.cube.CubeList:
2767 """Plot a power spectrum plot for each vertical level provided.
2769 A power spectrum plot can be plotted, but if the sequence_coordinate (i.e. time)
2770 is present then a sequence of plots will be produced using the time slider
2771 functionality to scroll through power spectrum against time. If a
2772 stamp_coordinate is present then postage stamp plots will be produced. If
2773 stamp_coordinate and single_plot is True, all postage stamp plots will be
2774 plotted in a single plot instead of separate postage stamp plots.
2776 Parameters
2777 ----------
2778 cubes: Cube | iris.cube.CubeList
2779 Iris cube or CubeList of the data to plot. It should have a single dimension other
2780 than the stamp coordinate.
2781 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
2782 We do not support different data such as temperature and humidity in the same CubeList for plotting.
2783 filename: str, optional
2784 Name of the plot to write, used as a prefix for plot sequences. Defaults
2785 to the recipe name.
2786 sequence_coordinate: str, optional
2787 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2788 This coordinate must exist in the cube and will be used for the time
2789 slider.
2790 stamp_coordinate: str, optional
2791 Coordinate about which to plot postage stamp plots. Defaults to
2792 ``"realization"``.
2793 single_plot: bool, optional
2794 If True, all postage stamp plots will be plotted in a single plot. If
2795 False, each postage stamp plot will be plotted separately. Is only valid
2796 if stamp_coordinate exists and has more than a single point.
2798 Returns
2799 -------
2800 iris.cube.Cube | iris.cube.CubeList
2801 The original Cube or CubeList (so further operations can be applied).
2802 Plotted data.
2804 Raises
2805 ------
2806 ValueError
2807 If the cube doesn't have the right dimensions.
2808 TypeError
2809 If the cube isn't a Cube or CubeList.
2810 """
2811 recipe_title = get_recipe_metadata().get("title", "Untitled")
2813 cubes = iter_maybe(cubes)
2814 # Ensure we have a name for the plot file.
2815 if filename is None:
2816 filename = slugify(recipe_title)
2818 # Internal plotting function.
2819 plotting_func = _plot_and_save_power_spectrum_series
2821 num_models = _get_num_models(cubes)
2823 _validate_cube_shape(cubes, num_models)
2825 # If several power spectra are plotted with time as sequence_coordinate for the
2826 # time slider option.
2827 for cube in cubes:
2828 try:
2829 cube.coord(sequence_coordinate)
2830 except iris.exceptions.CoordinateNotFoundError as err:
2831 raise ValueError(
2832 f"Cube must have a {sequence_coordinate} coordinate."
2833 ) from err
2835 # Make postage stamp plots if stamp_coordinate exists and has more than a
2836 # single point. If single_plot is True:
2837 # -- all postage stamp plots will be plotted in a single plot instead of
2838 # separate postage stamp plots.
2839 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are
2840 # produced per single model only
2841 if num_models == 1: 2841 ↛ 2854line 2841 didn't jump to line 2854 because the condition on line 2841 was always true
2842 if ( 2842 ↛ 2846line 2842 didn't jump to line 2846 because the condition on line 2842 was never true
2843 stamp_coordinate in [c.name() for c in cubes[0].coords()]
2844 and cubes[0].coord(stamp_coordinate).shape[0] > 1
2845 ):
2846 if single_plot:
2847 plotting_func = (
2848 _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series
2849 )
2850 else:
2851 plotting_func = _plot_and_save_postage_stamp_power_spectrum_series
2852 cube_iterables = cubes[0].slices_over(sequence_coordinate)
2853 else:
2854 all_points = sorted(
2855 set(
2856 itertools.chain.from_iterable(
2857 cb.coord(sequence_coordinate).points for cb in cubes
2858 )
2859 )
2860 )
2861 all_slices = list(
2862 itertools.chain.from_iterable(
2863 cb.slices_over(sequence_coordinate) for cb in cubes
2864 )
2865 )
2866 # Matched slices (matched by seq coord point; it may happen that
2867 # evaluated models do not cover the same seq coord range, hence matching
2868 # necessary)
2869 cube_iterables = [
2870 iris.cube.CubeList(
2871 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point
2872 )
2873 for point in all_points
2874 ]
2876 plot_index = []
2877 nplot = np.size(cube.coord(sequence_coordinate).points)
2878 # Create a plot for each value of the sequence coordinate. Allowing for
2879 # multiple cubes in a CubeList to be plotted in the same plot for similar
2880 # sequence values. Passing a CubeList into the internal plotting function
2881 # for similar values of the sequence coordinate. cube_slice can be an
2882 # iris.cube.Cube or an iris.cube.CubeList.
2883 for cube_slice in cube_iterables:
2884 single_cube = cube_slice
2885 if isinstance(cube_slice, iris.cube.CubeList): 2885 ↛ 2886line 2885 didn't jump to line 2886 because the condition on line 2885 was never true
2886 single_cube = cube_slice[0]
2888 # Use sequence value so multiple sequences can merge.
2889 sequence_value = single_cube.coord(sequence_coordinate).points[0]
2890 plot_filename = f"{filename.rsplit('.', 1)[0]}_{sequence_value}.png"
2891 coord = single_cube.coord(sequence_coordinate)
2892 # Format the coordinate value in a unit appropriate way.
2893 title = f"{recipe_title}\n [{coord.units.title(coord.points[0])}]"
2894 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs
2895 if nplot == 1 and coord.has_bounds: 2895 ↛ 2899line 2895 didn't jump to line 2899 because the condition on line 2895 was always true
2896 if np.size(coord.bounds) > 1:
2897 title = f"{recipe_title}\n [{coord.units.title(coord.bounds[0][0])} to {coord.units.title(coord.bounds[0][1])}]"
2898 # Do the actual plotting.
2899 plotting_func(
2900 cube_slice,
2901 filename=plot_filename,
2902 stamp_coordinate=stamp_coordinate,
2903 title=title,
2904 )
2905 plot_index.append(plot_filename)
2907 # Add list of plots to plot metadata.
2908 complete_plot_index = _append_to_plot_index(plot_index)
2910 # Make a page to display the plots.
2911 _make_plot_html_page(complete_plot_index)
2913 return cubes
2916def _DCT_ps(y_3d):
2917 """Calculate power spectra for regional domains.
2919 Parameters
2920 ----------
2921 y_3d: 3D array
2922 3 dimensional array to calculate spectrum for.
2923 (2D field data with 3rd dimension of time)
2925 Returns: ps_array
2926 Array of power spectra values calculated for input field (for each time)
2928 Method for regional domains:
2929 Calculate power spectra over limited area domain using Discrete Cosine Transform (DCT)
2930 as described in Denis et al 2002 [Denis_etal_2002]_.
2932 References
2933 ----------
2934 .. [Denis_etal_2002] Bertrand Denis, Jean Côté and René Laprise (2002)
2935 "Spectral Decomposition of Two-Dimensional Atmospheric Fields on
2936 Limited-Area Domains Using the Discrete Cosine Transform (DCT)"
2937 Monthly Weather Review, Vol. 130, 1812-1828
2938 doi: https://doi.org/10.1175/1520-0493(2002)130<1812:SDOTDA>2.0.CO;2
2939 """
2940 Nt, Ny, Nx = y_3d.shape
2942 # Max coefficient
2943 Nmin = min(Nx - 1, Ny - 1)
2945 # Create alpha matrix (of wavenumbers)
2946 alpha_matrix = _create_alpha_matrix(Ny, Nx)
2948 # Prepare output array
2949 ps_array = np.zeros((Nt, Nmin))
2951 # Loop over time to get spectrum for each time.
2952 for t in range(Nt):
2953 y_2d = y_3d[t]
2955 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space.
2956 # fkk is a 2D array of DCT coefficients, representing the amplitudes of
2957 # cosine basis functions at different spatial frequencies.
2959 # normalise spectrum to allow comparison between models.
2960 fkk = fft.dctn(y_2d, norm="ortho")
2962 # Normalise fkk
2963 fkk = fkk / np.sqrt(Ny * Nx)
2965 # calculate variance of spectral coefficient
2966 sigma_2 = fkk**2 / Nx / Ny
2968 # Group ellipses of alphas into the same wavenumber k/Nmin
2969 for k in range(1, Nmin + 1):
2970 alpha = k / Nmin
2971 alpha_p1 = (k + 1) / Nmin
2973 # Sum up elements matching k
2974 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1))
2975 ps_array[t, k - 1] = np.sum(sigma_2[mask_k])
2977 return ps_array
2980def _create_alpha_matrix(Ny, Nx):
2981 """Construct an array of 2D wavenumbers from 2D wavenumber pair.
2983 Parameters
2984 ----------
2985 Ny, Nx:
2986 Dimensions of the 2D field for which the power spectra is calculated. Used to
2987 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a
2988 single-scale parameter.
2990 Returns: alpha_matrix
2991 normalisation of 2D wavenumber axes, transforming the spectral domain into
2992 an elliptic coordinate system.
2994 """
2995 # Create x_indices: each row is [1, 2, ..., Nx]
2996 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1))
2998 # Create y_indices: each column is [1, 2, ..., Ny]
2999 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx))
3001 # Compute alpha_matrix
3002 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2)
3004 return alpha_matrix