Coverage for src / CSET / operators / plot.py: 83%
1008 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-03-31 13:21 +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
25from typing import Literal
27import cartopy.crs as ccrs
28import iris
29import iris.coords
30import iris.cube
31import iris.exceptions
32import iris.plot as iplt
33import matplotlib as mpl
34import matplotlib.colors as mcolors
35import matplotlib.pyplot as plt
36import numpy as np
37import scipy.fft as fft
38from iris.cube import Cube
39from markdown_it import MarkdownIt
41from CSET._common import (
42 combine_dicts,
43 filename_slugify,
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 slice_over_maybe,
54)
55from CSET.operators.collapse import collapse
56from CSET.operators.misc import _extract_common_time_points
57from CSET.operators.regrid import regrid_onto_cube
59# Use a non-interactive plotting backend.
60mpl.use("agg")
62DEFAULT_DISCRETE_COLORS = mpl.colormaps["tab10"].colors + mpl.colormaps["Accent"].colors
64############################
65# Private helper functions #
66############################
69def _append_to_plot_index(plot_index: list) -> list:
70 """Add plots into the plot index, returning the complete plot index."""
71 with open("meta.json", "r+t", encoding="UTF-8") as fp:
72 fcntl.flock(fp, fcntl.LOCK_EX)
73 fp.seek(0)
74 meta = json.load(fp)
75 complete_plot_index = meta.get("plots", [])
76 complete_plot_index = complete_plot_index + plot_index
77 meta["plots"] = complete_plot_index
78 if os.getenv("CYLC_TASK_CYCLE_POINT") and not bool(
79 os.getenv("DO_CASE_AGGREGATION")
80 ):
81 meta["case_date"] = os.getenv("CYLC_TASK_CYCLE_POINT", "")
82 fp.seek(0)
83 fp.truncate()
84 json.dump(meta, fp, indent=2)
85 return complete_plot_index
88def _check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube:
89 """Ensure a single cube is given.
91 If a CubeList of length one is given that the contained cube is returned,
92 otherwise an error is raised.
94 Parameters
95 ----------
96 cube: Cube | CubeList
97 The cube to check.
99 Returns
100 -------
101 cube: Cube
102 The checked cube.
104 Raises
105 ------
106 TypeError
107 If the input cube is not a Cube or CubeList of a single Cube.
108 """
109 if isinstance(cube, iris.cube.Cube):
110 return cube
111 if isinstance(cube, iris.cube.CubeList):
112 if len(cube) == 1:
113 return cube[0]
114 raise TypeError("Must have a single cube", cube)
117def _make_plot_html_page(plots: list):
118 """Create a HTML page to display a plot image."""
119 # Debug check that plots actually contains some strings.
120 assert isinstance(plots[0], str)
122 # Load HTML template file.
123 operator_files = importlib.resources.files()
124 template_file = operator_files.joinpath("_plot_page_template.html")
126 # Get some metadata.
127 meta = get_recipe_metadata()
128 title = meta.get("title", "Untitled")
129 description = MarkdownIt().render(meta.get("description", "*No description.*"))
131 # Prepare template variables.
132 variables = {
133 "title": title,
134 "description": description,
135 "initial_plot": plots[0],
136 "plots": plots,
137 "title_slug": slugify(title),
138 }
140 # Render template.
141 html = render_file(template_file, **variables)
143 # Save completed HTML.
144 with open("index.html", "wt", encoding="UTF-8") as fp:
145 fp.write(html)
148@functools.cache
149def _load_colorbar_map(user_colorbar_file: str = None) -> dict:
150 """Load the colorbar definitions from a file.
152 This is a separate function to make it cacheable.
153 """
154 colorbar_file = importlib.resources.files().joinpath("_colorbar_definition.json")
155 with open(colorbar_file, "rt", encoding="UTF-8") as fp:
156 colorbar = json.load(fp)
158 logging.debug("User colour bar file: %s", user_colorbar_file)
159 override_colorbar = {}
160 if user_colorbar_file:
161 try:
162 with open(user_colorbar_file, "rt", encoding="UTF-8") as fp:
163 override_colorbar = json.load(fp)
164 except FileNotFoundError:
165 logging.warning("Colorbar file does not exist. Using default values.")
167 # Overwrite values with the user supplied colorbar definition.
168 colorbar = combine_dicts(colorbar, override_colorbar)
169 return colorbar
172def _get_model_colors_map(cubes: iris.cube.CubeList | iris.cube.Cube) -> dict:
173 """Get an appropriate colors for model lines in line plots.
175 For each model in the list of cubes colors either from user provided
176 color definition file (so-called style file) or from default colors are mapped
177 to model_name attribute.
179 Parameters
180 ----------
181 cubes: CubeList or Cube
182 Cubes with model_name attribute
184 Returns
185 -------
186 model_colors_map:
187 Dictionary mapping model_name attribute to colors
188 """
189 user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
190 colorbar = _load_colorbar_map(user_colorbar_file)
191 model_names = sorted(
192 filter(
193 lambda x: x is not None,
194 (cube.attributes.get("model_name", None) for cube in iter_maybe(cubes)),
195 )
196 )
197 if not model_names:
198 return {}
199 use_user_colors = all(mname in colorbar.keys() for mname in model_names)
200 if use_user_colors: 200 ↛ 201line 200 didn't jump to line 201 because the condition on line 200 was never true
201 return {mname: colorbar[mname] for mname in model_names}
203 color_list = itertools.cycle(DEFAULT_DISCRETE_COLORS)
204 return {mname: color for mname, color in zip(model_names, color_list, strict=False)}
207def _colorbar_map_levels(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
208 """Get an appropriate colorbar for the given cube.
210 For the given variable the appropriate colorbar is looked up from a
211 combination of the built-in CSET colorbar definitions, and any user supplied
212 definitions. As well as varying on variables, these definitions may also
213 exist for specific pressure levels to account for variables with
214 significantly different ranges at different heights. The colorbars also exist
215 for masks and mask differences for considering variable presence diagnostics.
216 Specific variable ranges can be separately set in user-supplied definition
217 for x- or y-axis limits, or indicate where automated range preferred.
219 Parameters
220 ----------
221 cube: Cube
222 Cube of variable for which the colorbar information is desired.
223 axis: "x", "y", optional
224 Select the levels for just this axis of a line plot. The min and max
225 can be set by xmin/xmax or ymin/ymax respectively. For variables where
226 setting a universal range is not desirable (e.g. temperature), users
227 can set ymin/ymax values to "auto" in the colorbar definitions file.
228 Where no additional xmin/xmax or ymin/ymax values are provided, the
229 axis bounds default to use the vmin/vmax values provided.
231 Returns
232 -------
233 cmap:
234 Matplotlib colormap.
235 levels:
236 List of levels to use for plotting. For continuous plots the min and max
237 should be taken as the range.
238 norm:
239 BoundaryNorm information.
240 """
241 # Grab the colorbar file from the recipe global metadata.
242 user_colorbar_file = get_recipe_metadata().get("style_file_path", None)
243 colorbar = _load_colorbar_map(user_colorbar_file)
244 cmap = None
246 try:
247 # We assume that pressure is a scalar coordinate here.
248 pressure_level_raw = cube.coord("pressure").points[0]
249 # Ensure pressure_level is a string, as it is used as a JSON key.
250 pressure_level = str(int(pressure_level_raw))
251 except iris.exceptions.CoordinateNotFoundError:
252 pressure_level = None
254 # First try long name, then standard name, then var name. This order is used
255 # as long name is the one we correct between models, so it most likely to be
256 # consistent.
257 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name]))
258 for varname in varnames:
259 # Get the colormap for this variable.
260 try:
261 var_colorbar = colorbar[varname]
262 cmap = plt.get_cmap(colorbar[varname]["cmap"], 51)
263 varname_key = varname
264 break
265 except KeyError:
266 logging.debug("Cube name %s has no colorbar definition.", varname)
268 # Get colormap if it is a mask.
269 if any("mask_for_" in name for name in varnames):
270 cmap, levels, norm = _custom_colormap_mask(cube, axis=axis)
271 return cmap, levels, norm
272 # If winds on Beaufort Scale use custom colorbar and levels
273 if any("Beaufort_Scale" in name for name in varnames):
274 cmap, levels, norm = _custom_beaufort_scale(cube, axis=axis)
275 return cmap, levels, norm
276 # If probability is plotted use custom colorbar and levels
277 if any("probability_of_" in name for name in varnames):
278 cmap, levels, norm = _custom_colormap_probability(cube, axis=axis)
279 return cmap, levels, norm
280 # If aviation colour state use custom colorbar and levels
281 if any("aviation_colour_state" in name for name in varnames): 281 ↛ 282line 281 didn't jump to line 282 because the condition on line 281 was never true
282 cmap, levels, norm = _custom_colormap_aviation_colour_state(cube)
283 return cmap, levels, norm
285 # If no valid colormap has been defined, use defaults and return.
286 if not cmap:
287 logging.warning("No colorbar definition exists for %s.", cube.name())
288 cmap, levels, norm = mpl.colormaps["viridis"], None, None
289 return cmap, levels, norm
291 # Test if pressure-level specific settings are provided for cube.
292 if pressure_level:
293 try:
294 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level]
295 except KeyError:
296 logging.debug(
297 "%s has no colorbar definition for pressure level %s.",
298 varname,
299 pressure_level,
300 )
302 # Check for availability of x-axis or y-axis user-specific overrides
303 # for setting level bounds for line plot types and return just levels.
304 # Line plots do not need a colormap, and just use the data range.
305 if axis:
306 if axis == "x":
307 try:
308 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"]
309 except KeyError:
310 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
311 if axis == "y":
312 try:
313 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"]
314 except KeyError:
315 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
316 # Check if user-specified auto-scaling for this variable
317 if vmin == "auto" or vmax == "auto":
318 levels = None
319 else:
320 levels = [vmin, vmax]
321 return None, levels, None
322 # Get and use the colorbar levels for this variable if spatial or histogram.
323 else:
324 try:
325 levels = var_colorbar["levels"]
326 # Use discrete bins when levels are specified, rather
327 # than a smooth range.
328 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N)
329 logging.debug("Using levels for %s colorbar.", varname)
330 logging.info("Using levels: %s", levels)
331 except KeyError:
332 # Get the range for this variable.
333 vmin, vmax = var_colorbar["min"], var_colorbar["max"]
334 logging.debug("Using min and max for %s colorbar.", varname)
335 # Calculate levels from range.
336 levels = np.linspace(vmin, vmax, 101)
337 norm = None
339 # Overwrite cmap, levels and norm for specific variables that
340 # require custom colorbar_map as these can not be defined in the
341 # JSON file.
342 cmap, levels, norm = _custom_colourmap_precipitation(cube, cmap, levels, norm)
343 cmap, levels, norm = _custom_colourmap_visibility_in_air(
344 cube, cmap, levels, norm
345 )
346 cmap, levels, norm = _custom_colormap_celsius(cube, cmap, levels, norm)
347 return cmap, levels, norm
350def _setup_spatial_map(
351 cube: iris.cube.Cube,
352 figure,
353 cmap,
354 grid_size: int | None = None,
355 subplot: int | None = None,
356):
357 """Define map projections, extent and add coastlines for spatial plots.
359 For spatial map plots, a relevant map projection for rotated or non-rotated inputs
360 is specified, and map extent defined based on the input data.
362 Parameters
363 ----------
364 cube: Cube
365 2 dimensional (lat and lon) Cube of the data to plot.
366 figure:
367 Matplotlib Figure object holding all plot elements.
368 cmap:
369 Matplotlib colormap.
370 grid_size: int, optional
371 Size of grid for subplots if multiple spatial subplots in figure.
372 subplot: int, optional
373 Subplot index if multiple spatial subplots in figure.
375 Returns
376 -------
377 axes:
378 Matplotlib GeoAxes definition.
379 """
380 # Identify min/max plot bounds.
381 try:
382 lat_axis, lon_axis = get_cube_yxcoordname(cube)
383 x1 = np.min(cube.coord(lon_axis).points)
384 x2 = np.max(cube.coord(lon_axis).points)
385 y1 = np.min(cube.coord(lat_axis).points)
386 y2 = np.max(cube.coord(lat_axis).points)
388 # Adjust bounds within +/- 180.0 if x dimension extends beyond half-globe.
389 if np.abs(x2 - x1) > 180.0:
390 x1 = x1 - 180.0
391 x2 = x2 - 180.0
392 logging.debug("Adjusting plot bounds to fit global extent.")
394 # Consider map projection orientation.
395 # Adapting orientation enables plotting across international dateline.
396 # Users can adapt the default central_longitude if alternative projections views.
397 if x2 > 180.0 or x1 < -180.0:
398 central_longitude = 180.0
399 else:
400 central_longitude = 0.0
402 # Define spatial map projection.
403 coord_system = cube.coord(lat_axis).coord_system
404 if isinstance(coord_system, iris.coord_systems.RotatedGeogCS):
405 # Define rotated pole map projection for rotated pole inputs.
406 projection = ccrs.RotatedPole(
407 pole_longitude=coord_system.grid_north_pole_longitude,
408 pole_latitude=coord_system.grid_north_pole_latitude,
409 central_rotated_longitude=central_longitude,
410 )
411 crs = projection
412 elif isinstance(coord_system, iris.coord_systems.TransverseMercator): 412 ↛ 414line 412 didn't jump to line 414 because the condition on line 412 was never true
413 # Define Transverse Mercator projection for TM inputs.
414 projection = ccrs.TransverseMercator(
415 central_longitude=coord_system.longitude_of_central_meridian,
416 central_latitude=coord_system.latitude_of_projection_origin,
417 false_easting=coord_system.false_easting,
418 false_northing=coord_system.false_northing,
419 scale_factor=coord_system.scale_factor_at_central_meridian,
420 )
421 crs = projection
422 else:
423 # Define regular map projection for non-rotated pole inputs.
424 # Alternatives might include e.g. for global model outputs:
425 # projection=ccrs.Robinson(central_longitude=X.y, globe=None)
426 # See also https://scitools.org.uk/cartopy/docs/v0.15/crs/projections.html.
427 projection = ccrs.PlateCarree(central_longitude=central_longitude)
428 crs = ccrs.PlateCarree()
430 # Define axes for plot (or subplot) with required map projection.
431 if subplot is not None:
432 axes = figure.add_subplot(
433 grid_size, grid_size, subplot, projection=projection
434 )
435 else:
436 axes = figure.add_subplot(projection=projection)
438 # Add coastlines if cube contains x and y map coordinates.
439 if cmap.name in ["viridis", "Greys"]:
440 coastcol = "magenta"
441 else:
442 coastcol = "black"
443 logging.debug("Plotting coastlines in colour %s.", coastcol)
444 axes.coastlines(resolution="10m", color=coastcol)
446 # If is lat/lon spatial map, fix extent to keep plot tight.
447 # Specifying crs within set_extent helps ensure only data region is shown.
448 if isinstance(coord_system, iris.coord_systems.GeogCS):
449 axes.set_extent([x1, x2, y1, y2], crs=crs)
451 except ValueError:
452 # Skip if not both x and y map coordinates.
453 axes = figure.gca()
454 pass
456 return axes
459def _get_plot_resolution() -> int:
460 """Get resolution of rasterised plots in pixels per inch."""
461 return get_recipe_metadata().get("plot_resolution", 100)
464def _set_title_and_filename(
465 seq_coord: iris.coords.Coord,
466 nplot: int,
467 recipe_title: str,
468 filename: str,
469):
470 """Set plot title and filename based on cube coordinate.
472 Parameters
473 ----------
474 sequence_coordinate: iris.coords.Coord
475 Coordinate about which to make a plot sequence.
476 nplot: int
477 Number of output plots to generate - controls title/naming.
478 recipe_title: str
479 Default plot title, potentially to update.
480 filename: str
481 Input plot filename, potentially to update.
483 Returns
484 -------
485 plot_title: str
486 Output formatted plot title string, based on plotted data.
487 plot_filename: str
488 Output formatted plot filename string.
489 """
490 ndim = seq_coord.ndim
491 npoints = np.size(seq_coord.points)
492 sequence_title = ""
493 sequence_fname = ""
495 # Account for case with multi-dimension sequence input (e.g. aggregation)
496 if ndim > 1:
497 sequence_title = f"\n [{ndim} cases]"
498 sequence_fname = f"_{ndim}cases"
500 else:
501 if npoints == 1:
502 if nplot > 1:
503 # Set default labels for sequence inputs
504 sequence_value = seq_coord.units.title(seq_coord.points[0])
505 sequence_title = f"\n [{sequence_value}]"
506 sequence_fname = f"_{filename_slugify(sequence_value)}"
507 elif seq_coord.has_bounds():
508 ncase = np.size(seq_coord.bounds)
509 sequence_title = f"\n [{ncase} cases]"
510 sequence_fname = f"_{ncase}cases"
511 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs
512 # Take title endpoints from coord points where series input (e.g. timeseries)
513 if npoints > 1:
514 if not seq_coord.has_bounds():
515 startstring = seq_coord.units.title(seq_coord.points[0])
516 endstring = seq_coord.units.title(seq_coord.points[-1])
517 else:
518 # Take title endpoint from coord bounds where single input (e.g. map)
519 startstring = seq_coord.units.title(seq_coord.bounds.flatten()[0])
520 endstring = seq_coord.units.title(seq_coord.bounds.flatten()[-1])
521 sequence_title = f"\n [{startstring} to {endstring}]"
522 sequence_fname = (
523 f"_{filename_slugify(startstring)}_{filename_slugify(endstring)}"
524 )
526 # Set plot title and filename
527 plot_title = f"{recipe_title}{sequence_title}"
529 # Set plot filename, defaulting to user input if provided.
530 if filename is None:
531 filename = slugify(recipe_title)
532 plot_filename = f"{filename.rsplit('.', 1)[0]}{sequence_fname}.png"
533 else:
534 if nplot > 1:
535 plot_filename = f"{filename.rsplit('.', 1)[0]}{sequence_fname}.png"
536 else:
537 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
539 return plot_title, plot_filename
542def _plot_and_save_spatial_plot(
543 cube: iris.cube.Cube,
544 filename: str,
545 title: str,
546 method: Literal["contourf", "pcolormesh"],
547 overlay_cube: iris.cube.Cube | None = None,
548 contour_cube: iris.cube.Cube | None = None,
549 **kwargs,
550):
551 """Plot and save a spatial plot.
553 Parameters
554 ----------
555 cube: Cube
556 2 dimensional (lat and lon) Cube of the data to plot.
557 filename: str
558 Filename of the plot to write.
559 title: str
560 Plot title.
561 method: "contourf" | "pcolormesh"
562 The plotting method to use.
563 overlay_cube: Cube, optional
564 Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube
565 contour_cube: Cube, optional
566 Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube
567 """
568 # Setup plot details, size, resolution, etc.
569 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
571 # Specify the color bar
572 cmap, levels, norm = _colorbar_map_levels(cube)
574 # If overplotting, set required colorbars
575 if overlay_cube:
576 over_cmap, over_levels, over_norm = _colorbar_map_levels(overlay_cube)
577 if contour_cube:
578 cntr_cmap, cntr_levels, cntr_norm = _colorbar_map_levels(contour_cube)
580 # Setup plot map projection, extent and coastlines.
581 axes = _setup_spatial_map(cube, fig, cmap)
583 # Plot the field.
584 if method == "contourf":
585 # Filled contour plot of the field.
586 plot = iplt.contourf(cube, cmap=cmap, levels=levels, norm=norm)
587 elif method == "pcolormesh":
588 try:
589 vmin = min(levels)
590 vmax = max(levels)
591 except TypeError:
592 vmin, vmax = None, None
593 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax
594 # if levels are defined.
595 if norm is not None:
596 vmin = None
597 vmax = None
598 logging.debug("Plotting using defined levels.")
599 plot = iplt.pcolormesh(cube, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax)
600 else:
601 raise ValueError(f"Unknown plotting method: {method}")
603 # Overplot overlay field, if required
604 if overlay_cube:
605 try:
606 over_vmin = min(over_levels)
607 over_vmax = max(over_levels)
608 except TypeError:
609 over_vmin, over_vmax = None, None
610 if over_norm is not None: 610 ↛ 611line 610 didn't jump to line 611 because the condition on line 610 was never true
611 over_vmin = None
612 over_vmax = None
613 overlay = iplt.pcolormesh(
614 overlay_cube,
615 cmap=over_cmap,
616 norm=over_norm,
617 alpha=0.8,
618 vmin=over_vmin,
619 vmax=over_vmax,
620 )
621 # Overplot contour field, if required, with contour labelling.
622 if contour_cube:
623 contour = iplt.contour(
624 contour_cube,
625 colors="darkgray",
626 levels=cntr_levels,
627 norm=cntr_norm,
628 alpha=0.5,
629 linestyles="--",
630 linewidths=1,
631 )
632 plt.clabel(contour)
634 # Check to see if transect, and if so, adjust y axis.
635 if is_transect(cube):
636 if "pressure" in [coord.name() for coord in cube.coords()]:
637 axes.invert_yaxis()
638 axes.set_yscale("log")
639 axes.set_ylim(1100, 100)
640 # If both model_level_number and level_height exists, iplt can construct
641 # plot as a function of height above orography (NOT sea level).
642 elif {"model_level_number", "level_height"}.issubset( 642 ↛ 647line 642 didn't jump to line 647 because the condition on line 642 was always true
643 {coord.name() for coord in cube.coords()}
644 ):
645 axes.set_yscale("log")
647 axes.set_title(
648 f"{title}\n"
649 f"Start Lat: {cube.attributes['transect_coords'].split('_')[0]}"
650 f" Start Lon: {cube.attributes['transect_coords'].split('_')[1]}"
651 f" End Lat: {cube.attributes['transect_coords'].split('_')[2]}"
652 f" End Lon: {cube.attributes['transect_coords'].split('_')[3]}",
653 fontsize=16,
654 )
656 else:
657 # Add title.
658 axes.set_title(title, fontsize=16)
660 # Add watermark with min/max/mean. Currently not user togglable.
661 # In the bbox dictionary, fc and ec are hex colour codes for grey shade.
662 axes.annotate(
663 f"Min: {np.min(cube.data):.3g} Max: {np.max(cube.data):.3g} Mean: {np.mean(cube.data):.3g}",
664 xy=(1, -0.05),
665 xycoords="axes fraction",
666 xytext=(-5, 5),
667 textcoords="offset points",
668 ha="right",
669 va="bottom",
670 size=11,
671 bbox=dict(boxstyle="round", fc="#cccccc", ec="#808080", alpha=0.9),
672 )
674 # Add secondary colour bar for overlay_cube field if required.
675 if overlay_cube:
676 cbarB = fig.colorbar(
677 overlay, orientation="horizontal", location="bottom", pad=0.0, shrink=0.7
678 )
679 cbarB.set_label(label=f"{overlay_cube.name()} ({overlay_cube.units})", size=14)
680 # add ticks and tick_labels for every levels if less than 20 levels exist
681 if over_levels is not None and len(over_levels) < 20: 681 ↛ 682line 681 didn't jump to line 682 because the condition on line 681 was never true
682 cbarB.set_ticks(over_levels)
683 cbarB.set_ticklabels([f"{level:.2f}" for level in over_levels])
684 if "rainfall" or "snowfall" or "visibility" in overlay_cube.name():
685 cbarB.set_ticklabels([f"{level:.3g}" for level in over_levels])
686 logging.debug("Set secondary colorbar ticks and labels.")
688 # Add main colour bar.
689 cbar = fig.colorbar(
690 plot, orientation="horizontal", location="bottom", pad=0.042, shrink=0.7
691 )
692 cbar.set_label(label=f"{cube.name()} ({cube.units})", size=14)
693 # add ticks and tick_labels for every levels if less than 20 levels exist
694 if levels is not None and len(levels) < 20:
695 cbar.set_ticks(levels)
696 cbar.set_ticklabels([f"{level:.2f}" for level in levels])
697 if "rainfall" or "snowfall" or "visibility" in cube.name(): 697 ↛ 699line 697 didn't jump to line 699 because the condition on line 697 was always true
698 cbar.set_ticklabels([f"{level:.3g}" for level in levels])
699 logging.debug("Set colorbar ticks and labels.")
701 # Save plot.
702 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
703 logging.info("Saved spatial plot to %s", filename)
704 plt.close(fig)
707def _plot_and_save_postage_stamp_spatial_plot(
708 cube: iris.cube.Cube,
709 filename: str,
710 stamp_coordinate: str,
711 title: str,
712 method: Literal["contourf", "pcolormesh"],
713 overlay_cube: iris.cube.Cube | None = None,
714 contour_cube: iris.cube.Cube | None = None,
715 **kwargs,
716):
717 """Plot postage stamp spatial plots from an ensemble.
719 Parameters
720 ----------
721 cube: Cube
722 Iris cube of data to be plotted. It must have the stamp coordinate.
723 filename: str
724 Filename of the plot to write.
725 stamp_coordinate: str
726 Coordinate that becomes different plots.
727 method: "contourf" | "pcolormesh"
728 The plotting method to use.
729 overlay_cube: Cube, optional
730 Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube
731 contour_cube: Cube, optional
732 Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube
734 Raises
735 ------
736 ValueError
737 If the cube doesn't have the right dimensions.
738 """
739 # Use the smallest square grid that will fit the members.
740 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
742 fig = plt.figure(figsize=(10, 10))
744 # Specify the color bar
745 cmap, levels, norm = _colorbar_map_levels(cube)
746 # If overplotting, set required colorbars
747 if overlay_cube: 747 ↛ 748line 747 didn't jump to line 748 because the condition on line 747 was never true
748 over_cmap, over_levels, over_norm = _colorbar_map_levels(overlay_cube)
749 if contour_cube: 749 ↛ 750line 749 didn't jump to line 750 because the condition on line 749 was never true
750 cntr_cmap, cntr_levels, cntr_norm = _colorbar_map_levels(contour_cube)
752 # Make a subplot for each member.
753 for member, subplot in zip(
754 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
755 ):
756 # Setup subplot map projection, extent and coastlines.
757 axes = _setup_spatial_map(
758 member, fig, cmap, grid_size=grid_size, subplot=subplot
759 )
760 if method == "contourf":
761 # Filled contour plot of the field.
762 plot = iplt.contourf(member, cmap=cmap, levels=levels, norm=norm)
763 elif method == "pcolormesh":
764 if levels is not None:
765 vmin = min(levels)
766 vmax = max(levels)
767 else:
768 raise TypeError("Unknown vmin and vmax range.")
769 vmin, vmax = None, None
770 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax
771 # if levels are defined.
772 if norm is not None: 772 ↛ 773line 772 didn't jump to line 773 because the condition on line 772 was never true
773 vmin = None
774 vmax = None
775 # pcolormesh plot of the field.
776 plot = iplt.pcolormesh(member, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax)
777 else:
778 raise ValueError(f"Unknown plotting method: {method}")
780 # Overplot overlay field, if required
781 if overlay_cube: 781 ↛ 782line 781 didn't jump to line 782 because the condition on line 781 was never true
782 try:
783 over_vmin = min(over_levels)
784 over_vmax = max(over_levels)
785 except TypeError:
786 over_vmin, over_vmax = None, None
787 if over_norm is not None:
788 over_vmin = None
789 over_vmax = None
790 iplt.pcolormesh(
791 overlay_cube[member.coord(stamp_coordinate).points[0]],
792 cmap=over_cmap,
793 norm=over_norm,
794 alpha=0.6,
795 vmin=over_vmin,
796 vmax=over_vmax,
797 )
798 # Overplot contour field, if required
799 if contour_cube: 799 ↛ 800line 799 didn't jump to line 800 because the condition on line 799 was never true
800 iplt.contour(
801 contour_cube[member.coord(stamp_coordinate).points[0]],
802 colors="darkgray",
803 levels=cntr_levels,
804 norm=cntr_norm,
805 alpha=0.6,
806 linestyles="--",
807 linewidths=1,
808 )
809 axes.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
810 axes.set_axis_off()
812 # Put the shared colorbar in its own axes.
813 colorbar_axes = fig.add_axes([0.15, 0.07, 0.7, 0.03])
814 colorbar = fig.colorbar(
815 plot, colorbar_axes, orientation="horizontal", pad=0.042, shrink=0.7
816 )
817 colorbar.set_label(f"{cube.name()} ({cube.units})", size=14)
819 # Overall figure title.
820 fig.suptitle(title, fontsize=16)
822 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
823 logging.info("Saved contour postage stamp plot to %s", filename)
824 plt.close(fig)
827def _plot_and_save_line_series(
828 cubes: iris.cube.CubeList,
829 coords: list[iris.coords.Coord],
830 ensemble_coord: str,
831 filename: str,
832 title: str,
833 **kwargs,
834):
835 """Plot and save a 1D line series.
837 Parameters
838 ----------
839 cubes: Cube or CubeList
840 Cube or CubeList containing the cubes to plot on the y-axis.
841 coords: list[Coord]
842 Coordinates to plot on the x-axis, one per cube.
843 ensemble_coord: str
844 Ensemble coordinate in the cube.
845 filename: str
846 Filename of the plot to write.
847 title: str
848 Plot title.
849 """
850 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
852 model_colors_map = _get_model_colors_map(cubes)
854 # Store min/max ranges.
855 y_levels = []
857 # Check match-up across sequence coords gives consistent sizes
858 _validate_cubes_coords(cubes, coords)
860 for cube, coord in zip(cubes, coords, strict=True):
861 label = None
862 color = "black"
863 if model_colors_map:
864 label = cube.attributes.get("model_name")
865 color = model_colors_map.get(label)
866 for cube_slice in cube.slices_over(ensemble_coord):
867 # Label with (control) if part of an ensemble or not otherwise.
868 if cube_slice.coord(ensemble_coord).points == [0]:
869 iplt.plot(
870 coord,
871 cube_slice,
872 color=color,
873 marker="o",
874 ls="-",
875 lw=3,
876 label=f"{label} (control)"
877 if len(cube.coord(ensemble_coord).points) > 1
878 else label,
879 )
880 # Label with (perturbed) if part of an ensemble and not the control.
881 else:
882 iplt.plot(
883 coord,
884 cube_slice,
885 color=color,
886 ls="-",
887 lw=1.5,
888 alpha=0.75,
889 label=f"{label} (member)",
890 )
892 # Calculate the global min/max if multiple cubes are given.
893 _, levels, _ = _colorbar_map_levels(cube, axis="y")
894 if levels is not None: 894 ↛ 895line 894 didn't jump to line 895 because the condition on line 894 was never true
895 y_levels.append(min(levels))
896 y_levels.append(max(levels))
898 # Get the current axes.
899 ax = plt.gca()
901 # Add some labels and tweak the style.
902 # check if cubes[0] works for single cube if not CubeList
903 if coords[0].name() == "time":
904 ax.set_xlabel(f"{coords[0].name()}", fontsize=14)
905 else:
906 ax.set_xlabel(f"{coords[0].name()} / {coords[0].units}", fontsize=14)
907 ax.set_ylabel(f"{cubes[0].name()} / {cubes[0].units}", fontsize=14)
908 ax.set_title(title, fontsize=16)
910 ax.ticklabel_format(axis="y", useOffset=False)
911 ax.tick_params(axis="x", labelrotation=15)
912 ax.tick_params(axis="both", labelsize=12)
914 # Set y limits to global min and max, autoscale if colorbar doesn't exist.
915 if y_levels: 915 ↛ 916line 915 didn't jump to line 916 because the condition on line 915 was never true
916 ax.set_ylim(min(y_levels), max(y_levels))
917 # Add zero line.
918 if min(y_levels) < 0.0 and max(y_levels) > 0.0:
919 ax.axhline(y=0, xmin=0, xmax=1, ls="-", color="grey", lw=2)
920 logging.debug(
921 "Line plot with y-axis limits %s-%s", min(y_levels), max(y_levels)
922 )
923 else:
924 ax.autoscale()
926 # Add gridlines
927 ax.grid(linestyle="--", color="grey", linewidth=1)
928 # Ientify unique labels for legend
929 handles = list(
930 {
931 label: handle
932 for (handle, label) in zip(*ax.get_legend_handles_labels(), strict=True)
933 }.values()
934 )
935 ax.legend(handles=handles, loc="best", ncol=1, frameon=False, fontsize=16)
937 # Save plot.
938 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
939 logging.info("Saved line plot to %s", filename)
940 plt.close(fig)
943def _plot_and_save_vertical_line_series(
944 cubes: iris.cube.CubeList,
945 coords: list[iris.coords.Coord],
946 ensemble_coord: str,
947 filename: str,
948 series_coordinate: str,
949 title: str,
950 vmin: float,
951 vmax: float,
952 **kwargs,
953):
954 """Plot and save a 1D line series in vertical.
956 Parameters
957 ----------
958 cubes: CubeList
959 1 dimensional Cube or CubeList of the data to plot on x-axis.
960 coord: list[Coord]
961 Coordinates to plot on the y-axis, one per cube.
962 ensemble_coord: str
963 Ensemble coordinate in the cube.
964 filename: str
965 Filename of the plot to write.
966 series_coordinate: str
967 Coordinate to use as vertical axis.
968 title: str
969 Plot title.
970 vmin: float
971 Minimum value for the x-axis.
972 vmax: float
973 Maximum value for the x-axis.
974 """
975 # plot the vertical pressure axis using log scale
976 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
978 model_colors_map = _get_model_colors_map(cubes)
980 # Check match-up across sequence coords gives consistent sizes
981 _validate_cubes_coords(cubes, coords)
983 for cube, coord in zip(cubes, coords, strict=True):
984 label = None
985 color = "black"
986 if model_colors_map: 986 ↛ 987line 986 didn't jump to line 987 because the condition on line 986 was never true
987 label = cube.attributes.get("model_name")
988 color = model_colors_map.get(label)
990 for cube_slice in cube.slices_over(ensemble_coord):
991 # If ensemble data given plot control member with (control)
992 # unless single forecast.
993 if cube_slice.coord(ensemble_coord).points == [0]:
994 iplt.plot(
995 cube_slice,
996 coord,
997 color=color,
998 marker="o",
999 ls="-",
1000 lw=3,
1001 label=f"{label} (control)"
1002 if len(cube.coord(ensemble_coord).points) > 1
1003 else label,
1004 )
1005 # If ensemble data given plot perturbed members with (perturbed).
1006 else:
1007 iplt.plot(
1008 cube_slice,
1009 coord,
1010 color=color,
1011 ls="-",
1012 lw=1.5,
1013 alpha=0.75,
1014 label=f"{label} (member)",
1015 )
1017 # Get the current axis
1018 ax = plt.gca()
1020 # Special handling for pressure level data.
1021 if series_coordinate == "pressure": 1021 ↛ 1043line 1021 didn't jump to line 1043 because the condition on line 1021 was always true
1022 # Invert y-axis and set to log scale.
1023 ax.invert_yaxis()
1024 ax.set_yscale("log")
1026 # Define y-ticks and labels for pressure log axis.
1027 y_tick_labels = [
1028 "1000",
1029 "850",
1030 "700",
1031 "500",
1032 "300",
1033 "200",
1034 "100",
1035 ]
1036 y_ticks = [1000, 850, 700, 500, 300, 200, 100]
1038 # Set y-axis limits and ticks.
1039 ax.set_ylim(1100, 100)
1041 # Test if series_coordinate is model level data. The UM data uses
1042 # model_level_number and lfric uses full_levels as coordinate.
1043 elif series_coordinate in ("model_level_number", "full_levels", "half_levels"):
1044 # Define y-ticks and labels for vertical axis.
1045 y_ticks = iter_maybe(cubes)[0].coord(series_coordinate).points
1046 y_tick_labels = [str(int(i)) for i in y_ticks]
1047 ax.set_ylim(min(y_ticks), max(y_ticks))
1049 ax.set_yticks(y_ticks)
1050 ax.set_yticklabels(y_tick_labels)
1052 # Set x-axis limits.
1053 ax.set_xlim(vmin, vmax)
1054 # Mark y=0 if present in plot.
1055 if vmin < 0.0 and vmax > 0.0: 1055 ↛ 1056line 1055 didn't jump to line 1056 because the condition on line 1055 was never true
1056 ax.axvline(x=0, ymin=0, ymax=1, ls="-", color="grey", lw=2)
1058 # Add some labels and tweak the style.
1059 ax.set_ylabel(f"{coord.name()} / {coord.units}", fontsize=14)
1060 ax.set_xlabel(
1061 f"{iter_maybe(cubes)[0].name()} / {iter_maybe(cubes)[0].units}", fontsize=14
1062 )
1063 ax.set_title(title, fontsize=16)
1064 ax.ticklabel_format(axis="x")
1065 ax.tick_params(axis="y")
1066 ax.tick_params(axis="both", labelsize=12)
1068 # Add gridlines
1069 ax.grid(linestyle="--", color="grey", linewidth=1)
1070 # Ientify unique labels for legend
1071 handles = list(
1072 {
1073 label: handle
1074 for (handle, label) in zip(*ax.get_legend_handles_labels(), strict=True)
1075 }.values()
1076 )
1077 ax.legend(handles=handles, loc="best", ncol=1, frameon=False, fontsize=16)
1079 # Save plot.
1080 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1081 logging.info("Saved line plot to %s", filename)
1082 plt.close(fig)
1085def _plot_and_save_scatter_plot(
1086 cube_x: iris.cube.Cube | iris.cube.CubeList,
1087 cube_y: iris.cube.Cube | iris.cube.CubeList,
1088 filename: str,
1089 title: str,
1090 one_to_one: bool,
1091 model_names: list[str] = None,
1092 **kwargs,
1093):
1094 """Plot and save a 2D scatter plot.
1096 Parameters
1097 ----------
1098 cube_x: Cube | CubeList
1099 1 dimensional Cube or CubeList of the data to plot on x-axis.
1100 cube_y: Cube | CubeList
1101 1 dimensional Cube or CubeList of the data to plot on y-axis.
1102 filename: str
1103 Filename of the plot to write.
1104 title: str
1105 Plot title.
1106 one_to_one: bool
1107 Whether a 1:1 line is plotted.
1108 """
1109 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1110 # plot the cube_x and cube_y 1D fields as a scatter plot. If they are CubeLists this ensures
1111 # to pair each cube from cube_x with the corresponding cube from cube_y, allowing to iterate
1112 # over the pairs simultaneously.
1114 # Ensure cube_x and cube_y are iterable
1115 cube_x_iterable = iter_maybe(cube_x)
1116 cube_y_iterable = iter_maybe(cube_y)
1118 for cube_x_iter, cube_y_iter in zip(cube_x_iterable, cube_y_iterable, strict=True):
1119 iplt.scatter(cube_x_iter, cube_y_iter)
1120 if one_to_one is True:
1121 plt.plot(
1122 [
1123 np.nanmin([np.nanmin(cube_y.data), np.nanmin(cube_x.data)]),
1124 np.nanmax([np.nanmax(cube_y.data), np.nanmax(cube_x.data)]),
1125 ],
1126 [
1127 np.nanmin([np.nanmin(cube_y.data), np.nanmin(cube_x.data)]),
1128 np.nanmax([np.nanmax(cube_y.data), np.nanmax(cube_x.data)]),
1129 ],
1130 "k",
1131 linestyle="--",
1132 )
1133 ax = plt.gca()
1135 # Add some labels and tweak the style.
1136 if model_names is None:
1137 ax.set_xlabel(f"{cube_x[0].name()} / {cube_x[0].units}", fontsize=14)
1138 ax.set_ylabel(f"{cube_y[0].name()} / {cube_y[0].units}", fontsize=14)
1139 else:
1140 # Add the model names, these should be order of base (x) and other (y).
1141 ax.set_xlabel(
1142 f"{model_names[0]}_{cube_x[0].name()} / {cube_x[0].units}", fontsize=14
1143 )
1144 ax.set_ylabel(
1145 f"{model_names[1]}_{cube_y[0].name()} / {cube_y[0].units}", fontsize=14
1146 )
1147 ax.set_title(title, fontsize=16)
1148 ax.ticklabel_format(axis="y", useOffset=False)
1149 ax.tick_params(axis="x", labelrotation=15)
1150 ax.tick_params(axis="both", labelsize=12)
1151 ax.autoscale()
1153 # Save plot.
1154 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1155 logging.info("Saved scatter plot to %s", filename)
1156 plt.close(fig)
1159def _plot_and_save_vector_plot(
1160 cube_u: iris.cube.Cube,
1161 cube_v: iris.cube.Cube,
1162 filename: str,
1163 title: str,
1164 method: Literal["contourf", "pcolormesh"],
1165 **kwargs,
1166):
1167 """Plot and save a 2D vector plot.
1169 Parameters
1170 ----------
1171 cube_u: Cube
1172 2 dimensional Cube of u component of the data.
1173 cube_v: Cube
1174 2 dimensional Cube of v component of the data.
1175 filename: str
1176 Filename of the plot to write.
1177 title: str
1178 Plot title.
1179 """
1180 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1182 # Create a cube containing the magnitude of the vector field.
1183 cube_vec_mag = (cube_u**2 + cube_v**2) ** 0.5
1184 cube_vec_mag.rename(f"{cube_u.name()}_{cube_v.name()}_magnitude")
1186 # Specify the color bar
1187 cmap, levels, norm = _colorbar_map_levels(cube_vec_mag)
1189 # Setup plot map projection, extent and coastlines.
1190 axes = _setup_spatial_map(cube_vec_mag, fig, cmap)
1192 if method == "contourf":
1193 # Filled contour plot of the field.
1194 plot = iplt.contourf(cube_vec_mag, cmap=cmap, levels=levels, norm=norm)
1195 elif method == "pcolormesh":
1196 try:
1197 vmin = min(levels)
1198 vmax = max(levels)
1199 except TypeError:
1200 vmin, vmax = None, None
1201 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax
1202 # if levels are defined.
1203 if norm is not None:
1204 vmin = None
1205 vmax = None
1206 plot = iplt.pcolormesh(cube_vec_mag, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax)
1207 else:
1208 raise ValueError(f"Unknown plotting method: {method}")
1210 # Check to see if transect, and if so, adjust y axis.
1211 if is_transect(cube_vec_mag):
1212 if "pressure" in [coord.name() for coord in cube_vec_mag.coords()]:
1213 axes.invert_yaxis()
1214 axes.set_yscale("log")
1215 axes.set_ylim(1100, 100)
1216 # If both model_level_number and level_height exists, iplt can construct
1217 # plot as a function of height above orography (NOT sea level).
1218 elif {"model_level_number", "level_height"}.issubset(
1219 {coord.name() for coord in cube_vec_mag.coords()}
1220 ):
1221 axes.set_yscale("log")
1223 axes.set_title(
1224 f"{title}\n"
1225 f"Start Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[0]}"
1226 f" Start Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[1]}"
1227 f" End Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[2]}"
1228 f" End Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[3]}",
1229 fontsize=16,
1230 )
1232 else:
1233 # Add title.
1234 axes.set_title(title, fontsize=16)
1236 # Add watermark with min/max/mean. Currently not user togglable.
1237 # In the bbox dictionary, fc and ec are hex colour codes for grey shade.
1238 axes.annotate(
1239 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}",
1240 xy=(1, -0.05),
1241 xycoords="axes fraction",
1242 xytext=(-5, 5),
1243 textcoords="offset points",
1244 ha="right",
1245 va="bottom",
1246 size=11,
1247 bbox=dict(boxstyle="round", fc="#cccccc", ec="#808080", alpha=0.9),
1248 )
1250 # Add colour bar.
1251 cbar = fig.colorbar(plot, orientation="horizontal", pad=0.042, shrink=0.7)
1252 cbar.set_label(label=f"{cube_vec_mag.name()} ({cube_vec_mag.units})", size=14)
1253 # add ticks and tick_labels for every levels if less than 20 levels exist
1254 if levels is not None and len(levels) < 20:
1255 cbar.set_ticks(levels)
1256 cbar.set_ticklabels([f"{level:.1f}" for level in levels])
1258 # 30 barbs along the longest axis of the plot, or a barb per point for data
1259 # with less than 30 points.
1260 step = max(max(cube_u.shape) // 30, 1)
1261 iplt.quiver(cube_u[::step, ::step], cube_v[::step, ::step], pivot="middle")
1263 # Save plot.
1264 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1265 logging.info("Saved vector plot to %s", filename)
1266 plt.close(fig)
1269def _plot_and_save_histogram_series(
1270 cubes: iris.cube.Cube | iris.cube.CubeList,
1271 filename: str,
1272 title: str,
1273 vmin: float,
1274 vmax: float,
1275 **kwargs,
1276):
1277 """Plot and save a histogram series.
1279 Parameters
1280 ----------
1281 cubes: Cube or CubeList
1282 2 dimensional Cube or CubeList of the data to plot as histogram.
1283 filename: str
1284 Filename of the plot to write.
1285 title: str
1286 Plot title.
1287 vmin: float
1288 minimum for colorbar
1289 vmax: float
1290 maximum for colorbar
1291 """
1292 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1293 ax = plt.gca()
1295 model_colors_map = _get_model_colors_map(cubes)
1297 # Set default that histograms will produce probability density function
1298 # at each bin (integral over range sums to 1).
1299 density = True
1301 for cube in iter_maybe(cubes):
1302 # Easier to check title (where var name originates)
1303 # than seeing if long names exist etc.
1304 # Exception case, where distribution better fits log scales/bins.
1305 if "surface_microphysical" in title:
1306 if "amount" in title: 1306 ↛ 1308line 1306 didn't jump to line 1308 because the condition on line 1306 was never true
1307 # Compute histogram following Klingaman et al. (2017): ASoP
1308 bin2 = np.exp(np.log(0.02) + 0.1 * np.linspace(0, 99, 100))
1309 bins = np.pad(bin2, (1, 0), "constant", constant_values=0)
1310 density = False
1311 else:
1312 bins = 10.0 ** (
1313 np.arange(-10, 27, 1) / 10.0
1314 ) # Suggestion from RMED toolbox.
1315 bins = np.insert(bins, 0, 0)
1316 ax.set_yscale("log")
1317 vmin = bins[1]
1318 vmax = bins[-1] # Manually set vmin/vmax to override json derived value.
1319 ax.set_xscale("log")
1320 elif "lightning" in title:
1321 bins = [0, 1, 2, 3, 4, 5]
1322 else:
1323 bins = np.linspace(vmin, vmax, 51)
1324 logging.debug(
1325 "Plotting histogram with %s bins %s - %s.",
1326 np.size(bins),
1327 np.min(bins),
1328 np.max(bins),
1329 )
1331 # Reshape cube data into a single array to allow for a single histogram.
1332 # Otherwise we plot xdim histograms stacked.
1333 cube_data_1d = (cube.data).flatten()
1335 label = None
1336 color = "black"
1337 if model_colors_map: 1337 ↛ 1338line 1337 didn't jump to line 1338 because the condition on line 1337 was never true
1338 label = cube.attributes.get("model_name")
1339 color = model_colors_map[label]
1340 x, y = np.histogram(cube_data_1d, bins=bins, density=density)
1342 # Compute area under curve.
1343 if "surface_microphysical" in title and "amount" in title: 1343 ↛ 1344line 1343 didn't jump to line 1344 because the condition on line 1343 was never true
1344 bin_mean = (bins[:-1] + bins[1:]) / 2.0
1345 x = x * bin_mean / x.sum()
1346 x = x[1:]
1347 y = y[1:]
1349 ax.plot(
1350 y[:-1], x, color=color, linewidth=3, marker="o", markersize=6, label=label
1351 )
1353 # Add some labels and tweak the style.
1354 ax.set_title(title, fontsize=16)
1355 ax.set_xlabel(
1356 f"{iter_maybe(cubes)[0].name()} / {iter_maybe(cubes)[0].units}", fontsize=14
1357 )
1358 ax.set_ylabel("Normalised probability density", fontsize=14)
1359 if "surface_microphysical" in title and "amount" in title: 1359 ↛ 1360line 1359 didn't jump to line 1360 because the condition on line 1359 was never true
1360 ax.set_ylabel(
1361 f"Contribution to mean ({iter_maybe(cubes)[0].units})", fontsize=14
1362 )
1363 ax.set_xlim(vmin, vmax)
1364 ax.tick_params(axis="both", labelsize=12)
1366 # Overlay grid-lines onto histogram plot.
1367 ax.grid(linestyle="--", color="grey", linewidth=1)
1368 if model_colors_map: 1368 ↛ 1369line 1368 didn't jump to line 1369 because the condition on line 1368 was never true
1369 ax.legend(loc="best", ncol=1, frameon=False, fontsize=16)
1371 # Save plot.
1372 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1373 logging.info("Saved histogram plot to %s", filename)
1374 plt.close(fig)
1377def _plot_and_save_postage_stamp_histogram_series(
1378 cube: iris.cube.Cube,
1379 filename: str,
1380 title: str,
1381 stamp_coordinate: str,
1382 vmin: float,
1383 vmax: float,
1384 **kwargs,
1385):
1386 """Plot and save postage (ensemble members) stamps for a histogram series.
1388 Parameters
1389 ----------
1390 cube: Cube
1391 2 dimensional Cube of the data to plot as histogram.
1392 filename: str
1393 Filename of the plot to write.
1394 title: str
1395 Plot title.
1396 stamp_coordinate: str
1397 Coordinate that becomes different plots.
1398 vmin: float
1399 minimum for pdf x-axis
1400 vmax: float
1401 maximum for pdf x-axis
1402 """
1403 # Use the smallest square grid that will fit the members.
1404 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
1406 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1407 # Make a subplot for each member.
1408 for member, subplot in zip(
1409 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
1410 ):
1411 # Implicit interface is much easier here, due to needing to have the
1412 # cartopy GeoAxes generated.
1413 plt.subplot(grid_size, grid_size, subplot)
1414 # Reshape cube data into a single array to allow for a single histogram.
1415 # Otherwise we plot xdim histograms stacked.
1416 member_data_1d = (member.data).flatten()
1417 plt.hist(member_data_1d, density=True, stacked=True)
1418 ax = plt.gca()
1419 ax.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
1420 ax.set_xlim(vmin, vmax)
1422 # Overall figure title.
1423 fig.suptitle(title, fontsize=16)
1425 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1426 logging.info("Saved histogram postage stamp plot to %s", filename)
1427 plt.close(fig)
1430def _plot_and_save_postage_stamps_in_single_plot_histogram_series(
1431 cube: iris.cube.Cube,
1432 filename: str,
1433 title: str,
1434 stamp_coordinate: str,
1435 vmin: float,
1436 vmax: float,
1437 **kwargs,
1438):
1439 fig, ax = plt.subplots(figsize=(10, 10), facecolor="w", edgecolor="k")
1440 ax.set_title(title, fontsize=16)
1441 ax.set_xlim(vmin, vmax)
1442 ax.set_xlabel(f"{cube.name()} / {cube.units}", fontsize=14)
1443 ax.set_ylabel("normalised probability density", fontsize=14)
1444 # Loop over all slices along the stamp_coordinate
1445 for member in cube.slices_over(stamp_coordinate):
1446 # Flatten the member data to 1D
1447 member_data_1d = member.data.flatten()
1448 # Plot the histogram using plt.hist
1449 plt.hist(
1450 member_data_1d,
1451 density=True,
1452 stacked=True,
1453 label=f"Member #{member.coord(stamp_coordinate).points[0]}",
1454 )
1456 # Add a legend
1457 ax.legend(fontsize=16)
1459 # Save the figure to a file
1460 plt.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1461 logging.info("Saved histogram postage stamp plot to %s", filename)
1463 # Close the figure
1464 plt.close(fig)
1467def _plot_and_save_scattermap_plot(
1468 cube: iris.cube.Cube, filename: str, title: str, projection=None, **kwargs
1469):
1470 """Plot and save a geographical scatter plot.
1472 Parameters
1473 ----------
1474 cube: Cube
1475 1 dimensional Cube of the data points with auxiliary latitude and
1476 longitude coordinates,
1477 filename: str
1478 Filename of the plot to write.
1479 title: str
1480 Plot title.
1481 projection: str
1482 Mapping projection to be used by cartopy.
1483 """
1484 # Setup plot details, size, resolution, etc.
1485 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1486 if projection is not None:
1487 # Apart from the default, the only projection we currently support is
1488 # a stereographic projection over the North Pole.
1489 if projection == "NP_Stereo":
1490 axes = plt.axes(projection=ccrs.NorthPolarStereo(central_longitude=0.0))
1491 else:
1492 raise ValueError(f"Unknown projection: {projection}")
1493 else:
1494 axes = plt.axes(projection=ccrs.PlateCarree())
1496 # Scatter plot of the field. The marker size is chosen to give
1497 # symbols that decrease in size as the number of observations
1498 # increases, although the fraction of the figure covered by
1499 # symbols increases roughly as N^(1/2), disregarding overlaps,
1500 # and has been selected for the default figure size of (10, 10).
1501 # Should this be changed, the marker size should be adjusted in
1502 # proportion to the area of the figure.
1503 mrk_size = int(np.sqrt(2500000.0 / len(cube.data)))
1504 klon = None
1505 klat = None
1506 for kc in range(len(cube.aux_coords)):
1507 if cube.aux_coords[kc].standard_name == "latitude":
1508 klat = kc
1509 elif cube.aux_coords[kc].standard_name == "longitude":
1510 klon = kc
1511 scatter_map = iplt.scatter(
1512 cube.aux_coords[klon],
1513 cube.aux_coords[klat],
1514 c=cube.data[:],
1515 s=mrk_size,
1516 cmap="jet",
1517 edgecolors="k",
1518 )
1520 # Add coastlines.
1521 try:
1522 axes.coastlines(resolution="10m")
1523 except AttributeError:
1524 pass
1526 # Add title.
1527 axes.set_title(title, fontsize=16)
1529 # Add colour bar.
1530 cbar = fig.colorbar(scatter_map)
1531 cbar.set_label(label=f"{cube.name()} ({cube.units})", size=20)
1533 # Save plot.
1534 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1535 logging.info("Saved geographical scatter plot to %s", filename)
1536 plt.close(fig)
1539def _plot_and_save_power_spectrum_series(
1540 cubes: iris.cube.Cube | iris.cube.CubeList,
1541 filename: str,
1542 title: str,
1543 **kwargs,
1544):
1545 """Plot and save a power spectrum series.
1547 Parameters
1548 ----------
1549 cubes: Cube or CubeList
1550 2 dimensional Cube or CubeList of the data to plot as power spectrum.
1551 filename: str
1552 Filename of the plot to write.
1553 title: str
1554 Plot title.
1555 """
1556 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1557 ax = plt.gca()
1559 model_colors_map = _get_model_colors_map(cubes)
1561 for cube in iter_maybe(cubes):
1562 # Calculate power spectrum
1564 # Extract time coordinate and convert to datetime
1565 time_coord = cube.coord("time")
1566 time_points = time_coord.units.num2date(time_coord.points)
1568 # Choose one time point (e.g., the first one)
1569 target_time = time_points[0]
1571 # Bind target_time inside the lambda using a default argument
1572 time_constraint = iris.Constraint(
1573 time=lambda cell, target_time=target_time: cell.point == target_time
1574 )
1576 cube = cube.extract(time_constraint)
1578 if cube.ndim == 2:
1579 cube_3d = cube.data[np.newaxis, :, :]
1580 logging.debug("Adding in new axis for a 2 dimensional cube.")
1581 elif cube.ndim == 3: 1581 ↛ 1582line 1581 didn't jump to line 1582 because the condition on line 1581 was never true
1582 cube_3d = cube.data
1583 else:
1584 raise ValueError("Cube dimensions unsuitable for power spectra code")
1585 raise ValueError(
1586 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional."
1587 )
1589 # Calculate spectra
1590 ps_array = _DCT_ps(cube_3d)
1592 ps_cube = iris.cube.Cube(
1593 ps_array,
1594 long_name="power_spectra",
1595 )
1597 ps_cube.attributes["model_name"] = cube.attributes.get("model_name")
1599 # Create a frequency/wavelength array for coordinate
1600 ps_len = ps_cube.data.shape[1]
1601 freqs = np.arange(1, ps_len + 1)
1602 freq_coord = iris.coords.DimCoord(freqs, long_name="frequency", units="1")
1604 # Convert datetime to numeric time using original units
1605 numeric_time = time_coord.units.date2num(time_points)
1606 # Create a new DimCoord with numeric time
1607 new_time_coord = iris.coords.DimCoord(
1608 numeric_time, standard_name="time", units=time_coord.units
1609 )
1611 # Add time and frequency coordinate to spectra cube.
1612 ps_cube.add_dim_coord(new_time_coord.copy(), 0)
1613 ps_cube.add_dim_coord(freq_coord.copy(), 1)
1615 # Extract data from the cube
1616 frequency = ps_cube.coord("frequency").points
1617 power_spectrum = ps_cube.data
1619 label = None
1620 color = "black"
1621 if model_colors_map: 1621 ↛ 1622line 1621 didn't jump to line 1622 because the condition on line 1621 was never true
1622 label = ps_cube.attributes.get("model_name")
1623 color = model_colors_map[label]
1624 ax.plot(frequency, power_spectrum[0], color=color, label=label)
1626 # Add some labels and tweak the style.
1627 ax.set_title(title, fontsize=16)
1628 ax.set_xlabel("Wavenumber", fontsize=14)
1629 ax.set_ylabel("Power", fontsize=14)
1630 ax.tick_params(axis="both", labelsize=12)
1632 # Set log-log scale
1633 ax.set_xscale("log")
1634 ax.set_yscale("log")
1636 # Overlay grid-lines onto power spectrum plot.
1637 ax.grid(linestyle="--", color="grey", linewidth=1)
1638 if model_colors_map: 1638 ↛ 1639line 1638 didn't jump to line 1639 because the condition on line 1638 was never true
1639 ax.legend(loc="best", ncol=1, frameon=False, fontsize=16)
1641 # Save plot.
1642 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1643 logging.info("Saved power spectrum plot to %s", filename)
1644 plt.close(fig)
1647def _plot_and_save_postage_stamp_power_spectrum_series(
1648 cube: iris.cube.Cube,
1649 filename: str,
1650 title: str,
1651 stamp_coordinate: str,
1652 **kwargs,
1653):
1654 """Plot and save postage (ensemble members) stamps for a power spectrum series.
1656 Parameters
1657 ----------
1658 cube: Cube
1659 2 dimensional Cube of the data to plot as power spectrum.
1660 filename: str
1661 Filename of the plot to write.
1662 title: str
1663 Plot title.
1664 stamp_coordinate: str
1665 Coordinate that becomes different plots.
1666 """
1667 # Use the smallest square grid that will fit the members.
1668 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points))))
1670 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k")
1671 # Make a subplot for each member.
1672 for member, subplot in zip(
1673 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False
1674 ):
1675 # Implicit interface is much easier here, due to needing to have the
1676 # cartopy GeoAxes generated.
1677 plt.subplot(grid_size, grid_size, subplot)
1679 frequency = member.coord("frequency").points
1681 ax = plt.gca()
1682 ax.plot(frequency, member.data)
1683 ax.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}")
1685 # Overall figure title.
1686 fig.suptitle(title, fontsize=16)
1688 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1689 logging.info("Saved power spectra postage stamp plot to %s", filename)
1690 plt.close(fig)
1693def _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series(
1694 cube: iris.cube.Cube,
1695 filename: str,
1696 title: str,
1697 stamp_coordinate: str,
1698 **kwargs,
1699):
1700 fig, ax = plt.subplots(figsize=(10, 10), facecolor="w", edgecolor="k")
1701 ax.set_title(title, fontsize=16)
1702 ax.set_xlabel(f"{cube.name()} / {cube.units}", fontsize=14)
1703 ax.set_ylabel("Power", fontsize=14)
1704 # Loop over all slices along the stamp_coordinate
1705 for member in cube.slices_over(stamp_coordinate):
1706 frequency = member.coord("frequency").points
1707 ax.plot(
1708 frequency,
1709 member.data,
1710 label=f"Member #{member.coord(stamp_coordinate).points[0]}",
1711 )
1713 # Add a legend
1714 ax.legend(fontsize=16)
1716 # Save the figure to a file
1717 plt.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution())
1718 logging.info("Saved power spectra plot to %s", filename)
1720 # Close the figure
1721 plt.close(fig)
1724def _spatial_plot(
1725 method: Literal["contourf", "pcolormesh"],
1726 cube: iris.cube.Cube,
1727 filename: str | None,
1728 sequence_coordinate: str,
1729 stamp_coordinate: str,
1730 overlay_cube: iris.cube.Cube | None = None,
1731 contour_cube: iris.cube.Cube | None = None,
1732 **kwargs,
1733):
1734 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
1736 A 2D spatial field can be plotted, but if the sequence_coordinate is present
1737 then a sequence of plots will be produced. Similarly if the stamp_coordinate
1738 is present then postage stamp plots will be produced.
1740 If an overlay_cube and/or contour_cube are specified, multiple variables can
1741 be overplotted on the same figure.
1743 Parameters
1744 ----------
1745 method: "contourf" | "pcolormesh"
1746 The plotting method to use.
1747 cube: Cube
1748 Iris cube of the data to plot. It should have two spatial dimensions,
1749 such as lat and lon, and may also have a another two dimension to be
1750 plotted sequentially and/or as postage stamp plots.
1751 filename: str | None
1752 Name of the plot to write, used as a prefix for plot sequences. If None
1753 uses the recipe name.
1754 sequence_coordinate: str
1755 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
1756 This coordinate must exist in the cube.
1757 stamp_coordinate: str
1758 Coordinate about which to plot postage stamp plots. Defaults to
1759 ``"realization"``.
1760 overlay_cube: Cube | None, optional
1761 Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube
1762 contour_cube: Cube | None, optional
1763 Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube
1765 Raises
1766 ------
1767 ValueError
1768 If the cube doesn't have the right dimensions.
1769 TypeError
1770 If the cube isn't a single cube.
1771 """
1772 recipe_title = get_recipe_metadata().get("title", "Untitled")
1774 # Ensure we've got a single cube.
1775 cube = _check_single_cube(cube)
1777 # Make postage stamp plots if stamp_coordinate exists and has more than a
1778 # single point.
1779 plotting_func = _plot_and_save_spatial_plot
1780 try:
1781 if cube.coord(stamp_coordinate).shape[0] > 1:
1782 plotting_func = _plot_and_save_postage_stamp_spatial_plot
1783 except iris.exceptions.CoordinateNotFoundError:
1784 pass
1786 # Produce a geographical scatter plot if the data have a
1787 # dimension called observation or model_obs_error
1788 if any( 1788 ↛ 1792line 1788 didn't jump to line 1792 because the condition on line 1788 was never true
1789 crd.var_name == "station" or crd.var_name == "model_obs_error"
1790 for crd in cube.coords()
1791 ):
1792 plotting_func = _plot_and_save_scattermap_plot
1794 # Must have a sequence coordinate.
1795 try:
1796 cube.coord(sequence_coordinate)
1797 except iris.exceptions.CoordinateNotFoundError as err:
1798 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err
1800 # Create a plot for each value of the sequence coordinate.
1801 plot_index = []
1802 nplot = np.size(cube.coord(sequence_coordinate).points)
1804 for iseq, cube_slice in enumerate(cube.slices_over(sequence_coordinate)):
1805 # Set plot titles and filename
1806 seq_coord = cube_slice.coord(sequence_coordinate)
1807 plot_title, plot_filename = _set_title_and_filename(
1808 seq_coord, nplot, recipe_title, filename
1809 )
1811 # Extract sequence slice for overlay_cube and contour_cube if required.
1812 overlay_slice = slice_over_maybe(overlay_cube, sequence_coordinate, iseq)
1813 contour_slice = slice_over_maybe(contour_cube, sequence_coordinate, iseq)
1815 # Do the actual plotting.
1816 plotting_func(
1817 cube_slice,
1818 filename=plot_filename,
1819 stamp_coordinate=stamp_coordinate,
1820 title=plot_title,
1821 method=method,
1822 overlay_cube=overlay_slice,
1823 contour_cube=contour_slice,
1824 **kwargs,
1825 )
1826 plot_index.append(plot_filename)
1828 # Add list of plots to plot metadata.
1829 complete_plot_index = _append_to_plot_index(plot_index)
1831 # Make a page to display the plots.
1832 _make_plot_html_page(complete_plot_index)
1835def _custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
1836 """Get colourmap for mask.
1838 If "mask_for_" appears anywhere in the name of a cube this function will be called
1839 regardless of the name of the variable to ensure a consistent plot.
1841 Parameters
1842 ----------
1843 cube: Cube
1844 Cube of variable for which the colorbar information is desired.
1845 axis: "x", "y", optional
1846 Select the levels for just this axis of a line plot. The min and max
1847 can be set by xmin/xmax or ymin/ymax respectively. For variables where
1848 setting a universal range is not desirable (e.g. temperature), users
1849 can set ymin/ymax values to "auto" in the colorbar definitions file.
1850 Where no additional xmin/xmax or ymin/ymax values are provided, the
1851 axis bounds default to use the vmin/vmax values provided.
1853 Returns
1854 -------
1855 cmap:
1856 Matplotlib colormap.
1857 levels:
1858 List of levels to use for plotting. For continuous plots the min and max
1859 should be taken as the range.
1860 norm:
1861 BoundaryNorm information.
1862 """
1863 if "difference" not in cube.long_name:
1864 if axis:
1865 levels = [0, 1]
1866 # Complete settings based on levels.
1867 return None, levels, None
1868 else:
1869 # Define the levels and colors.
1870 levels = [0, 1, 2]
1871 colors = ["white", "dodgerblue"]
1872 # Create a custom color map.
1873 cmap = mcolors.ListedColormap(colors)
1874 # Normalize the levels.
1875 norm = mcolors.BoundaryNorm(levels, cmap.N)
1876 logging.debug("Colourmap for %s.", cube.long_name)
1877 return cmap, levels, norm
1878 else:
1879 if axis:
1880 levels = [-1, 1]
1881 return None, levels, None
1882 else:
1883 # Search for if mask difference, set to +/- 0.5 as values plotted <
1884 # not <=.
1885 levels = [-2, -0.5, 0.5, 2]
1886 colors = ["goldenrod", "white", "teal"]
1887 cmap = mcolors.ListedColormap(colors)
1888 norm = mcolors.BoundaryNorm(levels, cmap.N)
1889 logging.debug("Colourmap for %s.", cube.long_name)
1890 return cmap, levels, norm
1893def _custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None):
1894 """Get a custom colorbar for a cube in the Beaufort Scale.
1896 Specific variable ranges can be separately set in user-supplied definition
1897 for x- or y-axis limits, or indicate where automated range preferred.
1899 Parameters
1900 ----------
1901 cube: Cube
1902 Cube of variable with Beaufort Scale in name.
1903 axis: "x", "y", optional
1904 Select the levels for just this axis of a line plot. The min and max
1905 can be set by xmin/xmax or ymin/ymax respectively. For variables where
1906 setting a universal range is not desirable (e.g. temperature), users
1907 can set ymin/ymax values to "auto" in the colorbar definitions file.
1908 Where no additional xmin/xmax or ymin/ymax values are provided, the
1909 axis bounds default to use the vmin/vmax values provided.
1911 Returns
1912 -------
1913 cmap:
1914 Matplotlib colormap.
1915 levels:
1916 List of levels to use for plotting. For continuous plots the min and max
1917 should be taken as the range.
1918 norm:
1919 BoundaryNorm information.
1920 """
1921 if "difference" not in cube.long_name:
1922 if axis:
1923 levels = [0, 12]
1924 return None, levels, None
1925 else:
1926 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
1927 colors = [
1928 "black",
1929 (0, 0, 0.6),
1930 "blue",
1931 "cyan",
1932 "green",
1933 "yellow",
1934 (1, 0.5, 0),
1935 "red",
1936 "pink",
1937 "magenta",
1938 "purple",
1939 "maroon",
1940 "white",
1941 ]
1942 cmap = mcolors.ListedColormap(colors)
1943 norm = mcolors.BoundaryNorm(levels, cmap.N)
1944 logging.info("change colormap for Beaufort Scale colorbar.")
1945 return cmap, levels, norm
1946 else:
1947 if axis:
1948 levels = [-4, 4]
1949 return None, levels, None
1950 else:
1951 levels = [
1952 -3.5,
1953 -2.5,
1954 -1.5,
1955 -0.5,
1956 0.5,
1957 1.5,
1958 2.5,
1959 3.5,
1960 ]
1961 cmap = plt.get_cmap("bwr", 8)
1962 norm = mcolors.BoundaryNorm(levels, cmap.N)
1963 return cmap, levels, norm
1966def _custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm):
1967 """Return altered colourmap for temperature with change in units to Celsius.
1969 If "Celsius" appears anywhere in the name of a cube this function will be called.
1971 Parameters
1972 ----------
1973 cube: Cube
1974 Cube of variable for which the colorbar information is desired.
1975 cmap: Matplotlib colormap.
1976 levels: List
1977 List of levels to use for plotting. For continuous plots the min and max
1978 should be taken as the range.
1979 norm: BoundaryNorm.
1981 Returns
1982 -------
1983 cmap: Matplotlib colormap.
1984 levels: List
1985 List of levels to use for plotting. For continuous plots the min and max
1986 should be taken as the range.
1987 norm: BoundaryNorm.
1988 """
1989 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
1990 if any("temperature" in name for name in varnames) and "Celsius" == cube.units:
1991 levels = np.array(levels)
1992 levels -= 273
1993 levels = levels.tolist()
1994 else:
1995 # Do nothing keep the existing colourbar attributes
1996 levels = levels
1997 cmap = cmap
1998 norm = norm
1999 return cmap, levels, norm
2002def _custom_colormap_probability(
2003 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None
2004):
2005 """Get a custom colorbar for a probability cube.
2007 Specific variable ranges can be separately set in user-supplied definition
2008 for x- or y-axis limits, or indicate where automated range preferred.
2010 Parameters
2011 ----------
2012 cube: Cube
2013 Cube of variable with probability in name.
2014 axis: "x", "y", optional
2015 Select the levels for just this axis of a line plot. The min and max
2016 can be set by xmin/xmax or ymin/ymax respectively. For variables where
2017 setting a universal range is not desirable (e.g. temperature), users
2018 can set ymin/ymax values to "auto" in the colorbar definitions file.
2019 Where no additional xmin/xmax or ymin/ymax values are provided, the
2020 axis bounds default to use the vmin/vmax values provided.
2022 Returns
2023 -------
2024 cmap:
2025 Matplotlib colormap.
2026 levels:
2027 List of levels to use for plotting. For continuous plots the min and max
2028 should be taken as the range.
2029 norm:
2030 BoundaryNorm information.
2031 """
2032 if axis:
2033 levels = [0, 1]
2034 return None, levels, None
2035 else:
2036 cmap = mcolors.ListedColormap(
2037 [
2038 "#FFFFFF",
2039 "#636363",
2040 "#e1dada",
2041 "#B5CAFF",
2042 "#8FB3FF",
2043 "#7F97FF",
2044 "#ABCF63",
2045 "#E8F59E",
2046 "#FFFA14",
2047 "#FFD121",
2048 "#FFA30A",
2049 ]
2050 )
2051 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]
2052 norm = mcolors.BoundaryNorm(levels, cmap.N)
2053 return cmap, levels, norm
2056def _custom_colourmap_precipitation(cube: iris.cube.Cube, cmap, levels, norm):
2057 """Return a custom colourmap for the current recipe."""
2058 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
2059 if (
2060 any("surface_microphysical" in name for name in varnames)
2061 and "difference" not in cube.long_name
2062 and "mask" not in cube.long_name
2063 ):
2064 # Define the levels and colors
2065 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256]
2066 colors = [
2067 "w",
2068 (0, 0, 0.6),
2069 "b",
2070 "c",
2071 "g",
2072 "y",
2073 (1, 0.5, 0),
2074 "r",
2075 "pink",
2076 "m",
2077 "purple",
2078 "maroon",
2079 "gray",
2080 ]
2081 # Create a custom colormap
2082 cmap = mcolors.ListedColormap(colors)
2083 # Normalize the levels
2084 norm = mcolors.BoundaryNorm(levels, cmap.N)
2085 logging.info("change colormap for surface_microphysical variable colorbar.")
2086 else:
2087 # do nothing and keep existing colorbar attributes
2088 cmap = cmap
2089 levels = levels
2090 norm = norm
2091 return cmap, levels, norm
2094def _custom_colormap_aviation_colour_state(cube: iris.cube.Cube):
2095 """Return custom colourmap for aviation colour state.
2097 If "aviation_colour_state" appears anywhere in the name of a cube
2098 this function will be called.
2100 Parameters
2101 ----------
2102 cube: Cube
2103 Cube of variable for which the colorbar information is desired.
2105 Returns
2106 -------
2107 cmap: Matplotlib colormap.
2108 levels: List
2109 List of levels to use for plotting. For continuous plots the min and max
2110 should be taken as the range.
2111 norm: BoundaryNorm.
2112 """
2113 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5]
2114 colors = [
2115 "#87ceeb",
2116 "#ffffff",
2117 "#8ced69",
2118 "#ffff00",
2119 "#ffd700",
2120 "#ffa500",
2121 "#fe3620",
2122 ]
2123 # Create a custom colormap
2124 cmap = mcolors.ListedColormap(colors)
2125 # Normalise the levels
2126 norm = mcolors.BoundaryNorm(levels, cmap.N)
2127 return cmap, levels, norm
2130def _custom_colourmap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm):
2131 """Return a custom colourmap for the current recipe."""
2132 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name])
2133 if (
2134 any("visibility_in_air" in name for name in varnames)
2135 and "difference" not in cube.long_name
2136 and "mask" not in cube.long_name
2137 ):
2138 # Define the levels and colors (in km)
2139 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]
2140 norm = mcolors.BoundaryNorm(levels, cmap.N)
2141 colours = [
2142 "#8f00d6",
2143 "#d10000",
2144 "#ff9700",
2145 "#ffff00",
2146 "#00007f",
2147 "#6c9ccd",
2148 "#aae8ff",
2149 "#37a648",
2150 "#8edc64",
2151 "#c5ffc5",
2152 "#dcdcdc",
2153 "#ffffff",
2154 ]
2155 # Create a custom colormap
2156 cmap = mcolors.ListedColormap(colours)
2157 # Normalize the levels
2158 norm = mcolors.BoundaryNorm(levels, cmap.N)
2159 logging.info("change colormap for visibility_in_air variable colorbar.")
2160 else:
2161 # do nothing and keep existing colorbar attributes
2162 cmap = cmap
2163 levels = levels
2164 norm = norm
2165 return cmap, levels, norm
2168def _get_num_models(cube: iris.cube.Cube | iris.cube.CubeList) -> int:
2169 """Return number of models based on cube attributes."""
2170 model_names = list(
2171 filter(
2172 lambda x: x is not None,
2173 {cb.attributes.get("model_name", None) for cb in iter_maybe(cube)},
2174 )
2175 )
2176 if not model_names:
2177 logging.debug("Missing model names. Will assume single model.")
2178 return 1
2179 else:
2180 return len(model_names)
2183def _validate_cube_shape(
2184 cube: iris.cube.Cube | iris.cube.CubeList, num_models: int
2185) -> None:
2186 """Check all cubes have a model name."""
2187 if isinstance(cube, iris.cube.CubeList) and len(cube) != num_models: 2187 ↛ 2188line 2187 didn't jump to line 2188 because the condition on line 2187 was never true
2188 raise ValueError(
2189 f"The number of model names ({num_models}) should equal the number "
2190 f"of cubes ({len(cube)})."
2191 )
2194def _validate_cubes_coords(
2195 cubes: iris.cube.CubeList, coords: list[iris.coords.Coord]
2196) -> None:
2197 """Check same number of cubes as sequence coordinate for zip functions."""
2198 if len(cubes) != len(coords): 2198 ↛ 2199line 2198 didn't jump to line 2199 because the condition on line 2198 was never true
2199 raise ValueError(
2200 f"The number of CubeList entries ({len(cubes)}) should equal the number "
2201 f"of sequence coordinates ({len(coords)})."
2202 f"Check that number of time entries in input data are consistent if "
2203 f"performing time-averaging steps prior to plotting outputs."
2204 )
2207####################
2208# Public functions #
2209####################
2212def spatial_contour_plot(
2213 cube: iris.cube.Cube,
2214 filename: str = None,
2215 sequence_coordinate: str = "time",
2216 stamp_coordinate: str = "realization",
2217 **kwargs,
2218) -> iris.cube.Cube:
2219 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
2221 A 2D spatial field can be plotted, but if the sequence_coordinate is present
2222 then a sequence of plots will be produced. Similarly if the stamp_coordinate
2223 is present then postage stamp plots will be produced.
2225 Parameters
2226 ----------
2227 cube: Cube
2228 Iris cube of the data to plot. It should have two spatial dimensions,
2229 such as lat and lon, and may also have a another two dimension to be
2230 plotted sequentially and/or as postage stamp plots.
2231 filename: str, optional
2232 Name of the plot to write, used as a prefix for plot sequences. Defaults
2233 to the recipe name.
2234 sequence_coordinate: str, optional
2235 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2236 This coordinate must exist in the cube.
2237 stamp_coordinate: str, optional
2238 Coordinate about which to plot postage stamp plots. Defaults to
2239 ``"realization"``.
2241 Returns
2242 -------
2243 Cube
2244 The original cube (so further operations can be applied).
2246 Raises
2247 ------
2248 ValueError
2249 If the cube doesn't have the right dimensions.
2250 TypeError
2251 If the cube isn't a single cube.
2252 """
2253 _spatial_plot(
2254 "contourf", cube, filename, sequence_coordinate, stamp_coordinate, **kwargs
2255 )
2256 return cube
2259def spatial_pcolormesh_plot(
2260 cube: iris.cube.Cube,
2261 filename: str = None,
2262 sequence_coordinate: str = "time",
2263 stamp_coordinate: str = "realization",
2264 **kwargs,
2265) -> iris.cube.Cube:
2266 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube.
2268 A 2D spatial field can be plotted, but if the sequence_coordinate is present
2269 then a sequence of plots will be produced. Similarly if the stamp_coordinate
2270 is present then postage stamp plots will be produced.
2272 This function is significantly faster than ``spatial_contour_plot``,
2273 especially at high resolutions, and should be preferred unless contiguous
2274 contour areas are important.
2276 Parameters
2277 ----------
2278 cube: Cube
2279 Iris cube of the data to plot. It should have two spatial dimensions,
2280 such as lat and lon, and may also have a another two dimension to be
2281 plotted sequentially and/or as postage stamp plots.
2282 filename: str, optional
2283 Name of the plot to write, used as a prefix for plot sequences. Defaults
2284 to the recipe name.
2285 sequence_coordinate: str, optional
2286 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2287 This coordinate must exist in the cube.
2288 stamp_coordinate: str, optional
2289 Coordinate about which to plot postage stamp plots. Defaults to
2290 ``"realization"``.
2292 Returns
2293 -------
2294 Cube
2295 The original cube (so further operations can be applied).
2297 Raises
2298 ------
2299 ValueError
2300 If the cube doesn't have the right dimensions.
2301 TypeError
2302 If the cube isn't a single cube.
2303 """
2304 _spatial_plot(
2305 "pcolormesh", cube, filename, sequence_coordinate, stamp_coordinate, **kwargs
2306 )
2307 return cube
2310def spatial_multi_pcolormesh_plot(
2311 cube: iris.cube.Cube,
2312 overlay_cube: iris.cube.Cube,
2313 contour_cube: iris.cube.Cube,
2314 filename: str = None,
2315 sequence_coordinate: str = "time",
2316 stamp_coordinate: str = "realization",
2317 **kwargs,
2318) -> iris.cube.Cube:
2319 """Plot a set of spatial variables onto a map from a 2D, 3D, or 4D cube.
2321 A 2D basis cube spatial field can be plotted, but if the sequence_coordinate is present
2322 then a sequence of plots will be produced. Similarly if the stamp_coordinate
2323 is present then postage stamp plots will be produced.
2325 If specified, a masked overlay_cube can be overplotted on top of the base cube.
2327 If specified, contours of a contour_cube can be overplotted on top of those.
2329 For single-variable equivalent of this routine, use spatial_pcolormesh_plot.
2331 This function is significantly faster than ``spatial_contour_plot``,
2332 especially at high resolutions, and should be preferred unless contiguous
2333 contour areas are important.
2335 Parameters
2336 ----------
2337 cube: Cube
2338 Iris cube of the data to plot. It should have two spatial dimensions,
2339 such as lat and lon, and may also have a another two dimension to be
2340 plotted sequentially and/or as postage stamp plots.
2341 overlay_cube: Cube
2342 Iris cube of the data to plot as an overlay on top of basis cube. It should have two spatial dimensions,
2343 such as lat and lon, and may also have a another two dimension to be
2344 plotted sequentially and/or as postage stamp plots. This is likely to be a masked cube in order not to hide the underlying basis cube.
2345 contour_cube: Cube
2346 Iris cube of the data to plot as a contour overlay on top of basis cube and overlay_cube. It should have two spatial dimensions,
2347 such as lat and lon, and may also have a another two dimension to be
2348 plotted sequentially and/or as postage stamp plots.
2349 filename: str, optional
2350 Name of the plot to write, used as a prefix for plot sequences. Defaults
2351 to the recipe name.
2352 sequence_coordinate: str, optional
2353 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2354 This coordinate must exist in the cube.
2355 stamp_coordinate: str, optional
2356 Coordinate about which to plot postage stamp plots. Defaults to
2357 ``"realization"``.
2359 Returns
2360 -------
2361 Cube
2362 The original cube (so further operations can be applied).
2364 Raises
2365 ------
2366 ValueError
2367 If the cube doesn't have the right dimensions.
2368 TypeError
2369 If the cube isn't a single cube.
2370 """
2371 _spatial_plot(
2372 "pcolormesh",
2373 cube,
2374 filename,
2375 sequence_coordinate,
2376 stamp_coordinate,
2377 overlay_cube=overlay_cube,
2378 contour_cube=contour_cube,
2379 )
2380 return cube, overlay_cube, contour_cube
2383# TODO: Expand function to handle ensemble data.
2384# line_coordinate: str, optional
2385# Coordinate about which to plot multiple lines. Defaults to
2386# ``"realization"``.
2387def plot_line_series(
2388 cube: iris.cube.Cube | iris.cube.CubeList,
2389 filename: str = None,
2390 series_coordinate: str = "time",
2391 # line_coordinate: str = "realization",
2392 **kwargs,
2393) -> iris.cube.Cube | iris.cube.CubeList:
2394 """Plot a line plot for the specified coordinate.
2396 The Cube or CubeList must be 1D.
2398 Parameters
2399 ----------
2400 iris.cube | iris.cube.CubeList
2401 Cube or CubeList of the data to plot. The individual cubes should have a single dimension.
2402 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
2403 We do not support different data such as temperature and humidity in the same CubeList for plotting.
2404 filename: str, optional
2405 Name of the plot to write, used as a prefix for plot sequences. Defaults
2406 to the recipe name.
2407 series_coordinate: str, optional
2408 Coordinate about which to make a series. Defaults to ``"time"``. This
2409 coordinate must exist in the cube.
2411 Returns
2412 -------
2413 iris.cube.Cube | iris.cube.CubeList
2414 The original Cube or CubeList (so further operations can be applied).
2415 plotted data.
2417 Raises
2418 ------
2419 ValueError
2420 If the cubes don't have the right dimensions.
2421 TypeError
2422 If the cube isn't a Cube or CubeList.
2423 """
2424 # Ensure we have a name for the plot file.
2425 recipe_title = get_recipe_metadata().get("title", "Untitled")
2427 num_models = _get_num_models(cube)
2429 _validate_cube_shape(cube, num_models)
2431 # Iterate over all cubes and extract coordinate to plot.
2432 cubes = iter_maybe(cube)
2433 coords = []
2434 for cube in cubes:
2435 try:
2436 coords.append(cube.coord(series_coordinate))
2437 except iris.exceptions.CoordinateNotFoundError as err:
2438 raise ValueError(
2439 f"Cube must have a {series_coordinate} coordinate."
2440 ) from err
2441 if cube.ndim > 2 or not cube.coords("realization"):
2442 raise ValueError("Cube must be 1D or 2D with a realization coordinate.")
2444 # Format the title and filename using plotted series coordinate
2445 nplot = 1
2446 seq_coord = coords[0]
2447 plot_title, plot_filename = _set_title_and_filename(
2448 seq_coord, nplot, recipe_title, filename
2449 )
2451 # Do the actual plotting.
2452 _plot_and_save_line_series(cubes, coords, "realization", plot_filename, plot_title)
2454 # Add list of plots to plot metadata.
2455 plot_index = _append_to_plot_index([plot_filename])
2457 # Make a page to display the plots.
2458 _make_plot_html_page(plot_index)
2460 return cube
2463def plot_vertical_line_series(
2464 cubes: iris.cube.Cube | iris.cube.CubeList,
2465 filename: str = None,
2466 series_coordinate: str = "model_level_number",
2467 sequence_coordinate: str = "time",
2468 # line_coordinate: str = "realization",
2469 **kwargs,
2470) -> iris.cube.Cube | iris.cube.CubeList:
2471 """Plot a line plot against a type of vertical coordinate.
2473 The Cube or CubeList must be 1D.
2475 A 1D line plot with y-axis as pressure coordinate can be plotted, but if the sequence_coordinate is present
2476 then a sequence of plots will be produced.
2478 Parameters
2479 ----------
2480 iris.cube | iris.cube.CubeList
2481 Cube or CubeList of the data to plot. The individual cubes should have a single dimension.
2482 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
2483 We do not support different data such as temperature and humidity in the same CubeList for plotting.
2484 filename: str, optional
2485 Name of the plot to write, used as a prefix for plot sequences. Defaults
2486 to the recipe name.
2487 series_coordinate: str, optional
2488 Coordinate to plot on the y-axis. Can be ``pressure`` or
2489 ``model_level_number`` for UM, or ``full_levels`` or ``half_levels``
2490 for LFRic. Defaults to ``model_level_number``.
2491 This coordinate must exist in the cube.
2492 sequence_coordinate: str, optional
2493 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2494 This coordinate must exist in the cube.
2496 Returns
2497 -------
2498 iris.cube.Cube | iris.cube.CubeList
2499 The original Cube or CubeList (so further operations can be applied).
2500 Plotted data.
2502 Raises
2503 ------
2504 ValueError
2505 If the cubes doesn't have the right dimensions.
2506 TypeError
2507 If the cube isn't a Cube or CubeList.
2508 """
2509 # Ensure we have a name for the plot file.
2510 recipe_title = get_recipe_metadata().get("title", "Untitled")
2512 cubes = iter_maybe(cubes)
2513 # Initialise empty list to hold all data from all cubes in a CubeList
2514 all_data = []
2516 # Store min/max ranges for x range.
2517 x_levels = []
2519 num_models = _get_num_models(cubes)
2521 _validate_cube_shape(cubes, num_models)
2523 # Iterate over all cubes in cube or CubeList and plot.
2524 coords = []
2525 for cube in cubes:
2526 # Test if series coordinate i.e. pressure level exist for any cube with cube.ndim >=1.
2527 try:
2528 coords.append(cube.coord(series_coordinate))
2529 except iris.exceptions.CoordinateNotFoundError as err:
2530 raise ValueError(
2531 f"Cube must have a {series_coordinate} coordinate."
2532 ) from err
2534 try:
2535 if cube.ndim > 1 or not cube.coords("realization"): 2535 ↛ 2543line 2535 didn't jump to line 2543 because the condition on line 2535 was always true
2536 cube.coord(sequence_coordinate)
2537 except iris.exceptions.CoordinateNotFoundError as err:
2538 raise ValueError(
2539 f"Cube must have a {sequence_coordinate} coordinate or be 1D, or 2D with a realization coordinate."
2540 ) from err
2542 # Get minimum and maximum from levels information.
2543 _, levels, _ = _colorbar_map_levels(cube, axis="x")
2544 if levels is not None: 2544 ↛ 2548line 2544 didn't jump to line 2548 because the condition on line 2544 was always true
2545 x_levels.append(min(levels))
2546 x_levels.append(max(levels))
2547 else:
2548 all_data.append(cube.data)
2550 if len(x_levels) == 0: 2550 ↛ 2552line 2550 didn't jump to line 2552 because the condition on line 2550 was never true
2551 # Combine all data into a single NumPy array
2552 combined_data = np.concatenate(all_data)
2554 # Set the lower and upper limit for the x-axis to ensure all plots have
2555 # same range. This needs to read the whole cube over the range of the
2556 # sequence and if applicable postage stamp coordinate.
2557 vmin = np.floor(combined_data.min())
2558 vmax = np.ceil(combined_data.max())
2559 else:
2560 vmin = min(x_levels)
2561 vmax = max(x_levels)
2563 # Matching the slices (matching by seq coord point; it may happen that
2564 # evaluated models do not cover the same seq coord range, hence matching
2565 # necessary)
2566 def filter_cube_iterables(cube_iterables) -> bool:
2567 return len(cube_iterables) == len(coords)
2569 cube_iterables = filter(
2570 filter_cube_iterables,
2571 (
2572 iris.cube.CubeList(
2573 s
2574 for s in itertools.chain.from_iterable(
2575 cb.slices_over(sequence_coordinate) for cb in cubes
2576 )
2577 if s.coord(sequence_coordinate).points[0] == point
2578 )
2579 for point in sorted(
2580 set(
2581 itertools.chain.from_iterable(
2582 cb.coord(sequence_coordinate).points for cb in cubes
2583 )
2584 )
2585 )
2586 ),
2587 )
2589 # Create a plot for each value of the sequence coordinate.
2590 # Allowing for multiple cubes in a CubeList to be plotted in the same plot for
2591 # similar sequence values. Passing a CubeList into the internal plotting function
2592 # for similar values of the sequence coordinate. cube_slice can be an iris.cube.Cube
2593 # or an iris.cube.CubeList.
2594 plot_index = []
2595 nplot = np.size(cubes[0].coord(sequence_coordinate).points)
2596 for cubes_slice in cube_iterables:
2597 # Format the coordinate value in a unit appropriate way.
2598 seq_coord = cubes_slice[0].coord(sequence_coordinate)
2599 plot_title, plot_filename = _set_title_and_filename(
2600 seq_coord, nplot, recipe_title, filename
2601 )
2603 # Do the actual plotting.
2604 _plot_and_save_vertical_line_series(
2605 cubes_slice,
2606 coords,
2607 "realization",
2608 plot_filename,
2609 series_coordinate,
2610 title=plot_title,
2611 vmin=vmin,
2612 vmax=vmax,
2613 )
2614 plot_index.append(plot_filename)
2616 # Add list of plots to plot metadata.
2617 complete_plot_index = _append_to_plot_index(plot_index)
2619 # Make a page to display the plots.
2620 _make_plot_html_page(complete_plot_index)
2622 return cubes
2625def qq_plot(
2626 cubes: iris.cube.CubeList,
2627 coordinates: list[str],
2628 percentiles: list[float],
2629 model_names: list[str],
2630 filename: str = None,
2631 one_to_one: bool = True,
2632 **kwargs,
2633) -> iris.cube.CubeList:
2634 """Plot a Quantile-Quantile plot between two models for common time points.
2636 The cubes will be normalised by collapsing each cube to its percentiles. Cubes are
2637 collapsed within the operator over all specified coordinates such as
2638 grid_latitude, grid_longitude, vertical levels, but also realisation representing
2639 ensemble members to ensure a 1D cube (array).
2641 Parameters
2642 ----------
2643 cubes: iris.cube.CubeList
2644 Two cubes of the same variable with different models.
2645 coordinate: list[str]
2646 The list of coordinates to collapse over. This list should be
2647 every coordinate within the cube to result in a 1D cube around
2648 the percentile coordinate.
2649 percent: list[float]
2650 A list of percentiles to appear in the plot.
2651 model_names: list[str]
2652 A list of model names to appear on the axis of the plot.
2653 filename: str, optional
2654 Filename of the plot to write.
2655 one_to_one: bool, optional
2656 If True a 1:1 line is plotted; if False it is not. Default is True.
2658 Raises
2659 ------
2660 ValueError
2661 When the cubes are not compatible.
2663 Notes
2664 -----
2665 The quantile-quantile plot is a variant on the scatter plot representing
2666 two datasets by their quantiles (percentiles) for common time points.
2667 This plot does not use a theoretical distribution to compare against, but
2668 compares percentiles of two datasets. This plot does
2669 not use all raw data points, but plots the selected percentiles (quantiles) of
2670 each variable instead for the two datasets, thereby normalising the data for a
2671 direct comparison between the selected percentiles of the two dataset distributions.
2673 Quantile-quantile plots are valuable for comparing against
2674 observations and other models. Identical percentiles between the variables
2675 will lie on the one-to-one line implying the values correspond well to each
2676 other. Where there is a deviation from the one-to-one line a range of
2677 possibilities exist depending on how and where the data is shifted (e.g.,
2678 Wilks 2011 [Wilks2011]_).
2680 For distributions above the one-to-one line the distribution is left-skewed;
2681 below is right-skewed. A distinct break implies a bimodal distribution, and
2682 closer values/values further apart at the tails imply poor representation of
2683 the extremes.
2685 References
2686 ----------
2687 .. [Wilks2011] Wilks, D.S., (2011) "Statistical Methods in the Atmospheric
2688 Sciences" Third Edition, vol. 100, Academic Press, Oxford, UK, 676 pp.
2689 """
2690 # Check cubes using same functionality as the difference operator.
2691 if len(cubes) != 2:
2692 raise ValueError("cubes should contain exactly 2 cubes.")
2693 base: Cube = cubes.extract_cube(iris.AttributeConstraint(cset_comparison_base=1))
2694 other: Cube = cubes.extract_cube(
2695 iris.Constraint(
2696 cube_func=lambda cube: "cset_comparison_base" not in cube.attributes
2697 )
2698 )
2700 # Get spatial coord names.
2701 base_lat_name, base_lon_name = get_cube_yxcoordname(base)
2702 other_lat_name, other_lon_name = get_cube_yxcoordname(other)
2704 # Ensure cubes to compare are on common differencing grid.
2705 # This is triggered if either
2706 # i) latitude and longitude shapes are not the same. Note grid points
2707 # are not compared directly as these can differ through rounding
2708 # errors.
2709 # ii) or variables are known to often sit on different grid staggering
2710 # in different models (e.g. cell center vs cell edge), as is the case
2711 # for UM and LFRic comparisons.
2712 # In future greater choice of regridding method might be applied depending
2713 # on variable type. Linear regridding can in general be appropriate for smooth
2714 # variables. Care should be taken with interpretation of differences
2715 # given this dependency on regridding.
2716 if (
2717 base.coord(base_lat_name).shape != other.coord(other_lat_name).shape
2718 or base.coord(base_lon_name).shape != other.coord(other_lon_name).shape
2719 ) or (
2720 base.long_name
2721 in [
2722 "eastward_wind_at_10m",
2723 "northward_wind_at_10m",
2724 "northward_wind_at_cell_centres",
2725 "eastward_wind_at_cell_centres",
2726 "zonal_wind_at_pressure_levels",
2727 "meridional_wind_at_pressure_levels",
2728 "potential_vorticity_at_pressure_levels",
2729 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",
2730 ]
2731 ):
2732 logging.debug(
2733 "Linear regridding base cube to other grid to compute differences"
2734 )
2735 base = regrid_onto_cube(base, other, method="Linear")
2737 # Extract just common time points.
2738 base, other = _extract_common_time_points(base, other)
2740 # Equalise attributes so we can merge.
2741 fully_equalise_attributes([base, other])
2742 logging.debug("Base: %s\nOther: %s", base, other)
2744 # Collapse cubes.
2745 base = collapse(
2746 base,
2747 coordinate=coordinates,
2748 method="PERCENTILE",
2749 additional_percent=percentiles,
2750 )
2751 other = collapse(
2752 other,
2753 coordinate=coordinates,
2754 method="PERCENTILE",
2755 additional_percent=percentiles,
2756 )
2758 # Ensure we have a name for the plot file.
2759 recipe_title = get_recipe_metadata().get("title", "Untitled")
2760 title = f"{recipe_title}"
2762 if filename is None:
2763 filename = slugify(recipe_title)
2765 # Add file extension.
2766 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
2768 # Do the actual plotting on a scatter plot
2769 _plot_and_save_scatter_plot(
2770 base, other, plot_filename, title, one_to_one, model_names
2771 )
2773 # Add list of plots to plot metadata.
2774 plot_index = _append_to_plot_index([plot_filename])
2776 # Make a page to display the plots.
2777 _make_plot_html_page(plot_index)
2779 return iris.cube.CubeList([base, other])
2782def scatter_plot(
2783 cube_x: iris.cube.Cube | iris.cube.CubeList,
2784 cube_y: iris.cube.Cube | iris.cube.CubeList,
2785 filename: str = None,
2786 one_to_one: bool = True,
2787 **kwargs,
2788) -> iris.cube.CubeList:
2789 """Plot a scatter plot between two variables.
2791 Both cubes must be 1D.
2793 Parameters
2794 ----------
2795 cube_x: Cube | CubeList
2796 1 dimensional Cube of the data to plot on y-axis.
2797 cube_y: Cube | CubeList
2798 1 dimensional Cube of the data to plot on x-axis.
2799 filename: str, optional
2800 Filename of the plot to write.
2801 one_to_one: bool, optional
2802 If True a 1:1 line is plotted; if False it is not. Default is True.
2804 Returns
2805 -------
2806 cubes: CubeList
2807 CubeList of the original x and y cubes for further processing.
2809 Raises
2810 ------
2811 ValueError
2812 If the cube doesn't have the right dimensions and cubes not the same
2813 size.
2814 TypeError
2815 If the cube isn't a single cube.
2817 Notes
2818 -----
2819 Scatter plots are used for determining if there is a relationship between
2820 two variables. Positive relations have a slope going from bottom left to top
2821 right; Negative relations have a slope going from top left to bottom right.
2822 """
2823 # Iterate over all cubes in cube or CubeList and plot.
2824 for cube_iter in iter_maybe(cube_x):
2825 # Check cubes are correct shape.
2826 cube_iter = _check_single_cube(cube_iter)
2827 if cube_iter.ndim > 1:
2828 raise ValueError("cube_x must be 1D.")
2830 # Iterate over all cubes in cube or CubeList and plot.
2831 for cube_iter in iter_maybe(cube_y):
2832 # Check cubes are correct shape.
2833 cube_iter = _check_single_cube(cube_iter)
2834 if cube_iter.ndim > 1:
2835 raise ValueError("cube_y must be 1D.")
2837 # Ensure we have a name for the plot file.
2838 recipe_title = get_recipe_metadata().get("title", "Untitled")
2839 title = f"{recipe_title}"
2841 if filename is None:
2842 filename = slugify(recipe_title)
2844 # Add file extension.
2845 plot_filename = f"{filename.rsplit('.', 1)[0]}.png"
2847 # Do the actual plotting.
2848 _plot_and_save_scatter_plot(cube_x, cube_y, plot_filename, title, one_to_one)
2850 # Add list of plots to plot metadata.
2851 plot_index = _append_to_plot_index([plot_filename])
2853 # Make a page to display the plots.
2854 _make_plot_html_page(plot_index)
2856 return iris.cube.CubeList([cube_x, cube_y])
2859def vector_plot(
2860 cube_u: iris.cube.Cube,
2861 cube_v: iris.cube.Cube,
2862 filename: str = None,
2863 sequence_coordinate: str = "time",
2864 **kwargs,
2865) -> iris.cube.CubeList:
2866 """Plot a vector plot based on the input u and v components."""
2867 recipe_title = get_recipe_metadata().get("title", "Untitled")
2869 # Cubes must have a matching sequence coordinate.
2870 try:
2871 # Check that the u and v cubes have the same sequence coordinate.
2872 if cube_u.coord(sequence_coordinate) != cube_v.coord(sequence_coordinate): 2872 ↛ anywhereline 2872 didn't jump anywhere: it always raised an exception.
2873 raise ValueError("Coordinates do not match.")
2874 except (iris.exceptions.CoordinateNotFoundError, ValueError) as err:
2875 raise ValueError(
2876 f"Cubes should have matching {sequence_coordinate} coordinate:\n{cube_u}\n{cube_v}"
2877 ) from err
2879 # Create a plot for each value of the sequence coordinate.
2880 plot_index = []
2881 nplot = np.size(cube_u[0].coord(sequence_coordinate).points)
2882 for cube_u_slice, cube_v_slice in zip(
2883 cube_u.slices_over(sequence_coordinate),
2884 cube_v.slices_over(sequence_coordinate),
2885 strict=True,
2886 ):
2887 # Format the coordinate value in a unit appropriate way.
2888 seq_coord = cube_u_slice.coord(sequence_coordinate)
2889 plot_title, plot_filename = _set_title_and_filename(
2890 seq_coord, nplot, recipe_title, filename
2891 )
2893 # Do the actual plotting.
2894 _plot_and_save_vector_plot(
2895 cube_u_slice,
2896 cube_v_slice,
2897 filename=plot_filename,
2898 title=plot_title,
2899 method="contourf",
2900 )
2901 plot_index.append(plot_filename)
2903 # Add list of plots to plot metadata.
2904 complete_plot_index = _append_to_plot_index(plot_index)
2906 # Make a page to display the plots.
2907 _make_plot_html_page(complete_plot_index)
2909 return iris.cube.CubeList([cube_u, cube_v])
2912def plot_histogram_series(
2913 cubes: iris.cube.Cube | iris.cube.CubeList,
2914 filename: str = None,
2915 sequence_coordinate: str = "time",
2916 stamp_coordinate: str = "realization",
2917 single_plot: bool = False,
2918 **kwargs,
2919) -> iris.cube.Cube | iris.cube.CubeList:
2920 """Plot a histogram plot for each vertical level provided.
2922 A histogram plot can be plotted, but if the sequence_coordinate (i.e. time)
2923 is present then a sequence of plots will be produced using the time slider
2924 functionality to scroll through histograms against time. If a
2925 stamp_coordinate is present then postage stamp plots will be produced. If
2926 stamp_coordinate and single_plot is True, all postage stamp plots will be
2927 plotted in a single plot instead of separate postage stamp plots.
2929 Parameters
2930 ----------
2931 cubes: Cube | iris.cube.CubeList
2932 Iris cube or CubeList of the data to plot. It should have a single dimension other
2933 than the stamp coordinate.
2934 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
2935 We do not support different data such as temperature and humidity in the same CubeList for plotting.
2936 filename: str, optional
2937 Name of the plot to write, used as a prefix for plot sequences. Defaults
2938 to the recipe name.
2939 sequence_coordinate: str, optional
2940 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
2941 This coordinate must exist in the cube and will be used for the time
2942 slider.
2943 stamp_coordinate: str, optional
2944 Coordinate about which to plot postage stamp plots. Defaults to
2945 ``"realization"``.
2946 single_plot: bool, optional
2947 If True, all postage stamp plots will be plotted in a single plot. If
2948 False, each postage stamp plot will be plotted separately. Is only valid
2949 if stamp_coordinate exists and has more than a single point.
2951 Returns
2952 -------
2953 iris.cube.Cube | iris.cube.CubeList
2954 The original Cube or CubeList (so further operations can be applied).
2955 Plotted data.
2957 Raises
2958 ------
2959 ValueError
2960 If the cube doesn't have the right dimensions.
2961 TypeError
2962 If the cube isn't a Cube or CubeList.
2963 """
2964 recipe_title = get_recipe_metadata().get("title", "Untitled")
2966 cubes = iter_maybe(cubes)
2968 # Internal plotting function.
2969 plotting_func = _plot_and_save_histogram_series
2971 num_models = _get_num_models(cubes)
2973 _validate_cube_shape(cubes, num_models)
2975 # If several histograms are plotted with time as sequence_coordinate for the
2976 # time slider option.
2977 for cube in cubes:
2978 try:
2979 cube.coord(sequence_coordinate)
2980 except iris.exceptions.CoordinateNotFoundError as err:
2981 raise ValueError(
2982 f"Cube must have a {sequence_coordinate} coordinate."
2983 ) from err
2985 # Get minimum and maximum from levels information.
2986 levels = None
2987 for cube in cubes: 2987 ↛ 3003line 2987 didn't jump to line 3003 because the loop on line 2987 didn't complete
2988 # First check if user-specified "auto" range variable.
2989 # This maintains the value of levels as None, so proceed.
2990 _, levels, _ = _colorbar_map_levels(cube, axis="y")
2991 if levels is None:
2992 break
2993 # If levels is changed, recheck to use the vmin,vmax or
2994 # levels-based ranges for histogram plots.
2995 _, levels, _ = _colorbar_map_levels(cube)
2996 logging.debug("levels: %s", levels)
2997 if levels is not None: 2997 ↛ 2987line 2997 didn't jump to line 2987 because the condition on line 2997 was always true
2998 vmin = min(levels)
2999 vmax = max(levels)
3000 logging.debug("Updated vmin, vmax: %s, %s", vmin, vmax)
3001 break
3003 if levels is None:
3004 vmin = min(cb.data.min() for cb in cubes)
3005 vmax = max(cb.data.max() for cb in cubes)
3007 # Make postage stamp plots if stamp_coordinate exists and has more than a
3008 # single point. If single_plot is True:
3009 # -- all postage stamp plots will be plotted in a single plot instead of
3010 # separate postage stamp plots.
3011 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are
3012 # produced per single model only
3013 if num_models == 1: 3013 ↛ 3026line 3013 didn't jump to line 3026 because the condition on line 3013 was always true
3014 if ( 3014 ↛ 3018line 3014 didn't jump to line 3018 because the condition on line 3014 was never true
3015 stamp_coordinate in [c.name() for c in cubes[0].coords()]
3016 and cubes[0].coord(stamp_coordinate).shape[0] > 1
3017 ):
3018 if single_plot:
3019 plotting_func = (
3020 _plot_and_save_postage_stamps_in_single_plot_histogram_series
3021 )
3022 else:
3023 plotting_func = _plot_and_save_postage_stamp_histogram_series
3024 cube_iterables = cubes[0].slices_over(sequence_coordinate)
3025 else:
3026 all_points = sorted(
3027 set(
3028 itertools.chain.from_iterable(
3029 cb.coord(sequence_coordinate).points for cb in cubes
3030 )
3031 )
3032 )
3033 all_slices = list(
3034 itertools.chain.from_iterable(
3035 cb.slices_over(sequence_coordinate) for cb in cubes
3036 )
3037 )
3038 # Matched slices (matched by seq coord point; it may happen that
3039 # evaluated models do not cover the same seq coord range, hence matching
3040 # necessary)
3041 cube_iterables = [
3042 iris.cube.CubeList(
3043 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point
3044 )
3045 for point in all_points
3046 ]
3048 plot_index = []
3049 nplot = np.size(cube.coord(sequence_coordinate).points)
3050 # Create a plot for each value of the sequence coordinate. Allowing for
3051 # multiple cubes in a CubeList to be plotted in the same plot for similar
3052 # sequence values. Passing a CubeList into the internal plotting function
3053 # for similar values of the sequence coordinate. cube_slice can be an
3054 # iris.cube.Cube or an iris.cube.CubeList.
3055 for cube_slice in cube_iterables:
3056 single_cube = cube_slice
3057 if isinstance(cube_slice, iris.cube.CubeList): 3057 ↛ 3058line 3057 didn't jump to line 3058 because the condition on line 3057 was never true
3058 single_cube = cube_slice[0]
3060 # Set plot titles and filename, based on sequence coordinate
3061 seq_coord = single_cube.coord(sequence_coordinate)
3062 # Use time coordinate in title and filename if single histogram output.
3063 if sequence_coordinate == "realization" and nplot == 1: 3063 ↛ 3064line 3063 didn't jump to line 3064 because the condition on line 3063 was never true
3064 seq_coord = single_cube.coord("time")
3065 plot_title, plot_filename = _set_title_and_filename(
3066 seq_coord, nplot, recipe_title, filename
3067 )
3069 # Do the actual plotting.
3070 plotting_func(
3071 cube_slice,
3072 filename=plot_filename,
3073 stamp_coordinate=stamp_coordinate,
3074 title=plot_title,
3075 vmin=vmin,
3076 vmax=vmax,
3077 )
3078 plot_index.append(plot_filename)
3080 # Add list of plots to plot metadata.
3081 complete_plot_index = _append_to_plot_index(plot_index)
3083 # Make a page to display the plots.
3084 _make_plot_html_page(complete_plot_index)
3086 return cubes
3089def plot_power_spectrum_series(
3090 cubes: iris.cube.Cube | iris.cube.CubeList,
3091 filename: str = None,
3092 sequence_coordinate: str = "time",
3093 stamp_coordinate: str = "realization",
3094 single_plot: bool = False,
3095 **kwargs,
3096) -> iris.cube.Cube | iris.cube.CubeList:
3097 """Plot a power spectrum plot for each vertical level provided.
3099 A power spectrum plot can be plotted, but if the sequence_coordinate (i.e. time)
3100 is present then a sequence of plots will be produced using the time slider
3101 functionality to scroll through power spectrum against time. If a
3102 stamp_coordinate is present then postage stamp plots will be produced. If
3103 stamp_coordinate and single_plot is True, all postage stamp plots will be
3104 plotted in a single plot instead of separate postage stamp plots.
3106 Parameters
3107 ----------
3108 cubes: Cube | iris.cube.CubeList
3109 Iris cube or CubeList of the data to plot. It should have a single dimension other
3110 than the stamp coordinate.
3111 The cubes should cover the same phenomenon i.e. all cubes contain temperature data.
3112 We do not support different data such as temperature and humidity in the same CubeList for plotting.
3113 filename: str, optional
3114 Name of the plot to write, used as a prefix for plot sequences. Defaults
3115 to the recipe name.
3116 sequence_coordinate: str, optional
3117 Coordinate about which to make a plot sequence. Defaults to ``"time"``.
3118 This coordinate must exist in the cube and will be used for the time
3119 slider.
3120 stamp_coordinate: str, optional
3121 Coordinate about which to plot postage stamp plots. Defaults to
3122 ``"realization"``.
3123 single_plot: bool, optional
3124 If True, all postage stamp plots will be plotted in a single plot. If
3125 False, each postage stamp plot will be plotted separately. Is only valid
3126 if stamp_coordinate exists and has more than a single point.
3128 Returns
3129 -------
3130 iris.cube.Cube | iris.cube.CubeList
3131 The original Cube or CubeList (so further operations can be applied).
3132 Plotted data.
3134 Raises
3135 ------
3136 ValueError
3137 If the cube doesn't have the right dimensions.
3138 TypeError
3139 If the cube isn't a Cube or CubeList.
3140 """
3141 recipe_title = get_recipe_metadata().get("title", "Untitled")
3143 cubes = iter_maybe(cubes)
3145 # Internal plotting function.
3146 plotting_func = _plot_and_save_power_spectrum_series
3148 num_models = _get_num_models(cubes)
3150 _validate_cube_shape(cubes, num_models)
3152 # If several power spectra are plotted with time as sequence_coordinate for the
3153 # time slider option.
3154 for cube in cubes:
3155 try:
3156 cube.coord(sequence_coordinate)
3157 except iris.exceptions.CoordinateNotFoundError as err:
3158 raise ValueError(
3159 f"Cube must have a {sequence_coordinate} coordinate."
3160 ) from err
3162 # Make postage stamp plots if stamp_coordinate exists and has more than a
3163 # single point. If single_plot is True:
3164 # -- all postage stamp plots will be plotted in a single plot instead of
3165 # separate postage stamp plots.
3166 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are
3167 # produced per single model only
3168 if num_models == 1: 3168 ↛ 3181line 3168 didn't jump to line 3181 because the condition on line 3168 was always true
3169 if ( 3169 ↛ 3173line 3169 didn't jump to line 3173 because the condition on line 3169 was never true
3170 stamp_coordinate in [c.name() for c in cubes[0].coords()]
3171 and cubes[0].coord(stamp_coordinate).shape[0] > 1
3172 ):
3173 if single_plot:
3174 plotting_func = (
3175 _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series
3176 )
3177 else:
3178 plotting_func = _plot_and_save_postage_stamp_power_spectrum_series
3179 cube_iterables = cubes[0].slices_over(sequence_coordinate)
3180 else:
3181 all_points = sorted(
3182 set(
3183 itertools.chain.from_iterable(
3184 cb.coord(sequence_coordinate).points for cb in cubes
3185 )
3186 )
3187 )
3188 all_slices = list(
3189 itertools.chain.from_iterable(
3190 cb.slices_over(sequence_coordinate) for cb in cubes
3191 )
3192 )
3193 # Matched slices (matched by seq coord point; it may happen that
3194 # evaluated models do not cover the same seq coord range, hence matching
3195 # necessary)
3196 cube_iterables = [
3197 iris.cube.CubeList(
3198 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point
3199 )
3200 for point in all_points
3201 ]
3203 plot_index = []
3204 nplot = np.size(cube.coord(sequence_coordinate).points)
3205 # Create a plot for each value of the sequence coordinate. Allowing for
3206 # multiple cubes in a CubeList to be plotted in the same plot for similar
3207 # sequence values. Passing a CubeList into the internal plotting function
3208 # for similar values of the sequence coordinate. cube_slice can be an
3209 # iris.cube.Cube or an iris.cube.CubeList.
3210 for cube_slice in cube_iterables:
3211 single_cube = cube_slice
3212 if isinstance(cube_slice, iris.cube.CubeList): 3212 ↛ 3213line 3212 didn't jump to line 3213 because the condition on line 3212 was never true
3213 single_cube = cube_slice[0]
3215 # Set plot title and filenames based on sequence values
3216 seq_coord = single_cube.coord(sequence_coordinate)
3217 plot_title, plot_filename = _set_title_and_filename(
3218 seq_coord, nplot, recipe_title, filename
3219 )
3221 # Do the actual plotting.
3222 plotting_func(
3223 cube_slice,
3224 filename=plot_filename,
3225 stamp_coordinate=stamp_coordinate,
3226 title=plot_title,
3227 )
3228 plot_index.append(plot_filename)
3230 # Add list of plots to plot metadata.
3231 complete_plot_index = _append_to_plot_index(plot_index)
3233 # Make a page to display the plots.
3234 _make_plot_html_page(complete_plot_index)
3236 return cubes
3239def _DCT_ps(y_3d):
3240 """Calculate power spectra for regional domains.
3242 Parameters
3243 ----------
3244 y_3d: 3D array
3245 3 dimensional array to calculate spectrum for.
3246 (2D field data with 3rd dimension of time)
3248 Returns: ps_array
3249 Array of power spectra values calculated for input field (for each time)
3251 Method for regional domains:
3252 Calculate power spectra over limited area domain using Discrete Cosine Transform (DCT)
3253 as described in Denis et al 2002 [Denis_etal_2002]_.
3255 References
3256 ----------
3257 .. [Denis_etal_2002] Bertrand Denis, Jean Côté and René Laprise (2002)
3258 "Spectral Decomposition of Two-Dimensional Atmospheric Fields on
3259 Limited-Area Domains Using the Discrete Cosine Transform (DCT)"
3260 Monthly Weather Review, Vol. 130, 1812-1828
3261 doi: https://doi.org/10.1175/1520-0493(2002)130<1812:SDOTDA>2.0.CO;2
3262 """
3263 Nt, Ny, Nx = y_3d.shape
3265 # Max coefficient
3266 Nmin = min(Nx - 1, Ny - 1)
3268 # Create alpha matrix (of wavenumbers)
3269 alpha_matrix = _create_alpha_matrix(Ny, Nx)
3271 # Prepare output array
3272 ps_array = np.zeros((Nt, Nmin))
3274 # Loop over time to get spectrum for each time.
3275 for t in range(Nt):
3276 y_2d = y_3d[t]
3278 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space.
3279 # fkk is a 2D array of DCT coefficients, representing the amplitudes of
3280 # cosine basis functions at different spatial frequencies.
3282 # normalise spectrum to allow comparison between models.
3283 fkk = fft.dctn(y_2d, norm="ortho")
3285 # Normalise fkk
3286 fkk = fkk / np.sqrt(Ny * Nx)
3288 # calculate variance of spectral coefficient
3289 sigma_2 = fkk**2 / Nx / Ny
3291 # Group ellipses of alphas into the same wavenumber k/Nmin
3292 for k in range(1, Nmin + 1):
3293 alpha = k / Nmin
3294 alpha_p1 = (k + 1) / Nmin
3296 # Sum up elements matching k
3297 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1))
3298 ps_array[t, k - 1] = np.sum(sigma_2[mask_k])
3300 return ps_array
3303def _create_alpha_matrix(Ny, Nx):
3304 """Construct an array of 2D wavenumbers from 2D wavenumber pair.
3306 Parameters
3307 ----------
3308 Ny, Nx:
3309 Dimensions of the 2D field for which the power spectra is calculated. Used to
3310 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a
3311 single-scale parameter.
3313 Returns: alpha_matrix
3314 normalisation of 2D wavenumber axes, transforming the spectral domain into
3315 an elliptic coordinate system.
3317 """
3318 # Create x_indices: each row is [1, 2, ..., Nx]
3319 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1))
3321 # Create y_indices: each column is [1, 2, ..., Ny]
3322 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx))
3324 # Compute alpha_matrix
3325 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2)
3327 return alpha_matrix