Coverage for src / CSET / operators / plot.py: 82%

1107 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-05-01 15: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. 

14 

15"""Operators to produce various kinds of plots.""" 

16 

17import fcntl 

18import functools 

19import importlib.resources 

20import itertools 

21import json 

22import logging 

23import math 

24import os 

25from typing import Literal 

26 

27import cartopy.crs as ccrs 

28import cartopy.feature as cfeature 

29import iris 

30import iris.coords 

31import iris.cube 

32import iris.exceptions 

33import iris.plot as iplt 

34import matplotlib as mpl 

35import matplotlib.colors as mcolors 

36import matplotlib.pyplot as plt 

37import numpy as np 

38import scipy.fft as fft 

39from cartopy.mpl.geoaxes import GeoAxes 

40from iris.cube import Cube 

41from markdown_it import MarkdownIt 

42from mpl_toolkits.axes_grid1.inset_locator import inset_axes 

43 

44from CSET._common import ( 

45 combine_dicts, 

46 filename_slugify, 

47 get_recipe_metadata, 

48 iter_maybe, 

49 render_file, 

50 slugify, 

51) 

52from CSET.operators._utils import ( 

53 check_stamp_coordinate, 

54 fully_equalise_attributes, 

55 get_cube_yxcoordname, 

56 is_transect, 

57 slice_over_maybe, 

58) 

59from CSET.operators.collapse import collapse 

60from CSET.operators.misc import _extract_common_time_points 

61from CSET.operators.regrid import regrid_onto_cube 

62 

63# Use a non-interactive plotting backend. 

64mpl.use("agg") 

65 

66DEFAULT_DISCRETE_COLORS = mpl.colormaps["tab10"].colors + mpl.colormaps["Accent"].colors 

67 

68############################ 

69# Private helper functions # 

70############################ 

71 

72 

73def _append_to_plot_index(plot_index: list) -> list: 

74 """Add plots into the plot index, returning the complete plot index.""" 

75 with open("meta.json", "r+t", encoding="UTF-8") as fp: 

76 fcntl.flock(fp, fcntl.LOCK_EX) 

77 fp.seek(0) 

78 meta = json.load(fp) 

79 complete_plot_index = meta.get("plots", []) 

80 complete_plot_index = complete_plot_index + plot_index 

81 meta["plots"] = complete_plot_index 

82 if os.getenv("CYLC_TASK_CYCLE_POINT") and not bool( 

83 os.getenv("DO_CASE_AGGREGATION") 

84 ): 

85 meta["case_date"] = os.getenv("CYLC_TASK_CYCLE_POINT", "") 

86 fp.seek(0) 

87 fp.truncate() 

88 json.dump(meta, fp, indent=2) 

89 return complete_plot_index 

90 

91 

92def _check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube: 

93 """Ensure a single cube is given. 

94 

95 If a CubeList of length one is given that the contained cube is returned, 

96 otherwise an error is raised. 

97 

98 Parameters 

99 ---------- 

100 cube: Cube | CubeList 

101 The cube to check. 

102 

103 Returns 

104 ------- 

105 cube: Cube 

106 The checked cube. 

107 

108 Raises 

109 ------ 

110 TypeError 

111 If the input cube is not a Cube or CubeList of a single Cube. 

112 """ 

113 if isinstance(cube, iris.cube.Cube): 

114 return cube 

115 if isinstance(cube, iris.cube.CubeList): 

116 if len(cube) == 1: 

117 return cube[0] 

118 raise TypeError("Must have a single cube", cube) 

119 

120 

121def _make_plot_html_page(plots: list): 

122 """Create a HTML page to display a plot image.""" 

123 # Debug check that plots actually contains some strings. 

124 assert isinstance(plots[0], str) 

125 

126 # Load HTML template file. 

127 operator_files = importlib.resources.files() 

128 template_file = operator_files.joinpath("_plot_page_template.html") 

129 

130 # Get some metadata. 

131 meta = get_recipe_metadata() 

132 title = meta.get("title", "Untitled") 

133 description = MarkdownIt().render(meta.get("description", "*No description.*")) 

134 

135 # Prepare template variables. 

136 variables = { 

137 "title": title, 

138 "description": description, 

139 "initial_plot": plots[0], 

140 "plots": plots, 

141 "title_slug": slugify(title), 

142 } 

143 

144 # Render template. 

145 html = render_file(template_file, **variables) 

146 

147 # Save completed HTML. 

148 with open("index.html", "wt", encoding="UTF-8") as fp: 

149 fp.write(html) 

150 

151 

152@functools.cache 

153def _load_colorbar_map(user_colorbar_file: str = None) -> dict: 

154 """Load the colorbar definitions from a file. 

155 

156 This is a separate function to make it cacheable. 

157 """ 

158 colorbar_file = importlib.resources.files().joinpath("_colorbar_definition.json") 

159 with open(colorbar_file, "rt", encoding="UTF-8") as fp: 

160 colorbar = json.load(fp) 

161 

162 logging.debug("User colour bar file: %s", user_colorbar_file) 

163 override_colorbar = {} 

164 if user_colorbar_file: 

165 try: 

166 with open(user_colorbar_file, "rt", encoding="UTF-8") as fp: 

167 override_colorbar = json.load(fp) 

168 except FileNotFoundError: 

169 logging.warning("Colorbar file does not exist. Using default values.") 

170 

171 # Overwrite values with the user supplied colorbar definition. 

172 colorbar = combine_dicts(colorbar, override_colorbar) 

173 return colorbar 

174 

175 

176def _get_model_colors_map(cubes: iris.cube.CubeList | iris.cube.Cube) -> dict: 

177 """Get an appropriate colors for model lines in line plots. 

178 

179 For each model in the list of cubes colors either from user provided 

180 color definition file (so-called style file) or from default colors are mapped 

181 to model_name attribute. 

182 

183 Parameters 

184 ---------- 

185 cubes: CubeList or Cube 

186 Cubes with model_name attribute 

187 

188 Returns 

189 ------- 

190 model_colors_map: 

191 Dictionary mapping model_name attribute to colors 

192 """ 

193 user_colorbar_file = get_recipe_metadata().get("style_file_path", None) 

194 colorbar = _load_colorbar_map(user_colorbar_file) 

195 model_names = sorted( 

196 filter( 

197 lambda x: x is not None, 

198 (cube.attributes.get("model_name", None) for cube in iter_maybe(cubes)), 

199 ) 

200 ) 

201 if not model_names: 

202 return {} 

203 use_user_colors = all(mname in colorbar.keys() for mname in model_names) 

204 if use_user_colors: 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true

205 return {mname: colorbar[mname] for mname in model_names} 

206 

207 color_list = itertools.cycle(DEFAULT_DISCRETE_COLORS) 

208 return {mname: color for mname, color in zip(model_names, color_list, strict=False)} 

209 

210 

211def _colorbar_map_levels(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None): 

212 """Get an appropriate colorbar for the given cube. 

213 

214 For the given variable the appropriate colorbar is looked up from a 

215 combination of the built-in CSET colorbar definitions, and any user supplied 

216 definitions. As well as varying on variables, these definitions may also 

217 exist for specific pressure levels to account for variables with 

218 significantly different ranges at different heights. The colorbars also exist 

219 for masks and mask differences for considering variable presence diagnostics. 

220 Specific variable ranges can be separately set in user-supplied definition 

221 for x- or y-axis limits, or indicate where automated range preferred. 

222 

223 Parameters 

224 ---------- 

225 cube: Cube 

226 Cube of variable for which the colorbar information is desired. 

227 axis: "x", "y", optional 

228 Select the levels for just this axis of a line plot. The min and max 

229 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

230 setting a universal range is not desirable (e.g. temperature), users 

231 can set ymin/ymax values to "auto" in the colorbar definitions file. 

232 Where no additional xmin/xmax or ymin/ymax values are provided, the 

233 axis bounds default to use the vmin/vmax values provided. 

234 

235 Returns 

236 ------- 

237 cmap: 

238 Matplotlib colormap. 

239 levels: 

240 List of levels to use for plotting. For continuous plots the min and max 

241 should be taken as the range. 

242 norm: 

243 BoundaryNorm information. 

244 """ 

245 # Grab the colorbar file from the recipe global metadata. 

246 user_colorbar_file = get_recipe_metadata().get("style_file_path", None) 

247 colorbar = _load_colorbar_map(user_colorbar_file) 

248 cmap = None 

249 

250 try: 

251 # We assume that pressure is a scalar coordinate here. 

252 pressure_level_raw = cube.coord("pressure").points[0] 

253 # Ensure pressure_level is a string, as it is used as a JSON key. 

254 pressure_level = str(int(pressure_level_raw)) 

255 except iris.exceptions.CoordinateNotFoundError: 

256 pressure_level = None 

257 

258 # First try long name, then standard name, then var name. This order is used 

259 # as long name is the one we correct between models, so it most likely to be 

260 # consistent. 

261 varnames = list(filter(None, [cube.long_name, cube.standard_name, cube.var_name])) 

262 for varname in varnames: 

263 # Get the colormap for this variable. 

264 try: 

265 var_colorbar = colorbar[varname] 

266 cmap = plt.get_cmap(colorbar[varname]["cmap"], 51) 

267 varname_key = varname 

268 break 

269 except KeyError: 

270 logging.debug("Cube name %s has no colorbar definition.", varname) 

271 

272 # Get colormap if it is a mask. 

273 if any("mask_for_" in name for name in varnames): 

274 cmap, levels, norm = _custom_colormap_mask(cube, axis=axis) 

275 return cmap, levels, norm 

276 # If winds on Beaufort Scale use custom colorbar and levels 

277 if any("Beaufort_Scale" in name for name in varnames): 

278 cmap, levels, norm = _custom_beaufort_scale(cube, axis=axis) 

279 return cmap, levels, norm 

280 # If probability is plotted use custom colorbar and levels 

281 if any("probability_of_" in name for name in varnames): 

282 cmap, levels, norm = _custom_colormap_probability(cube, axis=axis) 

283 return cmap, levels, norm 

284 # If aviation colour state use custom colorbar and levels 

285 if any("aviation_colour_state" in name for name in varnames): 

286 cmap, levels, norm = _custom_colormap_aviation_colour_state(cube) 

287 return cmap, levels, norm 

288 

289 # If no valid colormap has been defined, use defaults and return. 

290 if not cmap: 

291 logging.warning("No colorbar definition exists for %s.", cube.name()) 

292 cmap, levels, norm = mpl.colormaps["viridis"], None, None 

293 return cmap, levels, norm 

294 

295 # Test if pressure-level specific settings are provided for cube. 

296 if pressure_level: 

297 try: 

298 var_colorbar = colorbar[varname_key]["pressure_levels"][pressure_level] 

299 except KeyError: 

300 logging.debug( 

301 "%s has no colorbar definition for pressure level %s.", 

302 varname, 

303 pressure_level, 

304 ) 

305 

306 # Check for availability of x-axis or y-axis user-specific overrides 

307 # for setting level bounds for line plot types and return just levels. 

308 # Line plots do not need a colormap, and just use the data range. 

309 if axis: 

310 if axis == "x": 

311 try: 

312 vmin, vmax = var_colorbar["xmin"], var_colorbar["xmax"] 

313 except KeyError: 

314 vmin, vmax = var_colorbar["min"], var_colorbar["max"] 

315 if axis == "y": 

316 try: 

317 vmin, vmax = var_colorbar["ymin"], var_colorbar["ymax"] 

318 except KeyError: 

319 vmin, vmax = var_colorbar["min"], var_colorbar["max"] 

320 # Check if user-specified auto-scaling for this variable 

321 if vmin == "auto" or vmax == "auto": 

322 levels = None 

323 else: 

324 levels = [vmin, vmax] 

325 return None, levels, None 

326 # Get and use the colorbar levels for this variable if spatial or histogram. 

327 else: 

328 try: 

329 levels = var_colorbar["levels"] 

330 # Use discrete bins when levels are specified, rather 

331 # than a smooth range. 

332 norm = mpl.colors.BoundaryNorm(levels, ncolors=cmap.N) 

333 logging.debug("Using levels for %s colorbar.", varname) 

334 logging.info("Using levels: %s", levels) 

335 except KeyError: 

336 # Get the range for this variable. 

337 vmin, vmax = var_colorbar["min"], var_colorbar["max"] 

338 logging.debug("Using min and max for %s colorbar.", varname) 

339 # Calculate levels from range. 

340 if vmin == "auto" or vmax == "auto": 340 ↛ 341line 340 didn't jump to line 341 because the condition on line 340 was never true

341 levels = None 

342 else: 

343 levels = np.linspace(vmin, vmax, 101) 

344 norm = None 

345 

346 # Overwrite cmap, levels and norm for specific variables that 

347 # require custom colorbar_map as these can not be defined in the 

348 # JSON file. 

349 cmap, levels, norm = _custom_colourmap_nimrod_weights(cube, cmap, levels, norm) 

350 cmap, levels, norm = _custom_colourmap_precipitation(cube, cmap, levels, norm) 

351 cmap, levels, norm = _custom_colourmap_visibility_in_air( 

352 cube, cmap, levels, norm 

353 ) 

354 cmap, levels, norm = _custom_colormap_celsius(cube, cmap, levels, norm) 

355 return cmap, levels, norm 

356 

357 

358def _setup_spatial_map( 

359 cube: iris.cube.Cube, 

360 figure, 

361 cmap, 

362 grid_size: tuple[int, int] | None = None, 

363 subplot: int | None = None, 

364): 

365 """Define map projections, extent and add coastlines and borderlines for spatial plots. 

366 

367 For spatial map plots, a relevant map projection for rotated or non-rotated inputs 

368 is specified, and map extent defined based on the input data. 

369 

370 Parameters 

371 ---------- 

372 cube: Cube 

373 2 dimensional (lat and lon) Cube of the data to plot. 

374 figure: 

375 Matplotlib Figure object holding all plot elements. 

376 cmap: 

377 Matplotlib colormap. 

378 grid_size: (int, int), optional 

379 Size of grid (rows, cols) for subplots if multiple spatial subplots in figure. 

380 subplot: int, optional 

381 Subplot index if multiple spatial subplots in figure. 

382 

383 Returns 

384 ------- 

385 axes: 

386 Matplotlib GeoAxes definition. 

387 """ 

388 # Identify min/max plot bounds. 

389 try: 

390 lat_axis, lon_axis = get_cube_yxcoordname(cube) 

391 x1 = np.min(cube.coord(lon_axis).points) 

392 x2 = np.max(cube.coord(lon_axis).points) 

393 y1 = np.min(cube.coord(lat_axis).points) 

394 y2 = np.max(cube.coord(lat_axis).points) 

395 

396 # Adjust bounds within +/- 180.0 if x dimension extends beyond half-globe. 

397 if np.abs(x2 - x1) > 180.0: 

398 x1 = x1 - 180.0 

399 x2 = x2 - 180.0 

400 logging.debug("Adjusting plot bounds to fit global extent.") 

401 

402 # Consider map projection orientation. 

403 # Adapting orientation enables plotting across international dateline. 

404 # Users can adapt the default central_longitude if alternative projections views. 

405 if x2 > 180.0 or x1 < -180.0: 

406 central_longitude = 180.0 

407 else: 

408 central_longitude = 0.0 

409 

410 # Define spatial map projection. 

411 coord_system = cube.coord(lat_axis).coord_system 

412 if isinstance(coord_system, iris.coord_systems.RotatedGeogCS): 

413 # Define rotated pole map projection for rotated pole inputs. 

414 projection = ccrs.RotatedPole( 

415 pole_longitude=coord_system.grid_north_pole_longitude, 

416 pole_latitude=coord_system.grid_north_pole_latitude, 

417 central_rotated_longitude=central_longitude, 

418 ) 

419 crs = projection 

420 elif isinstance(coord_system, iris.coord_systems.TransverseMercator): 420 ↛ 422line 420 didn't jump to line 422 because the condition on line 420 was never true

421 # Define Transverse Mercator projection for TM inputs. 

422 projection = ccrs.TransverseMercator( 

423 central_longitude=coord_system.longitude_of_central_meridian, 

424 central_latitude=coord_system.latitude_of_projection_origin, 

425 false_easting=coord_system.false_easting, 

426 false_northing=coord_system.false_northing, 

427 scale_factor=coord_system.scale_factor_at_central_meridian, 

428 ) 

429 crs = projection 

430 else: 

431 # Define regular map projection for non-rotated pole inputs. 

432 # Alternatives might include e.g. for global model outputs: 

433 # projection=ccrs.Robinson(central_longitude=X.y, globe=None) 

434 # See also https://scitools.org.uk/cartopy/docs/v0.15/crs/projections.html. 

435 projection = ccrs.PlateCarree(central_longitude=central_longitude) 

436 crs = ccrs.PlateCarree() 

437 

438 # Define axes for plot (or subplot) with required map projection. 

439 if subplot is not None: 

440 axes = figure.add_subplot( 

441 grid_size[0], grid_size[1], subplot, projection=projection 

442 ) 

443 else: 

444 axes = figure.add_subplot(projection=projection) 

445 

446 # Add coastlines and borderlines if cube contains x and y map coordinates. 

447 # Avoid adding lines for masked data or specific fixed ancillary spatial plots. 

448 if iris.util.is_masked(cube.data) or any( 448 ↛ 451line 448 didn't jump to line 451 because the condition on line 448 was never true

449 name in cube.name() for name in ["land_", "orography", "altitude"] 

450 ): 

451 pass 

452 else: 

453 if cmap.name in ["viridis", "Greys"]: 

454 coastcol = "magenta" 

455 else: 

456 coastcol = "black" 

457 logging.debug("Plotting coastlines and borderlines in colour %s.", coastcol) 

458 axes.coastlines(resolution="10m", color=coastcol) 

459 axes.add_feature(cfeature.BORDERS, edgecolor=coastcol) 

460 

461 # Add gridlines. 

462 gl = axes.gridlines( 

463 alpha=0.3, 

464 draw_labels=True, 

465 dms=False, 

466 x_inline=False, 

467 y_inline=False, 

468 ) 

469 gl.top_labels = False 

470 gl.right_labels = False 

471 if subplot: 

472 gl.bottom_labels = False 

473 gl.left_labels = False 

474 if subplot % grid_size[1] == 1: 

475 gl.left_labels = True 

476 if subplot > ((grid_size[0] - 1) * grid_size[1]): 476 ↛ 481line 476 didn't jump to line 481 because the condition on line 476 was always true

477 gl.bottom_labels = True 

478 

479 # If is lat/lon spatial map, fix extent to keep plot tight. 

480 # Specifying crs within set_extent helps ensure only data region is shown. 

481 if isinstance(coord_system, iris.coord_systems.GeogCS): 

482 axes.set_extent([x1, x2, y1, y2], crs=crs) 

483 

484 except ValueError: 

485 # Skip if not both x and y map coordinates. 

486 axes = figure.gca() 

487 pass 

488 

489 return axes 

490 

491 

492def _get_plot_resolution() -> int: 

493 """Get resolution of rasterised plots in pixels per inch.""" 

494 return get_recipe_metadata().get("plot_resolution", 100) 

495 

496 

497def _get_start_end_strings(seq_coord: iris.coords.Coord, use_bounds: bool): 

498 """Return title and filename based on start and end points or bounds.""" 

499 if use_bounds and seq_coord.has_bounds(): 

500 vals = seq_coord.bounds.flatten() 

501 else: 

502 vals = seq_coord.points 

503 start = seq_coord.units.title(vals[0]) 

504 end = seq_coord.units.title(vals[-1]) 

505 

506 if start == end: 

507 sequence_title = f"\n [{start}]" 

508 sequence_fname = f"_{filename_slugify(start)}" 

509 else: 

510 sequence_title = f"\n [{start} to {end}]" 

511 sequence_fname = f"_{filename_slugify(start)}_{filename_slugify(end)}" 

512 

513 # Do not include time if coord set to zero. 

514 if ( 

515 seq_coord.units == "hours since 0001-01-01 00:00:00" 

516 and vals[0] == 0 

517 and vals[-1] == 0 

518 ): 

519 sequence_title = "" 

520 sequence_fname = "" 

521 

522 return sequence_title, sequence_fname 

523 

524 

525def _set_title_and_filename( 

526 seq_coord: iris.coords.Coord, 

527 nplot: int, 

528 recipe_title: str, 

529 filename: str, 

530): 

531 """Set plot title and filename based on cube coordinate. 

532 

533 Parameters 

534 ---------- 

535 sequence_coordinate: iris.coords.Coord 

536 Coordinate about which to make a plot sequence. 

537 nplot: int 

538 Number of output plots to generate - controls title/naming. 

539 recipe_title: str 

540 Default plot title, potentially to update. 

541 filename: str 

542 Input plot filename, potentially to update. 

543 

544 Returns 

545 ------- 

546 plot_title: str 

547 Output formatted plot title string, based on plotted data. 

548 plot_filename: str 

549 Output formatted plot filename string. 

550 """ 

551 ndim = seq_coord.ndim 

552 npoints = np.size(seq_coord.points) 

553 sequence_title = "" 

554 sequence_fname = "" 

555 

556 # Case 1: Multiple dimension sequence input - list number of aggregated cases 

557 # (e.g. aggregation histogram plots) 

558 if ndim > 1: 

559 ncase = np.shape(seq_coord)[0] 

560 sequence_title = f"\n [{ncase} cases]" 

561 sequence_fname = f"_{ncase}cases" 

562 

563 # Case 2: Single dimension input 

564 else: 

565 # Single sequence point 

566 if npoints == 1: 

567 if nplot > 1: 

568 # Default labels for sequence inputs 

569 sequence_value = seq_coord.units.title(seq_coord.points[0]) 

570 sequence_title = f"\n [{sequence_value}]" 

571 sequence_fname = f"_{filename_slugify(sequence_value)}" 

572 else: 

573 # Aggregated attribute available where input collapsed over aggregation 

574 try: 

575 ncase = seq_coord.attributes["number_reference_times"] 

576 sequence_title = f"\n [{ncase} cases]" 

577 sequence_fname = f"_{ncase}cases" 

578 except KeyError: 

579 sequence_title, sequence_fname = _get_start_end_strings( 

580 seq_coord, use_bounds=seq_coord.has_bounds() 

581 ) 

582 # Multiple sequence (e.g. time) points 

583 else: 

584 sequence_title, sequence_fname = _get_start_end_strings( 

585 seq_coord, use_bounds=False 

586 ) 

587 

588 # Set plot title and filename 

589 plot_title = f"{recipe_title}{sequence_title}" 

590 

591 # Set plot filename, defaulting to user input if provided. 

592 if filename is None: 

593 filename = slugify(recipe_title) 

594 plot_filename = f"{filename.rsplit('.', 1)[0]}{sequence_fname}.png" 

595 else: 

596 if nplot > 1: 

597 plot_filename = f"{filename.rsplit('.', 1)[0]}{sequence_fname}.png" 

598 else: 

599 plot_filename = f"{filename.rsplit('.', 1)[0]}.png" 

600 

601 return plot_title, plot_filename 

602 

603 

604def _set_postage_stamp_title(stamp_coord: iris.coords.Coord) -> str: 

605 """Control postage stamp plot output titles based on stamp coordinate.""" 

606 if stamp_coord.name() == "realization": 

607 mtitle = "Member" 

608 else: 

609 mtitle = stamp_coord.name().capitalize() 

610 

611 if stamp_coord.name() == "time": 

612 mtitle = f"{stamp_coord.units.title(stamp_coord.points[0])}" 

613 else: 

614 mtitle = f"{mtitle} #{stamp_coord.points[0]}" 

615 

616 return mtitle 

617 

618 

619def _plot_and_save_spatial_plot( 

620 cube: iris.cube.Cube, 

621 filename: str, 

622 title: str, 

623 method: Literal["contourf", "pcolormesh"], 

624 overlay_cube: iris.cube.Cube | None = None, 

625 contour_cube: iris.cube.Cube | None = None, 

626 **kwargs, 

627): 

628 """Plot and save a spatial plot. 

629 

630 Parameters 

631 ---------- 

632 cube: Cube 

633 2 dimensional (lat and lon) Cube of the data to plot. 

634 filename: str 

635 Filename of the plot to write. 

636 title: str 

637 Plot title. 

638 method: "contourf" | "pcolormesh" 

639 The plotting method to use. 

640 overlay_cube: Cube, optional 

641 Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube 

642 contour_cube: Cube, optional 

643 Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube 

644 """ 

645 # Setup plot details, size, resolution, etc. 

646 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

647 

648 # Specify the color bar 

649 cmap, levels, norm = _colorbar_map_levels(cube) 

650 

651 # If overplotting, set required colorbars 

652 if overlay_cube: 

653 over_cmap, over_levels, over_norm = _colorbar_map_levels(overlay_cube) 

654 if contour_cube: 

655 cntr_cmap, cntr_levels, cntr_norm = _colorbar_map_levels(contour_cube) 

656 

657 # Setup plot map projection, extent and coastlines and borderlines. 

658 axes = _setup_spatial_map(cube, fig, cmap) 

659 

660 # Plot the field. 

661 if method == "contourf": 

662 # Filled contour plot of the field. 

663 plot = iplt.contourf(cube, cmap=cmap, levels=levels, norm=norm) 

664 elif method == "pcolormesh": 

665 try: 

666 vmin = min(levels) 

667 vmax = max(levels) 

668 except TypeError: 

669 vmin, vmax = None, None 

670 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax 

671 # if levels are defined. 

672 if norm is not None: 

673 vmin = None 

674 vmax = None 

675 logging.debug("Plotting using defined levels.") 

676 plot = iplt.pcolormesh(cube, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax) 

677 else: 

678 raise ValueError(f"Unknown plotting method: {method}") 

679 

680 # Overplot overlay field, if required 

681 if overlay_cube: 

682 try: 

683 over_vmin = min(over_levels) 

684 over_vmax = max(over_levels) 

685 except TypeError: 

686 over_vmin, over_vmax = None, None 

687 if over_norm is not None: 687 ↛ 688line 687 didn't jump to line 688 because the condition on line 687 was never true

688 over_vmin = None 

689 over_vmax = None 

690 overlay = iplt.pcolormesh( 

691 overlay_cube, 

692 cmap=over_cmap, 

693 norm=over_norm, 

694 alpha=0.8, 

695 vmin=over_vmin, 

696 vmax=over_vmax, 

697 ) 

698 # Overplot contour field, if required, with contour labelling. 

699 if contour_cube: 

700 contour = iplt.contour( 

701 contour_cube, 

702 colors="darkgray", 

703 levels=cntr_levels, 

704 norm=cntr_norm, 

705 alpha=0.5, 

706 linestyles="--", 

707 linewidths=1, 

708 ) 

709 plt.clabel(contour) 

710 

711 # Check to see if transect, and if so, adjust y axis. 

712 if is_transect(cube): 

713 if "pressure" in [coord.name() for coord in cube.coords()]: 

714 axes.invert_yaxis() 

715 axes.set_yscale("log") 

716 axes.set_ylim(1100, 100) 

717 # If both model_level_number and level_height exists, iplt can construct 

718 # plot as a function of height above orography (NOT sea level). 

719 elif {"model_level_number", "level_height"}.issubset( 719 ↛ 724line 719 didn't jump to line 724 because the condition on line 719 was always true

720 {coord.name() for coord in cube.coords()} 

721 ): 

722 axes.set_yscale("log") 

723 

724 axes.set_title( 

725 f"{title}\n" 

726 f"Start Lat: {cube.attributes['transect_coords'].split('_')[0]}" 

727 f" Start Lon: {cube.attributes['transect_coords'].split('_')[1]}" 

728 f" End Lat: {cube.attributes['transect_coords'].split('_')[2]}" 

729 f" End Lon: {cube.attributes['transect_coords'].split('_')[3]}", 

730 fontsize=16, 

731 ) 

732 

733 # Inset code 

734 axins = inset_axes( 

735 axes, 

736 width="20%", 

737 height="20%", 

738 loc="upper right", 

739 axes_class=GeoAxes, 

740 axes_kwargs={"map_projection": ccrs.PlateCarree()}, 

741 ) 

742 

743 # Slightly transparent to reduce plot blocking. 

744 axins.patch.set_alpha(0.4) 

745 

746 axins.coastlines(resolution="50m") 

747 axins.add_feature(cfeature.BORDERS, linewidth=0.3) 

748 

749 SLat, SLon, ELat, ELon = ( 

750 float(coord) for coord in cube.attributes["transect_coords"].split("_") 

751 ) 

752 

753 # Draw line between them 

754 axins.plot( 

755 [SLon, ELon], [SLat, ELat], color="black", transform=ccrs.PlateCarree() 

756 ) 

757 

758 # Plot points (note: lon, lat order for Cartopy) 

759 axins.plot(SLon, SLat, marker="x", color="green", transform=ccrs.PlateCarree()) 

760 axins.plot(ELon, ELat, marker="x", color="red", transform=ccrs.PlateCarree()) 

761 

762 lon_min, lon_max = sorted([SLon, ELon]) 

763 lat_min, lat_max = sorted([SLat, ELat]) 

764 

765 # Midpoints 

766 lon_mid = (lon_min + lon_max) / 2 

767 lat_mid = (lat_min + lat_max) / 2 

768 

769 # Maximum half-range 

770 half_range = max(lon_max - lon_min, lat_max - lat_min) / 2 

771 if half_range == 0: # points identical → provide small default 771 ↛ 775line 771 didn't jump to line 775 because the condition on line 771 was always true

772 half_range = 1 

773 

774 # Set square extent 

775 axins.set_extent( 

776 [ 

777 lon_mid - half_range, 

778 lon_mid + half_range, 

779 lat_mid - half_range, 

780 lat_mid + half_range, 

781 ], 

782 crs=ccrs.PlateCarree(), 

783 ) 

784 

785 # Ensure square aspect 

786 axins.set_aspect("equal") 

787 

788 else: 

789 # Add title. 

790 axes.set_title(title, fontsize=16) 

791 

792 # Adjust padding if spatial plot or transect 

793 if is_transect(cube): 

794 yinfopad = -0.1 

795 ycbarpad = 0.1 

796 else: 

797 yinfopad = 0.01 

798 ycbarpad = 0.042 

799 

800 # Add watermark with min/max/mean. Currently not user togglable. 

801 # In the bbox dictionary, fc and ec are hex colour codes for grey shade. 

802 axes.annotate( 

803 f"Min: {np.min(cube.data):.3g} Max: {np.max(cube.data):.3g} Mean: {np.mean(cube.data):.3g}", 

804 xy=(0.025, yinfopad), 

805 xycoords="axes fraction", 

806 xytext=(-5, 5), 

807 textcoords="offset points", 

808 ha="left", 

809 va="bottom", 

810 size=11, 

811 bbox=dict(boxstyle="round", fc="#cccccc", ec="#808080", alpha=0.9), 

812 ) 

813 

814 # Add secondary colour bar for overlay_cube field if required. 

815 if overlay_cube: 

816 cbarB = fig.colorbar( 

817 overlay, orientation="horizontal", location="bottom", pad=0.0, shrink=0.7 

818 ) 

819 cbarB.set_label(label=f"{overlay_cube.name()} ({overlay_cube.units})", size=14) 

820 # add ticks and tick_labels for every levels if less than 20 levels exist 

821 if over_levels is not None and len(over_levels) < 20: 821 ↛ 822line 821 didn't jump to line 822 because the condition on line 821 was never true

822 cbarB.set_ticks(over_levels) 

823 cbarB.set_ticklabels([f"{level:.2f}" for level in over_levels]) 

824 if "rainfall" or "snowfall" or "visibility" in overlay_cube.name(): 

825 cbarB.set_ticklabels([f"{level:.3g}" for level in over_levels]) 

826 logging.debug("Set secondary colorbar ticks and labels.") 

827 

828 # Add main colour bar. 

829 cbar = fig.colorbar( 

830 plot, orientation="horizontal", location="bottom", pad=ycbarpad, shrink=0.7 

831 ) 

832 

833 cbar.set_label(label=f"{cube.name()} ({cube.units})", size=14) 

834 # add ticks and tick_labels for every levels if less than 20 levels exist 

835 if levels is not None and len(levels) < 20: 

836 cbar.set_ticks(levels) 

837 cbar.set_ticklabels([f"{level:.2f}" for level in levels]) 

838 if "rainfall" or "snowfall" or "visibility" in cube.name(): 838 ↛ 841line 838 didn't jump to line 841 because the condition on line 838 was always true

839 cbar.set_ticklabels([f"{level:.3g}" for level in levels]) 

840 # Tick labels for rainfall rates from Nimrod radar data. 

841 if "rainfall rate composite" in cube.name(): 841 ↛ 842line 841 didn't jump to line 842 because the condition on line 841 was never true

842 cbar.set_ticklabels([f"{level:.3g}" for level in levels]) 

843 # Tick labels for rain accumulations from Nimrod radar data. 

844 if "rain accumulation" in cube.name(): 844 ↛ 845line 844 didn't jump to line 845 because the condition on line 844 was never true

845 cbar.set_ticklabels([f"{level:.3g}" for level in levels]) 

846 if "wts accumulation" in cube.name(): 846 ↛ 847line 846 didn't jump to line 847 because the condition on line 846 was never true

847 tick_levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] 

848 cbar.minorticks_off() 

849 cbar.set_ticks(tick_levels) 

850 cbar.set_ticklabels([f"{level:.3g}" for level in tick_levels]) 

851 cbar.set_label(label=f"{cube.name()}", size=14) 

852 # Tick labels for model rainfall data. 

853 if "surface_microphysical" in cube.name(): 853 ↛ 856line 853 didn't jump to line 856 because the condition on line 853 was always true

854 cbar.set_ticklabels([f"{level:.3g}" for level in levels]) 

855 # Tick labels for Nimrod weights data. 

856 logging.debug("Set colorbar ticks and labels.") 

857 

858 # Save plot. 

859 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

860 logging.info("Saved spatial plot to %s", filename) 

861 plt.close(fig) 

862 

863 

864def _plot_and_save_postage_stamp_spatial_plot( 

865 cube: iris.cube.Cube, 

866 filename: str, 

867 stamp_coordinate: str, 

868 title: str, 

869 method: Literal["contourf", "pcolormesh"], 

870 overlay_cube: iris.cube.Cube | None = None, 

871 contour_cube: iris.cube.Cube | None = None, 

872 **kwargs, 

873): 

874 """Plot postage stamp spatial plots from an ensemble. 

875 

876 Parameters 

877 ---------- 

878 cube: Cube 

879 Iris cube of data to be plotted. It must have the stamp coordinate. 

880 filename: str 

881 Filename of the plot to write. 

882 stamp_coordinate: str 

883 Coordinate that becomes different plots. 

884 method: "contourf" | "pcolormesh" 

885 The plotting method to use. 

886 overlay_cube: Cube, optional 

887 Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube 

888 contour_cube: Cube, optional 

889 Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube 

890 

891 Raises 

892 ------ 

893 ValueError 

894 If the cube doesn't have the right dimensions. 

895 """ 

896 # Use the smallest square grid that will fit the members. 

897 nmember = len(cube.coord(stamp_coordinate).points) 

898 grid_rows = int(math.sqrt(nmember)) 

899 grid_size = math.ceil(nmember / grid_rows) 

900 

901 fig = plt.figure( 

902 figsize=(10, 10 * max(grid_rows / grid_size, 0.5)), facecolor="w", edgecolor="k" 

903 ) 

904 

905 # Specify the color bar 

906 cmap, levels, norm = _colorbar_map_levels(cube) 

907 # If overplotting, set required colorbars 

908 if overlay_cube: 908 ↛ 909line 908 didn't jump to line 909 because the condition on line 908 was never true

909 over_cmap, over_levels, over_norm = _colorbar_map_levels(overlay_cube) 

910 if contour_cube: 910 ↛ 911line 910 didn't jump to line 911 because the condition on line 910 was never true

911 cntr_cmap, cntr_levels, cntr_norm = _colorbar_map_levels(contour_cube) 

912 

913 # Make a subplot for each member. 

914 for member, subplot in zip( 

915 cube.slices_over(stamp_coordinate), 

916 range(1, grid_size * grid_rows + 1), 

917 strict=False, 

918 ): 

919 # Setup subplot map projection, extent and coastlines and borderlines. 

920 axes = _setup_spatial_map( 

921 member, fig, cmap, grid_size=(grid_rows, grid_size), subplot=subplot 

922 ) 

923 if method == "contourf": 

924 # Filled contour plot of the field. 

925 plot = iplt.contourf(member, cmap=cmap, levels=levels, norm=norm) 

926 elif method == "pcolormesh": 

927 if levels is not None: 

928 vmin = min(levels) 

929 vmax = max(levels) 

930 else: 

931 raise TypeError("Unknown vmin and vmax range.") 

932 vmin, vmax = None, None 

933 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax 

934 # if levels are defined. 

935 if norm is not None: 935 ↛ 936line 935 didn't jump to line 936 because the condition on line 935 was never true

936 vmin = None 

937 vmax = None 

938 # pcolormesh plot of the field. 

939 plot = iplt.pcolormesh(member, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax) 

940 else: 

941 raise ValueError(f"Unknown plotting method: {method}") 

942 

943 # Overplot overlay field, if required 

944 if overlay_cube: 944 ↛ 945line 944 didn't jump to line 945 because the condition on line 944 was never true

945 try: 

946 over_vmin = min(over_levels) 

947 over_vmax = max(over_levels) 

948 except TypeError: 

949 over_vmin, over_vmax = None, None 

950 if over_norm is not None: 

951 over_vmin = None 

952 over_vmax = None 

953 iplt.pcolormesh( 

954 overlay_cube[member.coord(stamp_coordinate).points[0]], 

955 cmap=over_cmap, 

956 norm=over_norm, 

957 alpha=0.6, 

958 vmin=over_vmin, 

959 vmax=over_vmax, 

960 ) 

961 # Overplot contour field, if required 

962 if contour_cube: 962 ↛ 963line 962 didn't jump to line 963 because the condition on line 962 was never true

963 iplt.contour( 

964 contour_cube[member.coord(stamp_coordinate).points[0]], 

965 colors="darkgray", 

966 levels=cntr_levels, 

967 norm=cntr_norm, 

968 alpha=0.6, 

969 linestyles="--", 

970 linewidths=1, 

971 ) 

972 mtitle = _set_postage_stamp_title(member.coord(stamp_coordinate)) 

973 axes.set_title(f"{mtitle}") 

974 

975 # Put the shared colorbar in its own axes. 

976 colorbar_axes = fig.add_axes([0.15, 0.05, 0.7, 0.03]) 

977 colorbar = fig.colorbar( 

978 plot, colorbar_axes, orientation="horizontal", pad=0.042, shrink=0.7 

979 ) 

980 colorbar.set_label(f"{cube.name()} ({cube.units})", size=14) 

981 

982 # Overall figure title. 

983 fig.suptitle(title, fontsize=16) 

984 

985 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

986 logging.info("Saved contour postage stamp plot to %s", filename) 

987 plt.close(fig) 

988 

989 

990def _plot_and_save_line_series( 

991 cubes: iris.cube.CubeList, 

992 coords: list[iris.coords.Coord], 

993 ensemble_coord: str, 

994 filename: str, 

995 title: str, 

996 **kwargs, 

997): 

998 """Plot and save a 1D line series. 

999 

1000 Parameters 

1001 ---------- 

1002 cubes: Cube or CubeList 

1003 Cube or CubeList containing the cubes to plot on the y-axis. 

1004 coords: list[Coord] 

1005 Coordinates to plot on the x-axis, one per cube. 

1006 ensemble_coord: str 

1007 Ensemble coordinate in the cube. 

1008 filename: str 

1009 Filename of the plot to write. 

1010 title: str 

1011 Plot title. 

1012 """ 

1013 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

1014 

1015 model_colors_map = _get_model_colors_map(cubes) 

1016 

1017 # Store min/max ranges. 

1018 y_levels = [] 

1019 

1020 # Check match-up across sequence coords gives consistent sizes 

1021 _validate_cubes_coords(cubes, coords) 

1022 

1023 for cube, coord in zip(cubes, coords, strict=True): 

1024 label = None 

1025 color = "black" 

1026 if model_colors_map: 

1027 label = cube.attributes.get("model_name") 

1028 color = model_colors_map.get(label) 

1029 for cube_slice in cube.slices_over(ensemble_coord): 

1030 # Label with (control) if part of an ensemble or not otherwise. 

1031 if cube_slice.coord(ensemble_coord).points == [0]: 

1032 iplt.plot( 

1033 coord, 

1034 cube_slice, 

1035 color=color, 

1036 marker="o", 

1037 ls="-", 

1038 lw=3, 

1039 label=f"{label} (control)" 

1040 if len(cube.coord(ensemble_coord).points) > 1 

1041 else label, 

1042 ) 

1043 # Label with (perturbed) if part of an ensemble and not the control. 

1044 else: 

1045 iplt.plot( 

1046 coord, 

1047 cube_slice, 

1048 color=color, 

1049 ls="-", 

1050 lw=1.5, 

1051 alpha=0.75, 

1052 label=f"{label} (member)", 

1053 ) 

1054 

1055 # Calculate the global min/max if multiple cubes are given. 

1056 _, levels, _ = _colorbar_map_levels(cube, axis="y") 

1057 if levels is not None: 1057 ↛ 1058line 1057 didn't jump to line 1058 because the condition on line 1057 was never true

1058 y_levels.append(min(levels)) 

1059 y_levels.append(max(levels)) 

1060 

1061 # Get the current axes. 

1062 ax = plt.gca() 

1063 

1064 # Add some labels and tweak the style. 

1065 # check if cubes[0] works for single cube if not CubeList 

1066 if coords[0].name() == "time": 

1067 ax.set_xlabel(f"{coords[0].name()}", fontsize=14) 

1068 else: 

1069 ax.set_xlabel(f"{coords[0].name()} / {coords[0].units}", fontsize=14) 

1070 ax.set_ylabel(f"{cubes[0].name()} / {cubes[0].units}", fontsize=14) 

1071 ax.set_title(title, fontsize=16) 

1072 

1073 ax.ticklabel_format(axis="y", useOffset=False) 

1074 ax.tick_params(axis="x", labelrotation=15) 

1075 ax.tick_params(axis="both", labelsize=12) 

1076 

1077 # Set y limits to global min and max, autoscale if colorbar doesn't exist. 

1078 if y_levels: 1078 ↛ 1079line 1078 didn't jump to line 1079 because the condition on line 1078 was never true

1079 ax.set_ylim(min(y_levels), max(y_levels)) 

1080 # Add zero line. 

1081 if min(y_levels) < 0.0 and max(y_levels) > 0.0: 

1082 ax.axhline(y=0, xmin=0, xmax=1, ls="-", color="grey", lw=2) 

1083 logging.debug( 

1084 "Line plot with y-axis limits %s-%s", min(y_levels), max(y_levels) 

1085 ) 

1086 else: 

1087 ax.autoscale() 

1088 

1089 # Add gridlines 

1090 ax.grid(linestyle="--", color="grey", linewidth=1) 

1091 # Ientify unique labels for legend 

1092 handles = list( 

1093 { 

1094 label: handle 

1095 for (handle, label) in zip(*ax.get_legend_handles_labels(), strict=True) 

1096 }.values() 

1097 ) 

1098 ax.legend(handles=handles, loc="best", ncol=1, frameon=False, fontsize=16) 

1099 

1100 # Save plot. 

1101 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1102 logging.info("Saved line plot to %s", filename) 

1103 plt.close(fig) 

1104 

1105 

1106def _plot_and_save_vertical_line_series( 

1107 cubes: iris.cube.CubeList, 

1108 coords: list[iris.coords.Coord], 

1109 ensemble_coord: str, 

1110 filename: str, 

1111 series_coordinate: str, 

1112 title: str, 

1113 vmin: float, 

1114 vmax: float, 

1115 **kwargs, 

1116): 

1117 """Plot and save a 1D line series in vertical. 

1118 

1119 Parameters 

1120 ---------- 

1121 cubes: CubeList 

1122 1 dimensional Cube or CubeList of the data to plot on x-axis. 

1123 coord: list[Coord] 

1124 Coordinates to plot on the y-axis, one per cube. 

1125 ensemble_coord: str 

1126 Ensemble coordinate in the cube. 

1127 filename: str 

1128 Filename of the plot to write. 

1129 series_coordinate: str 

1130 Coordinate to use as vertical axis. 

1131 title: str 

1132 Plot title. 

1133 vmin: float 

1134 Minimum value for the x-axis. 

1135 vmax: float 

1136 Maximum value for the x-axis. 

1137 """ 

1138 # plot the vertical pressure axis using log scale 

1139 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

1140 

1141 model_colors_map = _get_model_colors_map(cubes) 

1142 

1143 # Check match-up across sequence coords gives consistent sizes 

1144 _validate_cubes_coords(cubes, coords) 

1145 

1146 for cube, coord in zip(cubes, coords, strict=True): 

1147 label = None 

1148 color = "black" 

1149 if model_colors_map: 1149 ↛ 1150line 1149 didn't jump to line 1150 because the condition on line 1149 was never true

1150 label = cube.attributes.get("model_name") 

1151 color = model_colors_map.get(label) 

1152 

1153 for cube_slice in cube.slices_over(ensemble_coord): 

1154 # If ensemble data given plot control member with (control) 

1155 # unless single forecast. 

1156 if cube_slice.coord(ensemble_coord).points == [0]: 

1157 iplt.plot( 

1158 cube_slice, 

1159 coord, 

1160 color=color, 

1161 marker="o", 

1162 ls="-", 

1163 lw=3, 

1164 label=f"{label} (control)" 

1165 if len(cube.coord(ensemble_coord).points) > 1 

1166 else label, 

1167 ) 

1168 # If ensemble data given plot perturbed members with (perturbed). 

1169 else: 

1170 iplt.plot( 

1171 cube_slice, 

1172 coord, 

1173 color=color, 

1174 ls="-", 

1175 lw=1.5, 

1176 alpha=0.75, 

1177 label=f"{label} (member)", 

1178 ) 

1179 

1180 # Get the current axis 

1181 ax = plt.gca() 

1182 

1183 # Special handling for pressure level data. 

1184 if series_coordinate == "pressure": 1184 ↛ 1206line 1184 didn't jump to line 1206 because the condition on line 1184 was always true

1185 # Invert y-axis and set to log scale. 

1186 ax.invert_yaxis() 

1187 ax.set_yscale("log") 

1188 

1189 # Define y-ticks and labels for pressure log axis. 

1190 y_tick_labels = [ 

1191 "1000", 

1192 "850", 

1193 "700", 

1194 "500", 

1195 "300", 

1196 "200", 

1197 "100", 

1198 ] 

1199 y_ticks = [1000, 850, 700, 500, 300, 200, 100] 

1200 

1201 # Set y-axis limits and ticks. 

1202 ax.set_ylim(1100, 100) 

1203 

1204 # Test if series_coordinate is model level data. The UM data uses 

1205 # model_level_number and lfric uses full_levels as coordinate. 

1206 elif series_coordinate in ("model_level_number", "full_levels", "half_levels"): 

1207 # Define y-ticks and labels for vertical axis. 

1208 y_ticks = iter_maybe(cubes)[0].coord(series_coordinate).points 

1209 y_tick_labels = [str(int(i)) for i in y_ticks] 

1210 ax.set_ylim(min(y_ticks), max(y_ticks)) 

1211 

1212 ax.set_yticks(y_ticks) 

1213 ax.set_yticklabels(y_tick_labels) 

1214 

1215 # Set x-axis limits. 

1216 ax.set_xlim(vmin, vmax) 

1217 # Mark y=0 if present in plot. 

1218 if vmin < 0.0 and vmax > 0.0: 1218 ↛ 1219line 1218 didn't jump to line 1219 because the condition on line 1218 was never true

1219 ax.axvline(x=0, ymin=0, ymax=1, ls="-", color="grey", lw=2) 

1220 

1221 # Add some labels and tweak the style. 

1222 ax.set_ylabel(f"{coord.name()} / {coord.units}", fontsize=14) 

1223 ax.set_xlabel( 

1224 f"{iter_maybe(cubes)[0].name()} / {iter_maybe(cubes)[0].units}", fontsize=14 

1225 ) 

1226 ax.set_title(title, fontsize=16) 

1227 ax.ticklabel_format(axis="x") 

1228 ax.tick_params(axis="y") 

1229 ax.tick_params(axis="both", labelsize=12) 

1230 

1231 # Add gridlines 

1232 ax.grid(linestyle="--", color="grey", linewidth=1) 

1233 # Ientify unique labels for legend 

1234 handles = list( 

1235 { 

1236 label: handle 

1237 for (handle, label) in zip(*ax.get_legend_handles_labels(), strict=True) 

1238 }.values() 

1239 ) 

1240 ax.legend(handles=handles, loc="best", ncol=1, frameon=False, fontsize=16) 

1241 

1242 # Save plot. 

1243 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1244 logging.info("Saved line plot to %s", filename) 

1245 plt.close(fig) 

1246 

1247 

1248def _plot_and_save_scatter_plot( 

1249 cube_x: iris.cube.Cube | iris.cube.CubeList, 

1250 cube_y: iris.cube.Cube | iris.cube.CubeList, 

1251 filename: str, 

1252 title: str, 

1253 one_to_one: bool, 

1254 model_names: list[str] = None, 

1255 **kwargs, 

1256): 

1257 """Plot and save a 2D scatter plot. 

1258 

1259 Parameters 

1260 ---------- 

1261 cube_x: Cube | CubeList 

1262 1 dimensional Cube or CubeList of the data to plot on x-axis. 

1263 cube_y: Cube | CubeList 

1264 1 dimensional Cube or CubeList of the data to plot on y-axis. 

1265 filename: str 

1266 Filename of the plot to write. 

1267 title: str 

1268 Plot title. 

1269 one_to_one: bool 

1270 Whether a 1:1 line is plotted. 

1271 """ 

1272 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

1273 # plot the cube_x and cube_y 1D fields as a scatter plot. If they are CubeLists this ensures 

1274 # to pair each cube from cube_x with the corresponding cube from cube_y, allowing to iterate 

1275 # over the pairs simultaneously. 

1276 

1277 # Ensure cube_x and cube_y are iterable 

1278 cube_x_iterable = iter_maybe(cube_x) 

1279 cube_y_iterable = iter_maybe(cube_y) 

1280 

1281 for cube_x_iter, cube_y_iter in zip(cube_x_iterable, cube_y_iterable, strict=True): 

1282 iplt.scatter(cube_x_iter, cube_y_iter) 

1283 if one_to_one is True: 

1284 plt.plot( 

1285 [ 

1286 np.nanmin([np.nanmin(cube_y.data), np.nanmin(cube_x.data)]), 

1287 np.nanmax([np.nanmax(cube_y.data), np.nanmax(cube_x.data)]), 

1288 ], 

1289 [ 

1290 np.nanmin([np.nanmin(cube_y.data), np.nanmin(cube_x.data)]), 

1291 np.nanmax([np.nanmax(cube_y.data), np.nanmax(cube_x.data)]), 

1292 ], 

1293 "k", 

1294 linestyle="--", 

1295 ) 

1296 ax = plt.gca() 

1297 

1298 # Add some labels and tweak the style. 

1299 if model_names is None: 

1300 ax.set_xlabel(f"{cube_x[0].name()} / {cube_x[0].units}", fontsize=14) 

1301 ax.set_ylabel(f"{cube_y[0].name()} / {cube_y[0].units}", fontsize=14) 

1302 else: 

1303 # Add the model names, these should be order of base (x) and other (y). 

1304 ax.set_xlabel( 

1305 f"{model_names[0]}_{cube_x[0].name()} / {cube_x[0].units}", fontsize=14 

1306 ) 

1307 ax.set_ylabel( 

1308 f"{model_names[1]}_{cube_y[0].name()} / {cube_y[0].units}", fontsize=14 

1309 ) 

1310 ax.set_title(title, fontsize=16) 

1311 ax.ticklabel_format(axis="y", useOffset=False) 

1312 ax.tick_params(axis="x", labelrotation=15) 

1313 ax.tick_params(axis="both", labelsize=12) 

1314 ax.autoscale() 

1315 

1316 # Save plot. 

1317 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1318 logging.info("Saved scatter plot to %s", filename) 

1319 plt.close(fig) 

1320 

1321 

1322def _plot_and_save_vector_plot( 

1323 cube_u: iris.cube.Cube, 

1324 cube_v: iris.cube.Cube, 

1325 filename: str, 

1326 title: str, 

1327 method: Literal["contourf", "pcolormesh"], 

1328 **kwargs, 

1329): 

1330 """Plot and save a 2D vector plot. 

1331 

1332 Parameters 

1333 ---------- 

1334 cube_u: Cube 

1335 2 dimensional Cube of u component of the data. 

1336 cube_v: Cube 

1337 2 dimensional Cube of v component of the data. 

1338 filename: str 

1339 Filename of the plot to write. 

1340 title: str 

1341 Plot title. 

1342 """ 

1343 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

1344 

1345 # Create a cube containing the magnitude of the vector field. 

1346 cube_vec_mag = (cube_u**2 + cube_v**2) ** 0.5 

1347 cube_vec_mag.rename(f"{cube_u.name()}_{cube_v.name()}_magnitude") 

1348 

1349 # Specify the color bar 

1350 cmap, levels, norm = _colorbar_map_levels(cube_vec_mag) 

1351 

1352 # Setup plot map projection, extent and coastlines and borderlines. 

1353 axes = _setup_spatial_map(cube_vec_mag, fig, cmap) 

1354 

1355 if method == "contourf": 

1356 # Filled contour plot of the field. 

1357 plot = iplt.contourf(cube_vec_mag, cmap=cmap, levels=levels, norm=norm) 

1358 elif method == "pcolormesh": 

1359 try: 

1360 vmin = min(levels) 

1361 vmax = max(levels) 

1362 except TypeError: 

1363 vmin, vmax = None, None 

1364 # pcolormesh plot of the field and ensure to use norm and not vmin/vmax 

1365 # if levels are defined. 

1366 if norm is not None: 

1367 vmin = None 

1368 vmax = None 

1369 plot = iplt.pcolormesh(cube_vec_mag, cmap=cmap, norm=norm, vmin=vmin, vmax=vmax) 

1370 else: 

1371 raise ValueError(f"Unknown plotting method: {method}") 

1372 

1373 # Check to see if transect, and if so, adjust y axis. 

1374 if is_transect(cube_vec_mag): 

1375 if "pressure" in [coord.name() for coord in cube_vec_mag.coords()]: 

1376 axes.invert_yaxis() 

1377 axes.set_yscale("log") 

1378 axes.set_ylim(1100, 100) 

1379 # If both model_level_number and level_height exists, iplt can construct 

1380 # plot as a function of height above orography (NOT sea level). 

1381 elif {"model_level_number", "level_height"}.issubset( 

1382 {coord.name() for coord in cube_vec_mag.coords()} 

1383 ): 

1384 axes.set_yscale("log") 

1385 

1386 axes.set_title( 

1387 f"{title}\n" 

1388 f"Start Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[0]}" 

1389 f" Start Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[1]}" 

1390 f" End Lat: {cube_vec_mag.attributes['transect_coords'].split('_')[2]}" 

1391 f" End Lon: {cube_vec_mag.attributes['transect_coords'].split('_')[3]}", 

1392 fontsize=16, 

1393 ) 

1394 

1395 else: 

1396 # Add title. 

1397 axes.set_title(title, fontsize=16) 

1398 

1399 # Add watermark with min/max/mean. Currently not user togglable. 

1400 # In the bbox dictionary, fc and ec are hex colour codes for grey shade. 

1401 axes.annotate( 

1402 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}", 

1403 xy=(0.05, -0.05), 

1404 xycoords="axes fraction", 

1405 xytext=(-5, 5), 

1406 textcoords="offset points", 

1407 ha="right", 

1408 va="bottom", 

1409 size=11, 

1410 bbox=dict(boxstyle="round", fc="#cccccc", ec="#808080", alpha=0.9), 

1411 ) 

1412 

1413 # Add colour bar. 

1414 cbar = fig.colorbar(plot, orientation="horizontal", pad=0.042, shrink=0.7) 

1415 cbar.set_label(label=f"{cube_vec_mag.name()} ({cube_vec_mag.units})", size=14) 

1416 # add ticks and tick_labels for every levels if less than 20 levels exist 

1417 if levels is not None and len(levels) < 20: 

1418 cbar.set_ticks(levels) 

1419 cbar.set_ticklabels([f"{level:.1f}" for level in levels]) 

1420 

1421 # 30 barbs along the longest axis of the plot, or a barb per point for data 

1422 # with less than 30 points. 

1423 step = max(max(cube_u.shape) // 30, 1) 

1424 iplt.quiver(cube_u[::step, ::step], cube_v[::step, ::step], pivot="middle") 

1425 

1426 # Save plot. 

1427 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1428 logging.info("Saved vector plot to %s", filename) 

1429 plt.close(fig) 

1430 

1431 

1432def _plot_and_save_histogram_series( 

1433 cubes: iris.cube.Cube | iris.cube.CubeList, 

1434 filename: str, 

1435 title: str, 

1436 vmin: float, 

1437 vmax: float, 

1438 **kwargs, 

1439): 

1440 """Plot and save a histogram series. 

1441 

1442 Parameters 

1443 ---------- 

1444 cubes: Cube or CubeList 

1445 2 dimensional Cube or CubeList of the data to plot as histogram. 

1446 filename: str 

1447 Filename of the plot to write. 

1448 title: str 

1449 Plot title. 

1450 vmin: float 

1451 minimum for colorbar 

1452 vmax: float 

1453 maximum for colorbar 

1454 """ 

1455 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

1456 ax = plt.gca() 

1457 

1458 model_colors_map = _get_model_colors_map(cubes) 

1459 

1460 # Set default that histograms will produce probability density function 

1461 # at each bin (integral over range sums to 1). 

1462 density = True 

1463 

1464 for cube in iter_maybe(cubes): 

1465 # Easier to check title (where var name originates) 

1466 # than seeing if long names exist etc. 

1467 # Exception case, where distribution better fits log scales/bins. 

1468 if ( 

1469 ("surface_microphysical" in title) 

1470 or ("rain accumulation" in title) 

1471 or ("Nimrod_5min" in title) 

1472 ): 

1473 if "amount" in title: 1473 ↛ 1475line 1473 didn't jump to line 1475 because the condition on line 1473 was never true

1474 # Compute histogram following Klingaman et al. (2017): ASoP 

1475 bin2 = np.exp(np.log(0.02) + 0.1 * np.linspace(0, 99, 100)) 

1476 bins = np.pad(bin2, (1, 0), "constant", constant_values=0) 

1477 density = False 

1478 else: 

1479 bins = 10.0 ** ( 

1480 np.arange(-10, 27, 1) / 10.0 

1481 ) # Suggestion from RMED toolbox. 

1482 bins = np.insert(bins, 0, 0) 

1483 ax.set_yscale("log") 

1484 vmin = bins[1] 

1485 vmax = bins[-1] # Manually set vmin/vmax to override json derived value. 

1486 ax.set_xscale("log") 

1487 elif "lightning" in title: 

1488 bins = [0, 1, 2, 3, 4, 5] 

1489 else: 

1490 bins = np.linspace(vmin, vmax, 51) 

1491 logging.debug( 

1492 "Plotting histogram with %s bins %s - %s.", 

1493 np.size(bins), 

1494 np.min(bins), 

1495 np.max(bins), 

1496 ) 

1497 

1498 # Reshape cube data into a single array to allow for a single histogram. 

1499 # Otherwise we plot xdim histograms stacked. 

1500 cube_data_1d = (cube.data).flatten() 

1501 

1502 label = None 

1503 color = "black" 

1504 if model_colors_map: 1504 ↛ 1505line 1504 didn't jump to line 1505 because the condition on line 1504 was never true

1505 label = cube.attributes.get("model_name") 

1506 color = model_colors_map[label] 

1507 x, y = np.histogram(cube_data_1d, bins=bins, density=density) 

1508 

1509 # Compute area under curve. 

1510 if ( 1510 ↛ 1515line 1510 didn't jump to line 1515 because the condition on line 1510 was never true

1511 ("surface_microphysical" in title and "amount" in title) 

1512 or ("rain_accumulation" in title) 

1513 or ("Nimrod_5min" in title) 

1514 ): 

1515 bin_mean = (bins[:-1] + bins[1:]) / 2.0 

1516 x = x * bin_mean / x.sum() 

1517 x = x[1:] 

1518 y = y[1:] 

1519 

1520 ax.plot( 

1521 y[:-1], x, color=color, linewidth=3, marker="o", markersize=6, label=label 

1522 ) 

1523 

1524 # Add some labels and tweak the style. 

1525 ax.set_title(title, fontsize=16) 

1526 ax.set_xlabel( 

1527 f"{iter_maybe(cubes)[0].name()} / {iter_maybe(cubes)[0].units}", fontsize=14 

1528 ) 

1529 ax.set_ylabel("Normalised probability density", fontsize=14) 

1530 if ( 1530 ↛ 1535line 1530 didn't jump to line 1535 because the condition on line 1530 was never true

1531 ("surface_microphysical" in title and "amount" in title) 

1532 or ("rain accumulation" in title) 

1533 or ("Nimrod_5min" in title) 

1534 ): 

1535 ax.set_ylabel( 

1536 f"Contribution to mean ({iter_maybe(cubes)[0].units})", fontsize=14 

1537 ) 

1538 ax.set_xlim(vmin, vmax) 

1539 ax.tick_params(axis="both", labelsize=12) 

1540 

1541 # Overlay grid-lines onto histogram plot. 

1542 ax.grid(linestyle="--", color="grey", linewidth=1) 

1543 if model_colors_map: 1543 ↛ 1544line 1543 didn't jump to line 1544 because the condition on line 1543 was never true

1544 ax.legend(loc="best", ncol=1, frameon=False, fontsize=16) 

1545 

1546 # Save plot. 

1547 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1548 logging.info("Saved histogram plot to %s", filename) 

1549 plt.close(fig) 

1550 

1551 

1552def _plot_and_save_postage_stamp_histogram_series( 

1553 cube: iris.cube.Cube, 

1554 filename: str, 

1555 title: str, 

1556 stamp_coordinate: str, 

1557 vmin: float, 

1558 vmax: float, 

1559 **kwargs, 

1560): 

1561 """Plot and save postage (ensemble members) stamps for a histogram series. 

1562 

1563 Parameters 

1564 ---------- 

1565 cube: Cube 

1566 2 dimensional Cube of the data to plot as histogram. 

1567 filename: str 

1568 Filename of the plot to write. 

1569 title: str 

1570 Plot title. 

1571 stamp_coordinate: str 

1572 Coordinate that becomes different plots. 

1573 vmin: float 

1574 minimum for pdf x-axis 

1575 vmax: float 

1576 maximum for pdf x-axis 

1577 """ 

1578 # Use the smallest square grid that will fit the members. 

1579 nmember = len(cube.coord(stamp_coordinate).points) 

1580 grid_rows = int(math.sqrt(nmember)) 

1581 grid_size = math.ceil(nmember / grid_rows) 

1582 

1583 fig = plt.figure( 

1584 figsize=(10, 10 * max(grid_rows / grid_size, 0.5)), facecolor="w", edgecolor="k" 

1585 ) 

1586 # Make a subplot for each member. 

1587 for member, subplot in zip( 

1588 cube.slices_over(stamp_coordinate), 

1589 range(1, grid_size * grid_rows + 1), 

1590 strict=False, 

1591 ): 

1592 # Implicit interface is much easier here, due to needing to have the 

1593 # cartopy GeoAxes generated. 

1594 plt.subplot(grid_rows, grid_size, subplot) 

1595 # Reshape cube data into a single array to allow for a single histogram. 

1596 # Otherwise we plot xdim histograms stacked. 

1597 member_data_1d = (member.data).flatten() 

1598 plt.hist(member_data_1d, density=True, stacked=True) 

1599 axes = plt.gca() 

1600 mtitle = _set_postage_stamp_title(member.coord(stamp_coordinate)) 

1601 axes.set_title(f"{mtitle}") 

1602 axes.set_xlim(vmin, vmax) 

1603 

1604 # Overall figure title. 

1605 fig.suptitle(title, fontsize=16) 

1606 

1607 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1608 logging.info("Saved histogram postage stamp plot to %s", filename) 

1609 plt.close(fig) 

1610 

1611 

1612def _plot_and_save_postage_stamps_in_single_plot_histogram_series( 

1613 cube: iris.cube.Cube, 

1614 filename: str, 

1615 title: str, 

1616 stamp_coordinate: str, 

1617 vmin: float, 

1618 vmax: float, 

1619 **kwargs, 

1620): 

1621 fig, ax = plt.subplots(figsize=(10, 10), facecolor="w", edgecolor="k") 

1622 ax.set_title(title, fontsize=16) 

1623 ax.set_xlim(vmin, vmax) 

1624 ax.set_xlabel(f"{cube.name()} / {cube.units}", fontsize=14) 

1625 ax.set_ylabel("normalised probability density", fontsize=14) 

1626 # Loop over all slices along the stamp_coordinate 

1627 for member in cube.slices_over(stamp_coordinate): 

1628 # Flatten the member data to 1D 

1629 member_data_1d = member.data.flatten() 

1630 # Plot the histogram using plt.hist 

1631 mtitle = _set_postage_stamp_title(member.coord(stamp_coordinate)) 

1632 plt.hist( 

1633 member_data_1d, 

1634 density=True, 

1635 stacked=True, 

1636 label=f"{mtitle}", 

1637 ) 

1638 

1639 # Add a legend 

1640 ax.legend(fontsize=16) 

1641 

1642 # Save the figure to a file 

1643 plt.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1644 logging.info("Saved histogram postage stamp plot to %s", filename) 

1645 

1646 # Close the figure 

1647 plt.close(fig) 

1648 

1649 

1650def _plot_and_save_scattermap_plot( 

1651 cube: iris.cube.Cube, filename: str, title: str, projection=None, **kwargs 

1652): 

1653 """Plot and save a geographical scatter plot. 

1654 

1655 Parameters 

1656 ---------- 

1657 cube: Cube 

1658 1 dimensional Cube of the data points with auxiliary latitude and 

1659 longitude coordinates, 

1660 filename: str 

1661 Filename of the plot to write. 

1662 title: str 

1663 Plot title. 

1664 projection: str 

1665 Mapping projection to be used by cartopy. 

1666 """ 

1667 # Setup plot details, size, resolution, etc. 

1668 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

1669 if projection is not None: 

1670 # Apart from the default, the only projection we currently support is 

1671 # a stereographic projection over the North Pole. 

1672 if projection == "NP_Stereo": 

1673 axes = plt.axes(projection=ccrs.NorthPolarStereo(central_longitude=0.0)) 

1674 else: 

1675 raise ValueError(f"Unknown projection: {projection}") 

1676 else: 

1677 axes = plt.axes(projection=ccrs.PlateCarree()) 

1678 

1679 # Scatter plot of the field. The marker size is chosen to give 

1680 # symbols that decrease in size as the number of observations 

1681 # increases, although the fraction of the figure covered by 

1682 # symbols increases roughly as N^(1/2), disregarding overlaps, 

1683 # and has been selected for the default figure size of (10, 10). 

1684 # Should this be changed, the marker size should be adjusted in 

1685 # proportion to the area of the figure. 

1686 mrk_size = int(np.sqrt(2500000.0 / len(cube.data))) 

1687 klon = None 

1688 klat = None 

1689 for kc in range(len(cube.aux_coords)): 

1690 if cube.aux_coords[kc].standard_name == "latitude": 

1691 klat = kc 

1692 elif cube.aux_coords[kc].standard_name == "longitude": 

1693 klon = kc 

1694 scatter_map = iplt.scatter( 

1695 cube.aux_coords[klon], 

1696 cube.aux_coords[klat], 

1697 c=cube.data[:], 

1698 s=mrk_size, 

1699 cmap="jet", 

1700 edgecolors="k", 

1701 ) 

1702 

1703 # Add coastlines and borderlines. 

1704 try: 

1705 axes.coastlines(resolution="10m") 

1706 axes.add_feature(cfeature.BORDERS) 

1707 except AttributeError: 

1708 pass 

1709 

1710 # Add title. 

1711 axes.set_title(title, fontsize=16) 

1712 

1713 # Add colour bar. 

1714 cbar = fig.colorbar(scatter_map) 

1715 cbar.set_label(label=f"{cube.name()} ({cube.units})", size=20) 

1716 

1717 # Save plot. 

1718 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1719 logging.info("Saved geographical scatter plot to %s", filename) 

1720 plt.close(fig) 

1721 

1722 

1723def _plot_and_save_power_spectrum_series( 

1724 cubes: iris.cube.Cube | iris.cube.CubeList, 

1725 filename: str, 

1726 title: str, 

1727 **kwargs, 

1728): 

1729 """Plot and save a power spectrum series. 

1730 

1731 Parameters 

1732 ---------- 

1733 cubes: Cube or CubeList 

1734 2 dimensional Cube or CubeList of the data to plot as power spectrum. 

1735 filename: str 

1736 Filename of the plot to write. 

1737 title: str 

1738 Plot title. 

1739 """ 

1740 fig = plt.figure(figsize=(10, 10), facecolor="w", edgecolor="k") 

1741 ax = plt.gca() 

1742 

1743 model_colors_map = _get_model_colors_map(cubes) 

1744 

1745 for cube in iter_maybe(cubes): 

1746 # Calculate power spectrum 

1747 

1748 # Extract time coordinate and convert to datetime 

1749 time_coord = cube.coord("time") 

1750 time_points = time_coord.units.num2date(time_coord.points) 

1751 

1752 # Choose one time point (e.g., the first one) 

1753 target_time = time_points[0] 

1754 

1755 # Bind target_time inside the lambda using a default argument 

1756 time_constraint = iris.Constraint( 

1757 time=lambda cell, target_time=target_time: cell.point == target_time 

1758 ) 

1759 

1760 cube = cube.extract(time_constraint) 

1761 

1762 if cube.ndim == 2: 

1763 cube_3d = cube.data[np.newaxis, :, :] 

1764 logging.debug("Adding in new axis for a 2 dimensional cube.") 

1765 elif cube.ndim == 3: 1765 ↛ 1766line 1765 didn't jump to line 1766 because the condition on line 1765 was never true

1766 cube_3d = cube.data 

1767 else: 

1768 raise ValueError("Cube dimensions unsuitable for power spectra code") 

1769 raise ValueError( 

1770 f"Cube is {cube.ndim} dimensional. Cube should be 2 or 3 dimensional." 

1771 ) 

1772 

1773 # Calculate spectra 

1774 ps_array = _DCT_ps(cube_3d) 

1775 

1776 ps_cube = iris.cube.Cube( 

1777 ps_array, 

1778 long_name="power_spectra", 

1779 ) 

1780 

1781 ps_cube.attributes["model_name"] = cube.attributes.get("model_name") 

1782 

1783 # Create a frequency/wavelength array for coordinate 

1784 ps_len = ps_cube.data.shape[1] 

1785 freqs = np.arange(1, ps_len + 1) 

1786 freq_coord = iris.coords.DimCoord(freqs, long_name="frequency", units="1") 

1787 

1788 # Convert datetime to numeric time using original units 

1789 numeric_time = time_coord.units.date2num(time_points) 

1790 # Create a new DimCoord with numeric time 

1791 new_time_coord = iris.coords.DimCoord( 

1792 numeric_time, standard_name="time", units=time_coord.units 

1793 ) 

1794 

1795 # Add time and frequency coordinate to spectra cube. 

1796 ps_cube.add_dim_coord(new_time_coord.copy(), 0) 

1797 ps_cube.add_dim_coord(freq_coord.copy(), 1) 

1798 

1799 # Extract data from the cube 

1800 frequency = ps_cube.coord("frequency").points 

1801 power_spectrum = ps_cube.data 

1802 

1803 label = None 

1804 color = "black" 

1805 if model_colors_map: 1805 ↛ 1806line 1805 didn't jump to line 1806 because the condition on line 1805 was never true

1806 label = ps_cube.attributes.get("model_name") 

1807 color = model_colors_map[label] 

1808 ax.plot(frequency, power_spectrum[0], color=color, label=label) 

1809 

1810 # Add some labels and tweak the style. 

1811 ax.set_title(title, fontsize=16) 

1812 ax.set_xlabel("Wavenumber", fontsize=14) 

1813 ax.set_ylabel("Power", fontsize=14) 

1814 ax.tick_params(axis="both", labelsize=12) 

1815 

1816 # Set log-log scale 

1817 ax.set_xscale("log") 

1818 ax.set_yscale("log") 

1819 

1820 # Overlay grid-lines onto power spectrum plot. 

1821 ax.grid(linestyle="--", color="grey", linewidth=1) 

1822 if model_colors_map: 1822 ↛ 1823line 1822 didn't jump to line 1823 because the condition on line 1822 was never true

1823 ax.legend(loc="best", ncol=1, frameon=False, fontsize=16) 

1824 

1825 # Save plot. 

1826 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1827 logging.info("Saved power spectrum plot to %s", filename) 

1828 plt.close(fig) 

1829 

1830 

1831def _plot_and_save_postage_stamp_power_spectrum_series( 

1832 cube: iris.cube.Cube, 

1833 filename: str, 

1834 title: str, 

1835 stamp_coordinate: str, 

1836 **kwargs, 

1837): 

1838 """Plot and save postage (ensemble members) stamps for a power spectrum series. 

1839 

1840 Parameters 

1841 ---------- 

1842 cube: Cube 

1843 2 dimensional Cube of the data to plot as power spectrum. 

1844 filename: str 

1845 Filename of the plot to write. 

1846 title: str 

1847 Plot title. 

1848 stamp_coordinate: str 

1849 Coordinate that becomes different plots. 

1850 """ 

1851 # Use the smallest square grid that will fit the members. 

1852 nmember = len(cube.coord(stamp_coordinate).points) 

1853 grid_rows = int(math.sqrt(nmember)) 

1854 grid_size = math.ceil(nmember / grid_rows) 

1855 

1856 fig = plt.figure( 

1857 figsize=(10, 10 * max(grid_rows / grid_size, 0.5)), facecolor="w", edgecolor="k" 

1858 ) 

1859 

1860 # Make a subplot for each member. 

1861 for member, subplot in zip( 

1862 cube.slices_over(stamp_coordinate), 

1863 range(1, grid_size * grid_rows + 1), 

1864 strict=False, 

1865 ): 

1866 # Implicit interface is much easier here, due to needing to have the 

1867 # cartopy GeoAxes generated. 

1868 plt.subplot(grid_rows, grid_size, subplot) 

1869 

1870 frequency = member.coord("frequency").points 

1871 

1872 axes = plt.gca() 

1873 axes.plot(frequency, member.data) 

1874 axes.set_title(f"Member #{member.coord(stamp_coordinate).points[0]}") 

1875 

1876 # Overall figure title. 

1877 fig.suptitle(title, fontsize=16) 

1878 

1879 fig.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1880 logging.info("Saved power spectra postage stamp plot to %s", filename) 

1881 plt.close(fig) 

1882 

1883 

1884def _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series( 

1885 cube: iris.cube.Cube, 

1886 filename: str, 

1887 title: str, 

1888 stamp_coordinate: str, 

1889 **kwargs, 

1890): 

1891 fig, ax = plt.subplots(figsize=(10, 10), facecolor="w", edgecolor="k") 

1892 ax.set_title(title, fontsize=16) 

1893 ax.set_xlabel(f"{cube.name()} / {cube.units}", fontsize=14) 

1894 ax.set_ylabel("Power", fontsize=14) 

1895 # Loop over all slices along the stamp_coordinate 

1896 for member in cube.slices_over(stamp_coordinate): 

1897 frequency = member.coord("frequency").points 

1898 ax.plot( 

1899 frequency, 

1900 member.data, 

1901 label=f"Member #{member.coord(stamp_coordinate).points[0]}", 

1902 ) 

1903 

1904 # Add a legend 

1905 ax.legend(fontsize=16) 

1906 

1907 # Save the figure to a file 

1908 plt.savefig(filename, bbox_inches="tight", dpi=_get_plot_resolution()) 

1909 logging.info("Saved power spectra plot to %s", filename) 

1910 

1911 # Close the figure 

1912 plt.close(fig) 

1913 

1914 

1915def _spatial_plot( 

1916 method: Literal["contourf", "pcolormesh"], 

1917 cube: iris.cube.Cube, 

1918 filename: str | None, 

1919 sequence_coordinate: str, 

1920 stamp_coordinate: str, 

1921 overlay_cube: iris.cube.Cube | None = None, 

1922 contour_cube: iris.cube.Cube | None = None, 

1923 **kwargs, 

1924): 

1925 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube. 

1926 

1927 A 2D spatial field can be plotted, but if the sequence_coordinate is present 

1928 then a sequence of plots will be produced. Similarly if the stamp_coordinate 

1929 is present then postage stamp plots will be produced. 

1930 

1931 If an overlay_cube and/or contour_cube are specified, multiple variables can 

1932 be overplotted on the same figure. 

1933 

1934 Parameters 

1935 ---------- 

1936 method: "contourf" | "pcolormesh" 

1937 The plotting method to use. 

1938 cube: Cube 

1939 Iris cube of the data to plot. It should have two spatial dimensions, 

1940 such as lat and lon, and may also have a another two dimension to be 

1941 plotted sequentially and/or as postage stamp plots. 

1942 filename: str | None 

1943 Name of the plot to write, used as a prefix for plot sequences. If None 

1944 uses the recipe name. 

1945 sequence_coordinate: str 

1946 Coordinate about which to make a plot sequence. Defaults to ``"time"``. 

1947 This coordinate must exist in the cube. 

1948 stamp_coordinate: str 

1949 Coordinate about which to plot postage stamp plots. Defaults to 

1950 ``"realization"``. 

1951 overlay_cube: Cube | None, optional 

1952 Optional 2 dimensional (lat and lon) Cube of data to overplot on top of base cube 

1953 contour_cube: Cube | None, optional 

1954 Optional 2 dimensional (lat and lon) Cube of data to overplot as contours over base cube 

1955 

1956 Raises 

1957 ------ 

1958 ValueError 

1959 If the cube doesn't have the right dimensions. 

1960 TypeError 

1961 If the cube isn't a single cube. 

1962 """ 

1963 recipe_title = get_recipe_metadata().get("title", "Untitled") 

1964 

1965 # Ensure we've got a single cube. 

1966 cube = _check_single_cube(cube) 

1967 

1968 # Check if there is a valid stamp coordinate in cube dimensions. 

1969 if stamp_coordinate == "realization": 1969 ↛ 1974line 1969 didn't jump to line 1974 because the condition on line 1969 was always true

1970 stamp_coordinate = check_stamp_coordinate(cube) 

1971 

1972 # Make postage stamp plots if stamp_coordinate exists and has more than a 

1973 # single point. 

1974 plotting_func = _plot_and_save_spatial_plot 

1975 try: 

1976 if cube.coord(stamp_coordinate).shape[0] > 1: 

1977 plotting_func = _plot_and_save_postage_stamp_spatial_plot 

1978 except iris.exceptions.CoordinateNotFoundError: 

1979 pass 

1980 

1981 # Produce a geographical scatter plot if the data have a 

1982 # dimension called observation or model_obs_error 

1983 if any( 1983 ↛ 1987line 1983 didn't jump to line 1987 because the condition on line 1983 was never true

1984 crd.var_name == "station" or crd.var_name == "model_obs_error" 

1985 for crd in cube.coords() 

1986 ): 

1987 plotting_func = _plot_and_save_scattermap_plot 

1988 

1989 # Must have a sequence coordinate. 

1990 try: 

1991 cube.coord(sequence_coordinate) 

1992 except iris.exceptions.CoordinateNotFoundError as err: 

1993 raise ValueError(f"Cube must have a {sequence_coordinate} coordinate.") from err 

1994 

1995 # Create a plot for each value of the sequence coordinate. 

1996 plot_index = [] 

1997 nplot = np.size(cube.coord(sequence_coordinate).points) 

1998 

1999 for iseq, cube_slice in enumerate(cube.slices_over(sequence_coordinate)): 

2000 # Set plot titles and filename 

2001 seq_coord = cube_slice.coord(sequence_coordinate) 

2002 plot_title, plot_filename = _set_title_and_filename( 

2003 seq_coord, nplot, recipe_title, filename 

2004 ) 

2005 

2006 # Extract sequence slice for overlay_cube and contour_cube if required. 

2007 overlay_slice = slice_over_maybe(overlay_cube, sequence_coordinate, iseq) 

2008 contour_slice = slice_over_maybe(contour_cube, sequence_coordinate, iseq) 

2009 

2010 # Do the actual plotting. 

2011 plotting_func( 

2012 cube_slice, 

2013 filename=plot_filename, 

2014 stamp_coordinate=stamp_coordinate, 

2015 title=plot_title, 

2016 method=method, 

2017 overlay_cube=overlay_slice, 

2018 contour_cube=contour_slice, 

2019 **kwargs, 

2020 ) 

2021 plot_index.append(plot_filename) 

2022 

2023 # Add list of plots to plot metadata. 

2024 complete_plot_index = _append_to_plot_index(plot_index) 

2025 

2026 # Make a page to display the plots. 

2027 _make_plot_html_page(complete_plot_index) 

2028 

2029 

2030def _custom_colormap_mask(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None): 

2031 """Get colourmap for mask. 

2032 

2033 If "mask_for_" appears anywhere in the name of a cube this function will be called 

2034 regardless of the name of the variable to ensure a consistent plot. 

2035 

2036 Parameters 

2037 ---------- 

2038 cube: Cube 

2039 Cube of variable for which the colorbar information is desired. 

2040 axis: "x", "y", optional 

2041 Select the levels for just this axis of a line plot. The min and max 

2042 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

2043 setting a universal range is not desirable (e.g. temperature), users 

2044 can set ymin/ymax values to "auto" in the colorbar definitions file. 

2045 Where no additional xmin/xmax or ymin/ymax values are provided, the 

2046 axis bounds default to use the vmin/vmax values provided. 

2047 

2048 Returns 

2049 ------- 

2050 cmap: 

2051 Matplotlib colormap. 

2052 levels: 

2053 List of levels to use for plotting. For continuous plots the min and max 

2054 should be taken as the range. 

2055 norm: 

2056 BoundaryNorm information. 

2057 """ 

2058 if "difference" not in cube.long_name: 

2059 if axis: 

2060 levels = [0, 1] 

2061 # Complete settings based on levels. 

2062 return None, levels, None 

2063 else: 

2064 # Define the levels and colors. 

2065 levels = [0, 1, 2] 

2066 colors = ["white", "dodgerblue"] 

2067 # Create a custom color map. 

2068 cmap = mcolors.ListedColormap(colors) 

2069 # Normalize the levels. 

2070 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2071 logging.debug("Colourmap for %s.", cube.long_name) 

2072 return cmap, levels, norm 

2073 else: 

2074 if axis: 

2075 levels = [-1, 1] 

2076 return None, levels, None 

2077 else: 

2078 # Search for if mask difference, set to +/- 0.5 as values plotted < 

2079 # not <=. 

2080 levels = [-2, -0.5, 0.5, 2] 

2081 colors = ["goldenrod", "white", "teal"] 

2082 cmap = mcolors.ListedColormap(colors) 

2083 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2084 logging.debug("Colourmap for %s.", cube.long_name) 

2085 return cmap, levels, norm 

2086 

2087 

2088def _custom_beaufort_scale(cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None): 

2089 """Get a custom colorbar for a cube in the Beaufort Scale. 

2090 

2091 Specific variable ranges can be separately set in user-supplied definition 

2092 for x- or y-axis limits, or indicate where automated range preferred. 

2093 

2094 Parameters 

2095 ---------- 

2096 cube: Cube 

2097 Cube of variable with Beaufort Scale in name. 

2098 axis: "x", "y", optional 

2099 Select the levels for just this axis of a line plot. The min and max 

2100 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

2101 setting a universal range is not desirable (e.g. temperature), users 

2102 can set ymin/ymax values to "auto" in the colorbar definitions file. 

2103 Where no additional xmin/xmax or ymin/ymax values are provided, the 

2104 axis bounds default to use the vmin/vmax values provided. 

2105 

2106 Returns 

2107 ------- 

2108 cmap: 

2109 Matplotlib colormap. 

2110 levels: 

2111 List of levels to use for plotting. For continuous plots the min and max 

2112 should be taken as the range. 

2113 norm: 

2114 BoundaryNorm information. 

2115 """ 

2116 if "difference" not in cube.long_name: 

2117 if axis: 

2118 levels = [0, 12] 

2119 return None, levels, None 

2120 else: 

2121 levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13] 

2122 colors = [ 

2123 "black", 

2124 (0, 0, 0.6), 

2125 "blue", 

2126 "cyan", 

2127 "green", 

2128 "yellow", 

2129 (1, 0.5, 0), 

2130 "red", 

2131 "pink", 

2132 "magenta", 

2133 "purple", 

2134 "maroon", 

2135 "white", 

2136 ] 

2137 cmap = mcolors.ListedColormap(colors) 

2138 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2139 logging.info("change colormap for Beaufort Scale colorbar.") 

2140 return cmap, levels, norm 

2141 else: 

2142 if axis: 

2143 levels = [-4, 4] 

2144 return None, levels, None 

2145 else: 

2146 levels = [ 

2147 -3.5, 

2148 -2.5, 

2149 -1.5, 

2150 -0.5, 

2151 0.5, 

2152 1.5, 

2153 2.5, 

2154 3.5, 

2155 ] 

2156 cmap = plt.get_cmap("bwr", 8) 

2157 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2158 return cmap, levels, norm 

2159 

2160 

2161def _custom_colormap_celsius(cube: iris.cube.Cube, cmap, levels, norm): 

2162 """Return altered colourmap for temperature with change in units to Celsius. 

2163 

2164 If "Celsius" appears anywhere in the name of a cube this function will be called. 

2165 

2166 Parameters 

2167 ---------- 

2168 cube: Cube 

2169 Cube of variable for which the colorbar information is desired. 

2170 cmap: Matplotlib colormap. 

2171 levels: List 

2172 List of levels to use for plotting. For continuous plots the min and max 

2173 should be taken as the range. 

2174 norm: BoundaryNorm. 

2175 

2176 Returns 

2177 ------- 

2178 cmap: Matplotlib colormap. 

2179 levels: List 

2180 List of levels to use for plotting. For continuous plots the min and max 

2181 should be taken as the range. 

2182 norm: BoundaryNorm. 

2183 """ 

2184 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name]) 

2185 if any("temperature" in name for name in varnames) and "Celsius" == cube.units: 

2186 levels = np.array(levels) 

2187 levels -= 273 

2188 levels = levels.tolist() 

2189 else: 

2190 # Do nothing keep the existing colourbar attributes 

2191 levels = levels 

2192 cmap = cmap 

2193 norm = norm 

2194 return cmap, levels, norm 

2195 

2196 

2197def _custom_colormap_probability( 

2198 cube: iris.cube.Cube, axis: Literal["x", "y"] | None = None 

2199): 

2200 """Get a custom colorbar for a probability cube. 

2201 

2202 Specific variable ranges can be separately set in user-supplied definition 

2203 for x- or y-axis limits, or indicate where automated range preferred. 

2204 

2205 Parameters 

2206 ---------- 

2207 cube: Cube 

2208 Cube of variable with probability in name. 

2209 axis: "x", "y", optional 

2210 Select the levels for just this axis of a line plot. The min and max 

2211 can be set by xmin/xmax or ymin/ymax respectively. For variables where 

2212 setting a universal range is not desirable (e.g. temperature), users 

2213 can set ymin/ymax values to "auto" in the colorbar definitions file. 

2214 Where no additional xmin/xmax or ymin/ymax values are provided, the 

2215 axis bounds default to use the vmin/vmax values provided. 

2216 

2217 Returns 

2218 ------- 

2219 cmap: 

2220 Matplotlib colormap. 

2221 levels: 

2222 List of levels to use for plotting. For continuous plots the min and max 

2223 should be taken as the range. 

2224 norm: 

2225 BoundaryNorm information. 

2226 """ 

2227 if axis: 

2228 levels = [0, 1] 

2229 return None, levels, None 

2230 else: 

2231 cmap = mcolors.ListedColormap( 

2232 [ 

2233 "#FFFFFF", 

2234 "#636363", 

2235 "#e1dada", 

2236 "#B5CAFF", 

2237 "#8FB3FF", 

2238 "#7F97FF", 

2239 "#ABCF63", 

2240 "#E8F59E", 

2241 "#FFFA14", 

2242 "#FFD121", 

2243 "#FFA30A", 

2244 ] 

2245 ) 

2246 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] 

2247 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2248 return cmap, levels, norm 

2249 

2250 

2251def _custom_colourmap_precipitation(cube: iris.cube.Cube, cmap, levels, norm): 

2252 """Return a custom colourmap for the current recipe.""" 

2253 varnames_lower = [ 

2254 n.lower() for n in (cube.long_name, cube.standard_name, cube.var_name) if n 

2255 ] 

2256 

2257 is_rainfall_var = any( 

2258 key in name 

2259 for name in varnames_lower 

2260 for key in ( 

2261 "surface_microphysical", 

2262 "rainfall rate composite", 

2263 "nimrod5min", 

2264 "nimrod_5min", 

2265 "rain_accumulation", 

2266 "rain accumulation", 

2267 ) 

2268 ) 

2269 

2270 if is_rainfall_var: 

2271 print("varnames_lower ", varnames_lower) 

2272 levels = [0, 0.125, 0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128, 256] 

2273 colors = [ 

2274 "w", 

2275 (0, 0, 0.6), 

2276 "b", 

2277 "c", 

2278 "g", 

2279 "y", 

2280 (1, 0.5, 0), 

2281 "r", 

2282 "pink", 

2283 "m", 

2284 "purple", 

2285 "maroon", 

2286 "gray", 

2287 ] 

2288 # Create a custom colormap 

2289 cmap = mcolors.ListedColormap(colors) 

2290 # Normalize the levels 

2291 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2292 logging.info("Using custom rainfall colourmap.") 

2293 

2294 return cmap, levels, norm 

2295 

2296 

2297def _custom_colormap_aviation_colour_state(cube: iris.cube.Cube): 

2298 """Return custom colourmap for aviation colour state. 

2299 

2300 If "aviation_colour_state" appears anywhere in the name of a cube 

2301 this function will be called. 

2302 

2303 Parameters 

2304 ---------- 

2305 cube: Cube 

2306 Cube of variable for which the colorbar information is desired. 

2307 

2308 Returns 

2309 ------- 

2310 cmap: Matplotlib colormap. 

2311 levels: List 

2312 List of levels to use for plotting. For continuous plots the min and max 

2313 should be taken as the range. 

2314 norm: BoundaryNorm. 

2315 """ 

2316 levels = [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5, 6.5] 

2317 colors = [ 

2318 "#87ceeb", 

2319 "#ffffff", 

2320 "#8ced69", 

2321 "#ffff00", 

2322 "#ffd700", 

2323 "#ffa500", 

2324 "#fe3620", 

2325 ] 

2326 # Create a custom colormap 

2327 cmap = mcolors.ListedColormap(colors) 

2328 # Normalise the levels 

2329 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2330 return cmap, levels, norm 

2331 

2332 

2333def _custom_colourmap_nimrod_weights(cube: iris.cube.Cube, cmap, levels, norm): 

2334 """Return a custom colourmap for the current recipe.""" 

2335 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name]) 

2336 if ( 2336 ↛ 2343line 2336 didn't jump to line 2343 because the condition on line 2336 was never true

2337 any("wts" in name for name in varnames) 

2338 and "difference" not in cube.long_name 

2339 and "mask" not in cube.long_name 

2340 ): 

2341 # Define the levels and colors. Remember the Nimrod weights vary over the 

2342 # range [0,13] and should be integer values. Optimum value is 13. 

2343 levels = [ 

2344 -0.5, 

2345 0.5, 

2346 1.5, 

2347 2.5, 

2348 3.5, 

2349 4.5, 

2350 5.5, 

2351 6.5, 

2352 7.5, 

2353 8.5, 

2354 9.5, 

2355 10.5, 

2356 11.5, 

2357 12.5, 

2358 13.5, 

2359 ] 

2360 # levels = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14] 

2361 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2362 colours = [ 

2363 "#d10000", 

2364 "purple", 

2365 "#8f00d6", 

2366 "#ff9700", 

2367 "pink", 

2368 "#ffff00", 

2369 "#00007f", 

2370 "#6c9ccd", 

2371 "#aae8ff", 

2372 "#37a648", 

2373 "#8edc64", 

2374 "#c5ffc5", 

2375 "#dcdcdc", 

2376 "#ffffff", 

2377 ] 

2378 # Create a custom colormap 

2379 cmap = mcolors.ListedColormap(colours) 

2380 # Normalize the levels 

2381 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2382 logging.info("Change colormap for Nimrod weights colorbar.") 

2383 else: 

2384 # do nothing and keep existing colorbar attributes 

2385 cmap = cmap 

2386 levels = levels 

2387 norm = norm 

2388 return cmap, levels, norm 

2389 

2390 

2391def _custom_colourmap_visibility_in_air(cube: iris.cube.Cube, cmap, levels, norm): 

2392 """Return a custom colourmap for the current recipe.""" 

2393 varnames = filter(None, [cube.long_name, cube.standard_name, cube.var_name]) 

2394 if ( 

2395 any("visibility_in_air" in name for name in varnames) 

2396 and "difference" not in cube.long_name 

2397 and "mask" not in cube.long_name 

2398 ): 

2399 # Define the levels and colors (in km) 

2400 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] 

2401 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2402 colours = [ 

2403 "#8f00d6", 

2404 "#d10000", 

2405 "#ff9700", 

2406 "#ffff00", 

2407 "#00007f", 

2408 "#6c9ccd", 

2409 "#aae8ff", 

2410 "#37a648", 

2411 "#8edc64", 

2412 "#c5ffc5", 

2413 "#dcdcdc", 

2414 "#ffffff", 

2415 ] 

2416 # Create a custom colormap 

2417 cmap = mcolors.ListedColormap(colours) 

2418 # Normalize the levels 

2419 norm = mcolors.BoundaryNorm(levels, cmap.N) 

2420 logging.info("change colormap for visibility_in_air variable colorbar.") 

2421 else: 

2422 # do nothing and keep existing colorbar attributes 

2423 cmap = cmap 

2424 levels = levels 

2425 norm = norm 

2426 return cmap, levels, norm 

2427 

2428 

2429def _get_num_models(cube: iris.cube.Cube | iris.cube.CubeList) -> int: 

2430 """Return number of models based on cube attributes.""" 

2431 model_names = list( 

2432 filter( 

2433 lambda x: x is not None, 

2434 {cb.attributes.get("model_name", None) for cb in iter_maybe(cube)}, 

2435 ) 

2436 ) 

2437 if not model_names: 

2438 logging.debug("Missing model names. Will assume single model.") 

2439 return 1 

2440 else: 

2441 return len(model_names) 

2442 

2443 

2444def _validate_cube_shape( 

2445 cube: iris.cube.Cube | iris.cube.CubeList, num_models: int 

2446) -> None: 

2447 """Check all cubes have a model name.""" 

2448 if isinstance(cube, iris.cube.CubeList) and len(cube) != num_models: 2448 ↛ 2449line 2448 didn't jump to line 2449 because the condition on line 2448 was never true

2449 raise ValueError( 

2450 f"The number of model names ({num_models}) should equal the number " 

2451 f"of cubes ({len(cube)})." 

2452 ) 

2453 

2454 

2455def _validate_cubes_coords( 

2456 cubes: iris.cube.CubeList, coords: list[iris.coords.Coord] 

2457) -> None: 

2458 """Check same number of cubes as sequence coordinate for zip functions.""" 

2459 if len(cubes) != len(coords): 2459 ↛ 2460line 2459 didn't jump to line 2460 because the condition on line 2459 was never true

2460 raise ValueError( 

2461 f"The number of CubeList entries ({len(cubes)}) should equal the number " 

2462 f"of sequence coordinates ({len(coords)})." 

2463 f"Check that number of time entries in input data are consistent if " 

2464 f"performing time-averaging steps prior to plotting outputs." 

2465 ) 

2466 

2467 

2468#################### 

2469# Public functions # 

2470#################### 

2471 

2472 

2473def spatial_contour_plot( 

2474 cube: iris.cube.Cube, 

2475 filename: str = None, 

2476 sequence_coordinate: str = "time", 

2477 stamp_coordinate: str = "realization", 

2478 **kwargs, 

2479) -> iris.cube.Cube: 

2480 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube. 

2481 

2482 A 2D spatial field can be plotted, but if the sequence_coordinate is present 

2483 then a sequence of plots will be produced. Similarly if the stamp_coordinate 

2484 is present then postage stamp plots will be produced. 

2485 

2486 Parameters 

2487 ---------- 

2488 cube: Cube 

2489 Iris cube of the data to plot. It should have two spatial dimensions, 

2490 such as lat and lon, and may also have a another two dimension to be 

2491 plotted sequentially and/or as postage stamp plots. 

2492 filename: str, optional 

2493 Name of the plot to write, used as a prefix for plot sequences. Defaults 

2494 to the recipe name. 

2495 sequence_coordinate: str, optional 

2496 Coordinate about which to make a plot sequence. Defaults to ``"time"``. 

2497 This coordinate must exist in the cube. 

2498 stamp_coordinate: str, optional 

2499 Coordinate about which to plot postage stamp plots. Defaults to 

2500 ``"realization"``. 

2501 

2502 Returns 

2503 ------- 

2504 Cube 

2505 The original cube (so further operations can be applied). 

2506 

2507 Raises 

2508 ------ 

2509 ValueError 

2510 If the cube doesn't have the right dimensions. 

2511 TypeError 

2512 If the cube isn't a single cube. 

2513 """ 

2514 _spatial_plot( 

2515 "contourf", cube, filename, sequence_coordinate, stamp_coordinate, **kwargs 

2516 ) 

2517 return cube 

2518 

2519 

2520def spatial_pcolormesh_plot( 

2521 cube: iris.cube.Cube, 

2522 filename: str = None, 

2523 sequence_coordinate: str = "time", 

2524 stamp_coordinate: str = "realization", 

2525 **kwargs, 

2526) -> iris.cube.Cube: 

2527 """Plot a spatial variable onto a map from a 2D, 3D, or 4D cube. 

2528 

2529 A 2D spatial field can be plotted, but if the sequence_coordinate is present 

2530 then a sequence of plots will be produced. Similarly if the stamp_coordinate 

2531 is present then postage stamp plots will be produced. 

2532 

2533 This function is significantly faster than ``spatial_contour_plot``, 

2534 especially at high resolutions, and should be preferred unless contiguous 

2535 contour areas are important. 

2536 

2537 Parameters 

2538 ---------- 

2539 cube: Cube 

2540 Iris cube of the data to plot. It should have two spatial dimensions, 

2541 such as lat and lon, and may also have a another two dimension to be 

2542 plotted sequentially and/or as postage stamp plots. 

2543 filename: str, optional 

2544 Name of the plot to write, used as a prefix for plot sequences. Defaults 

2545 to the recipe name. 

2546 sequence_coordinate: str, optional 

2547 Coordinate about which to make a plot sequence. Defaults to ``"time"``. 

2548 This coordinate must exist in the cube. 

2549 stamp_coordinate: str, optional 

2550 Coordinate about which to plot postage stamp plots. Defaults to 

2551 ``"realization"``. 

2552 

2553 Returns 

2554 ------- 

2555 Cube 

2556 The original cube (so further operations can be applied). 

2557 

2558 Raises 

2559 ------ 

2560 ValueError 

2561 If the cube doesn't have the right dimensions. 

2562 TypeError 

2563 If the cube isn't a single cube. 

2564 """ 

2565 _spatial_plot( 

2566 "pcolormesh", cube, filename, sequence_coordinate, stamp_coordinate, **kwargs 

2567 ) 

2568 return cube 

2569 

2570 

2571def spatial_multi_pcolormesh_plot( 

2572 cube: iris.cube.Cube, 

2573 overlay_cube: iris.cube.Cube, 

2574 contour_cube: iris.cube.Cube, 

2575 filename: str = None, 

2576 sequence_coordinate: str = "time", 

2577 stamp_coordinate: str = "realization", 

2578 **kwargs, 

2579) -> iris.cube.Cube: 

2580 """Plot a set of spatial variables onto a map from a 2D, 3D, or 4D cube. 

2581 

2582 A 2D basis cube spatial field can be plotted, but if the sequence_coordinate is present 

2583 then a sequence of plots will be produced. Similarly if the stamp_coordinate 

2584 is present then postage stamp plots will be produced. 

2585 

2586 If specified, a masked overlay_cube can be overplotted on top of the base cube. 

2587 

2588 If specified, contours of a contour_cube can be overplotted on top of those. 

2589 

2590 For single-variable equivalent of this routine, use spatial_pcolormesh_plot. 

2591 

2592 This function is significantly faster than ``spatial_contour_plot``, 

2593 especially at high resolutions, and should be preferred unless contiguous 

2594 contour areas are important. 

2595 

2596 Parameters 

2597 ---------- 

2598 cube: Cube 

2599 Iris cube of the data to plot. It should have two spatial dimensions, 

2600 such as lat and lon, and may also have a another two dimension to be 

2601 plotted sequentially and/or as postage stamp plots. 

2602 overlay_cube: Cube 

2603 Iris cube of the data to plot as an overlay on top of basis cube. It should have two spatial dimensions, 

2604 such as lat and lon, and may also have a another two dimension to be 

2605 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. 

2606 contour_cube: Cube 

2607 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, 

2608 such as lat and lon, and may also have a another two dimension to be 

2609 plotted sequentially and/or as postage stamp plots. 

2610 filename: str, optional 

2611 Name of the plot to write, used as a prefix for plot sequences. Defaults 

2612 to the recipe name. 

2613 sequence_coordinate: str, optional 

2614 Coordinate about which to make a plot sequence. Defaults to ``"time"``. 

2615 This coordinate must exist in the cube. 

2616 stamp_coordinate: str, optional 

2617 Coordinate about which to plot postage stamp plots. Defaults to 

2618 ``"realization"``. 

2619 

2620 Returns 

2621 ------- 

2622 Cube 

2623 The original cube (so further operations can be applied). 

2624 

2625 Raises 

2626 ------ 

2627 ValueError 

2628 If the cube doesn't have the right dimensions. 

2629 TypeError 

2630 If the cube isn't a single cube. 

2631 """ 

2632 _spatial_plot( 

2633 "pcolormesh", 

2634 cube, 

2635 filename, 

2636 sequence_coordinate, 

2637 stamp_coordinate, 

2638 overlay_cube=overlay_cube, 

2639 contour_cube=contour_cube, 

2640 ) 

2641 return cube, overlay_cube, contour_cube 

2642 

2643 

2644# TODO: Expand function to handle ensemble data. 

2645# line_coordinate: str, optional 

2646# Coordinate about which to plot multiple lines. Defaults to 

2647# ``"realization"``. 

2648def plot_line_series( 

2649 cube: iris.cube.Cube | iris.cube.CubeList, 

2650 filename: str = None, 

2651 series_coordinate: str = "time", 

2652 # line_coordinate: str = "realization", 

2653 **kwargs, 

2654) -> iris.cube.Cube | iris.cube.CubeList: 

2655 """Plot a line plot for the specified coordinate. 

2656 

2657 The Cube or CubeList must be 1D. 

2658 

2659 Parameters 

2660 ---------- 

2661 iris.cube | iris.cube.CubeList 

2662 Cube or CubeList of the data to plot. The individual cubes should have a single dimension. 

2663 The cubes should cover the same phenomenon i.e. all cubes contain temperature data. 

2664 We do not support different data such as temperature and humidity in the same CubeList for plotting. 

2665 filename: str, optional 

2666 Name of the plot to write, used as a prefix for plot sequences. Defaults 

2667 to the recipe name. 

2668 series_coordinate: str, optional 

2669 Coordinate about which to make a series. Defaults to ``"time"``. This 

2670 coordinate must exist in the cube. 

2671 

2672 Returns 

2673 ------- 

2674 iris.cube.Cube | iris.cube.CubeList 

2675 The original Cube or CubeList (so further operations can be applied). 

2676 Plotted data. 

2677 

2678 Raises 

2679 ------ 

2680 ValueError 

2681 If the cubes don't have the right dimensions. 

2682 TypeError 

2683 If the cube isn't a Cube or CubeList. 

2684 """ 

2685 # Ensure we have a name for the plot file. 

2686 recipe_title = get_recipe_metadata().get("title", "Untitled") 

2687 

2688 num_models = _get_num_models(cube) 

2689 

2690 _validate_cube_shape(cube, num_models) 

2691 

2692 # Iterate over all cubes and extract coordinate to plot. 

2693 cubes = iter_maybe(cube) 

2694 coords = [] 

2695 for cube in cubes: 

2696 try: 

2697 coords.append(cube.coord(series_coordinate)) 

2698 except iris.exceptions.CoordinateNotFoundError as err: 

2699 raise ValueError( 

2700 f"Cube must have a {series_coordinate} coordinate." 

2701 ) from err 

2702 if cube.ndim > 2 or not cube.coords("realization"): 

2703 raise ValueError("Cube must be 1D or 2D with a realization coordinate.") 

2704 

2705 # Format the title and filename using plotted series coordinate 

2706 nplot = 1 

2707 seq_coord = coords[0] 

2708 plot_title, plot_filename = _set_title_and_filename( 

2709 seq_coord, nplot, recipe_title, filename 

2710 ) 

2711 

2712 # Do the actual plotting. 

2713 _plot_and_save_line_series(cubes, coords, "realization", plot_filename, plot_title) 

2714 

2715 # Add list of plots to plot metadata. 

2716 plot_index = _append_to_plot_index([plot_filename]) 

2717 

2718 # Make a page to display the plots. 

2719 _make_plot_html_page(plot_index) 

2720 

2721 return cube 

2722 

2723 

2724def plot_vertical_line_series( 

2725 cubes: iris.cube.Cube | iris.cube.CubeList, 

2726 filename: str = None, 

2727 series_coordinate: str = "model_level_number", 

2728 sequence_coordinate: str = "time", 

2729 # line_coordinate: str = "realization", 

2730 **kwargs, 

2731) -> iris.cube.Cube | iris.cube.CubeList: 

2732 """Plot a line plot against a type of vertical coordinate. 

2733 

2734 The Cube or CubeList must be 1D. 

2735 

2736 A 1D line plot with y-axis as pressure coordinate can be plotted, but if the sequence_coordinate is present 

2737 then a sequence of plots will be produced. 

2738 

2739 Parameters 

2740 ---------- 

2741 iris.cube | iris.cube.CubeList 

2742 Cube or CubeList of the data to plot. The individual cubes should have a single dimension. 

2743 The cubes should cover the same phenomenon i.e. all cubes contain temperature data. 

2744 We do not support different data such as temperature and humidity in the same CubeList for plotting. 

2745 filename: str, optional 

2746 Name of the plot to write, used as a prefix for plot sequences. Defaults 

2747 to the recipe name. 

2748 series_coordinate: str, optional 

2749 Coordinate to plot on the y-axis. Can be ``pressure`` or 

2750 ``model_level_number`` for UM, or ``full_levels`` or ``half_levels`` 

2751 for LFRic. Defaults to ``model_level_number``. 

2752 This coordinate must exist in the cube. 

2753 sequence_coordinate: str, optional 

2754 Coordinate about which to make a plot sequence. Defaults to ``"time"``. 

2755 This coordinate must exist in the cube. 

2756 

2757 Returns 

2758 ------- 

2759 iris.cube.Cube | iris.cube.CubeList 

2760 The original Cube or CubeList (so further operations can be applied). 

2761 Plotted data. 

2762 

2763 Raises 

2764 ------ 

2765 ValueError 

2766 If the cubes doesn't have the right dimensions. 

2767 TypeError 

2768 If the cube isn't a Cube or CubeList. 

2769 """ 

2770 # Ensure we have a name for the plot file. 

2771 recipe_title = get_recipe_metadata().get("title", "Untitled") 

2772 

2773 cubes = iter_maybe(cubes) 

2774 # Initialise empty list to hold all data from all cubes in a CubeList 

2775 all_data = [] 

2776 

2777 # Store min/max ranges for x range. 

2778 x_levels = [] 

2779 

2780 num_models = _get_num_models(cubes) 

2781 

2782 _validate_cube_shape(cubes, num_models) 

2783 

2784 # Iterate over all cubes in cube or CubeList and plot. 

2785 coords = [] 

2786 for cube in cubes: 

2787 # Test if series coordinate i.e. pressure level exist for any cube with cube.ndim >=1. 

2788 try: 

2789 coords.append(cube.coord(series_coordinate)) 

2790 except iris.exceptions.CoordinateNotFoundError as err: 

2791 raise ValueError( 

2792 f"Cube must have a {series_coordinate} coordinate." 

2793 ) from err 

2794 

2795 try: 

2796 if cube.ndim > 1 or not cube.coords("realization"): 2796 ↛ 2804line 2796 didn't jump to line 2804 because the condition on line 2796 was always true

2797 cube.coord(sequence_coordinate) 

2798 except iris.exceptions.CoordinateNotFoundError as err: 

2799 raise ValueError( 

2800 f"Cube must have a {sequence_coordinate} coordinate or be 1D, or 2D with a realization coordinate." 

2801 ) from err 

2802 

2803 # Get minimum and maximum from levels information. 

2804 _, levels, _ = _colorbar_map_levels(cube, axis="x") 

2805 if levels is not None: 2805 ↛ 2809line 2805 didn't jump to line 2809 because the condition on line 2805 was always true

2806 x_levels.append(min(levels)) 

2807 x_levels.append(max(levels)) 

2808 else: 

2809 all_data.append(cube.data) 

2810 

2811 if len(x_levels) == 0: 2811 ↛ 2813line 2811 didn't jump to line 2813 because the condition on line 2811 was never true

2812 # Combine all data into a single NumPy array 

2813 combined_data = np.concatenate(all_data) 

2814 

2815 # Set the lower and upper limit for the x-axis to ensure all plots have 

2816 # same range. This needs to read the whole cube over the range of the 

2817 # sequence and if applicable postage stamp coordinate. 

2818 vmin = np.floor(combined_data.min()) 

2819 vmax = np.ceil(combined_data.max()) 

2820 else: 

2821 vmin = min(x_levels) 

2822 vmax = max(x_levels) 

2823 

2824 # Matching the slices (matching by seq coord point; it may happen that 

2825 # evaluated models do not cover the same seq coord range, hence matching 

2826 # necessary) 

2827 def filter_cube_iterables(cube_iterables) -> bool: 

2828 return len(cube_iterables) == len(coords) 

2829 

2830 cube_iterables = filter( 

2831 filter_cube_iterables, 

2832 ( 

2833 iris.cube.CubeList( 

2834 s 

2835 for s in itertools.chain.from_iterable( 

2836 cb.slices_over(sequence_coordinate) for cb in cubes 

2837 ) 

2838 if s.coord(sequence_coordinate).points[0] == point 

2839 ) 

2840 for point in sorted( 

2841 set( 

2842 itertools.chain.from_iterable( 

2843 cb.coord(sequence_coordinate).points for cb in cubes 

2844 ) 

2845 ) 

2846 ) 

2847 ), 

2848 ) 

2849 

2850 # Create a plot for each value of the sequence coordinate. 

2851 # Allowing for multiple cubes in a CubeList to be plotted in the same plot for 

2852 # similar sequence values. Passing a CubeList into the internal plotting function 

2853 # for similar values of the sequence coordinate. cube_slice can be an iris.cube.Cube 

2854 # or an iris.cube.CubeList. 

2855 plot_index = [] 

2856 nplot = np.size(cubes[0].coord(sequence_coordinate).points) 

2857 for cubes_slice in cube_iterables: 

2858 # Format the coordinate value in a unit appropriate way. 

2859 seq_coord = cubes_slice[0].coord(sequence_coordinate) 

2860 plot_title, plot_filename = _set_title_and_filename( 

2861 seq_coord, nplot, recipe_title, filename 

2862 ) 

2863 

2864 # Do the actual plotting. 

2865 _plot_and_save_vertical_line_series( 

2866 cubes_slice, 

2867 coords, 

2868 "realization", 

2869 plot_filename, 

2870 series_coordinate, 

2871 title=plot_title, 

2872 vmin=vmin, 

2873 vmax=vmax, 

2874 ) 

2875 plot_index.append(plot_filename) 

2876 

2877 # Add list of plots to plot metadata. 

2878 complete_plot_index = _append_to_plot_index(plot_index) 

2879 

2880 # Make a page to display the plots. 

2881 _make_plot_html_page(complete_plot_index) 

2882 

2883 return cubes 

2884 

2885 

2886def qq_plot( 

2887 cubes: iris.cube.CubeList, 

2888 coordinates: list[str], 

2889 percentiles: list[float], 

2890 model_names: list[str], 

2891 filename: str = None, 

2892 one_to_one: bool = True, 

2893 **kwargs, 

2894) -> iris.cube.CubeList: 

2895 """Plot a Quantile-Quantile plot between two models for common time points. 

2896 

2897 The cubes will be normalised by collapsing each cube to its percentiles. Cubes are 

2898 collapsed within the operator over all specified coordinates such as 

2899 grid_latitude, grid_longitude, vertical levels, but also realisation representing 

2900 ensemble members to ensure a 1D cube (array). 

2901 

2902 Parameters 

2903 ---------- 

2904 cubes: iris.cube.CubeList 

2905 Two cubes of the same variable with different models. 

2906 coordinate: list[str] 

2907 The list of coordinates to collapse over. This list should be 

2908 every coordinate within the cube to result in a 1D cube around 

2909 the percentile coordinate. 

2910 percent: list[float] 

2911 A list of percentiles to appear in the plot. 

2912 model_names: list[str] 

2913 A list of model names to appear on the axis of the plot. 

2914 filename: str, optional 

2915 Filename of the plot to write. 

2916 one_to_one: bool, optional 

2917 If True a 1:1 line is plotted; if False it is not. Default is True. 

2918 

2919 Raises 

2920 ------ 

2921 ValueError 

2922 When the cubes are not compatible. 

2923 

2924 Notes 

2925 ----- 

2926 The quantile-quantile plot is a variant on the scatter plot representing 

2927 two datasets by their quantiles (percentiles) for common time points. 

2928 This plot does not use a theoretical distribution to compare against, but 

2929 compares percentiles of two datasets. This plot does 

2930 not use all raw data points, but plots the selected percentiles (quantiles) of 

2931 each variable instead for the two datasets, thereby normalising the data for a 

2932 direct comparison between the selected percentiles of the two dataset distributions. 

2933 

2934 Quantile-quantile plots are valuable for comparing against 

2935 observations and other models. Identical percentiles between the variables 

2936 will lie on the one-to-one line implying the values correspond well to each 

2937 other. Where there is a deviation from the one-to-one line a range of 

2938 possibilities exist depending on how and where the data is shifted (e.g., 

2939 Wilks 2011 [Wilks2011]_). 

2940 

2941 For distributions above the one-to-one line the distribution is left-skewed; 

2942 below is right-skewed. A distinct break implies a bimodal distribution, and 

2943 closer values/values further apart at the tails imply poor representation of 

2944 the extremes. 

2945 

2946 References 

2947 ---------- 

2948 .. [Wilks2011] Wilks, D.S., (2011) "Statistical Methods in the Atmospheric 

2949 Sciences" Third Edition, vol. 100, Academic Press, Oxford, UK, 676 pp. 

2950 """ 

2951 # Check cubes using same functionality as the difference operator. 

2952 if len(cubes) != 2: 

2953 raise ValueError("cubes should contain exactly 2 cubes.") 

2954 base: Cube = cubes.extract_cube(iris.AttributeConstraint(cset_comparison_base=1)) 

2955 other: Cube = cubes.extract_cube( 

2956 iris.Constraint( 

2957 cube_func=lambda cube: "cset_comparison_base" not in cube.attributes 

2958 ) 

2959 ) 

2960 

2961 # Get spatial coord names. 

2962 base_lat_name, base_lon_name = get_cube_yxcoordname(base) 

2963 other_lat_name, other_lon_name = get_cube_yxcoordname(other) 

2964 

2965 # Ensure cubes to compare are on common differencing grid. 

2966 # This is triggered if either 

2967 # i) latitude and longitude shapes are not the same. Note grid points 

2968 # are not compared directly as these can differ through rounding 

2969 # errors. 

2970 # ii) or variables are known to often sit on different grid staggering 

2971 # in different models (e.g. cell center vs cell edge), as is the case 

2972 # for UM and LFRic comparisons. 

2973 # In future greater choice of regridding method might be applied depending 

2974 # on variable type. Linear regridding can in general be appropriate for smooth 

2975 # variables. Care should be taken with interpretation of differences 

2976 # given this dependency on regridding. 

2977 if ( 

2978 base.coord(base_lat_name).shape != other.coord(other_lat_name).shape 

2979 or base.coord(base_lon_name).shape != other.coord(other_lon_name).shape 

2980 ) or ( 

2981 base.long_name 

2982 in [ 

2983 "eastward_wind_at_10m", 

2984 "northward_wind_at_10m", 

2985 "northward_wind_at_cell_centres", 

2986 "eastward_wind_at_cell_centres", 

2987 "zonal_wind_at_pressure_levels", 

2988 "meridional_wind_at_pressure_levels", 

2989 "potential_vorticity_at_pressure_levels", 

2990 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging", 

2991 ] 

2992 ): 

2993 logging.debug( 

2994 "Linear regridding base cube to other grid to compute differences" 

2995 ) 

2996 base = regrid_onto_cube(base, other, method="Linear") 

2997 

2998 # Extract just common time points. 

2999 base, other = _extract_common_time_points(base, other) 

3000 

3001 # Equalise attributes so we can merge. 

3002 fully_equalise_attributes([base, other]) 

3003 logging.debug("Base: %s\nOther: %s", base, other) 

3004 

3005 # Collapse cubes. 

3006 base = collapse( 

3007 base, 

3008 coordinate=coordinates, 

3009 method="PERCENTILE", 

3010 additional_percent=percentiles, 

3011 ) 

3012 other = collapse( 

3013 other, 

3014 coordinate=coordinates, 

3015 method="PERCENTILE", 

3016 additional_percent=percentiles, 

3017 ) 

3018 

3019 # Ensure we have a name for the plot file. 

3020 recipe_title = get_recipe_metadata().get("title", "Untitled") 

3021 title = f"{recipe_title}" 

3022 

3023 if filename is None: 

3024 filename = slugify(recipe_title) 

3025 

3026 # Add file extension. 

3027 plot_filename = f"{filename.rsplit('.', 1)[0]}.png" 

3028 

3029 # Do the actual plotting on a scatter plot 

3030 _plot_and_save_scatter_plot( 

3031 base, other, plot_filename, title, one_to_one, model_names 

3032 ) 

3033 

3034 # Add list of plots to plot metadata. 

3035 plot_index = _append_to_plot_index([plot_filename]) 

3036 

3037 # Make a page to display the plots. 

3038 _make_plot_html_page(plot_index) 

3039 

3040 return iris.cube.CubeList([base, other]) 

3041 

3042 

3043def scatter_plot( 

3044 cube_x: iris.cube.Cube | iris.cube.CubeList, 

3045 cube_y: iris.cube.Cube | iris.cube.CubeList, 

3046 filename: str = None, 

3047 one_to_one: bool = True, 

3048 **kwargs, 

3049) -> iris.cube.CubeList: 

3050 """Plot a scatter plot between two variables. 

3051 

3052 Both cubes must be 1D. 

3053 

3054 Parameters 

3055 ---------- 

3056 cube_x: Cube | CubeList 

3057 1 dimensional Cube of the data to plot on y-axis. 

3058 cube_y: Cube | CubeList 

3059 1 dimensional Cube of the data to plot on x-axis. 

3060 filename: str, optional 

3061 Filename of the plot to write. 

3062 one_to_one: bool, optional 

3063 If True a 1:1 line is plotted; if False it is not. Default is True. 

3064 

3065 Returns 

3066 ------- 

3067 cubes: CubeList 

3068 CubeList of the original x and y cubes for further processing. 

3069 

3070 Raises 

3071 ------ 

3072 ValueError 

3073 If the cube doesn't have the right dimensions and cubes not the same 

3074 size. 

3075 TypeError 

3076 If the cube isn't a single cube. 

3077 

3078 Notes 

3079 ----- 

3080 Scatter plots are used for determining if there is a relationship between 

3081 two variables. Positive relations have a slope going from bottom left to top 

3082 right; Negative relations have a slope going from top left to bottom right. 

3083 """ 

3084 # Iterate over all cubes in cube or CubeList and plot. 

3085 for cube_iter in iter_maybe(cube_x): 

3086 # Check cubes are correct shape. 

3087 cube_iter = _check_single_cube(cube_iter) 

3088 if cube_iter.ndim > 1: 

3089 raise ValueError("cube_x must be 1D.") 

3090 

3091 # Iterate over all cubes in cube or CubeList and plot. 

3092 for cube_iter in iter_maybe(cube_y): 

3093 # Check cubes are correct shape. 

3094 cube_iter = _check_single_cube(cube_iter) 

3095 if cube_iter.ndim > 1: 

3096 raise ValueError("cube_y must be 1D.") 

3097 

3098 # Ensure we have a name for the plot file. 

3099 recipe_title = get_recipe_metadata().get("title", "Untitled") 

3100 title = f"{recipe_title}" 

3101 

3102 if filename is None: 

3103 filename = slugify(recipe_title) 

3104 

3105 # Add file extension. 

3106 plot_filename = f"{filename.rsplit('.', 1)[0]}.png" 

3107 

3108 # Do the actual plotting. 

3109 _plot_and_save_scatter_plot(cube_x, cube_y, plot_filename, title, one_to_one) 

3110 

3111 # Add list of plots to plot metadata. 

3112 plot_index = _append_to_plot_index([plot_filename]) 

3113 

3114 # Make a page to display the plots. 

3115 _make_plot_html_page(plot_index) 

3116 

3117 return iris.cube.CubeList([cube_x, cube_y]) 

3118 

3119 

3120def vector_plot( 

3121 cube_u: iris.cube.Cube, 

3122 cube_v: iris.cube.Cube, 

3123 filename: str = None, 

3124 sequence_coordinate: str = "time", 

3125 **kwargs, 

3126) -> iris.cube.CubeList: 

3127 """Plot a vector plot based on the input u and v components.""" 

3128 recipe_title = get_recipe_metadata().get("title", "Untitled") 

3129 

3130 # Cubes must have a matching sequence coordinate. 

3131 try: 

3132 # Check that the u and v cubes have the same sequence coordinate. 

3133 if cube_u.coord(sequence_coordinate) != cube_v.coord(sequence_coordinate): 3133 ↛ anywhereline 3133 didn't jump anywhere: it always raised an exception.

3134 raise ValueError("Coordinates do not match.") 

3135 except (iris.exceptions.CoordinateNotFoundError, ValueError) as err: 

3136 raise ValueError( 

3137 f"Cubes should have matching {sequence_coordinate} coordinate:\n{cube_u}\n{cube_v}" 

3138 ) from err 

3139 

3140 # Create a plot for each value of the sequence coordinate. 

3141 plot_index = [] 

3142 nplot = np.size(cube_u[0].coord(sequence_coordinate).points) 

3143 for cube_u_slice, cube_v_slice in zip( 

3144 cube_u.slices_over(sequence_coordinate), 

3145 cube_v.slices_over(sequence_coordinate), 

3146 strict=True, 

3147 ): 

3148 # Format the coordinate value in a unit appropriate way. 

3149 seq_coord = cube_u_slice.coord(sequence_coordinate) 

3150 plot_title, plot_filename = _set_title_and_filename( 

3151 seq_coord, nplot, recipe_title, filename 

3152 ) 

3153 

3154 # Do the actual plotting. 

3155 _plot_and_save_vector_plot( 

3156 cube_u_slice, 

3157 cube_v_slice, 

3158 filename=plot_filename, 

3159 title=plot_title, 

3160 method="contourf", 

3161 ) 

3162 plot_index.append(plot_filename) 

3163 

3164 # Add list of plots to plot metadata. 

3165 complete_plot_index = _append_to_plot_index(plot_index) 

3166 

3167 # Make a page to display the plots. 

3168 _make_plot_html_page(complete_plot_index) 

3169 

3170 return iris.cube.CubeList([cube_u, cube_v]) 

3171 

3172 

3173def plot_histogram_series( 

3174 cubes: iris.cube.Cube | iris.cube.CubeList, 

3175 filename: str = None, 

3176 sequence_coordinate: str = "time", 

3177 stamp_coordinate: str = "realization", 

3178 single_plot: bool = False, 

3179 **kwargs, 

3180) -> iris.cube.Cube | iris.cube.CubeList: 

3181 """Plot a histogram plot for each vertical level provided. 

3182 

3183 A histogram plot can be plotted, but if the sequence_coordinate (i.e. time) 

3184 is present then a sequence of plots will be produced using the time slider 

3185 functionality to scroll through histograms against time. If a 

3186 stamp_coordinate is present then postage stamp plots will be produced. If 

3187 stamp_coordinate and single_plot is True, all postage stamp plots will be 

3188 plotted in a single plot instead of separate postage stamp plots. 

3189 

3190 Parameters 

3191 ---------- 

3192 cubes: Cube | iris.cube.CubeList 

3193 Iris cube or CubeList of the data to plot. It should have a single dimension other 

3194 than the stamp coordinate. 

3195 The cubes should cover the same phenomenon i.e. all cubes contain temperature data. 

3196 We do not support different data such as temperature and humidity in the same CubeList for plotting. 

3197 filename: str, optional 

3198 Name of the plot to write, used as a prefix for plot sequences. Defaults 

3199 to the recipe name. 

3200 sequence_coordinate: str, optional 

3201 Coordinate about which to make a plot sequence. Defaults to ``"time"``. 

3202 This coordinate must exist in the cube and will be used for the time 

3203 slider. 

3204 stamp_coordinate: str, optional 

3205 Coordinate about which to plot postage stamp plots. Defaults to 

3206 ``"realization"``. 

3207 single_plot: bool, optional 

3208 If True, all postage stamp plots will be plotted in a single plot. If 

3209 False, each postage stamp plot will be plotted separately. Is only valid 

3210 if stamp_coordinate exists and has more than a single point. 

3211 

3212 Returns 

3213 ------- 

3214 iris.cube.Cube | iris.cube.CubeList 

3215 The original Cube or CubeList (so further operations can be applied). 

3216 Plotted data. 

3217 

3218 Raises 

3219 ------ 

3220 ValueError 

3221 If the cube doesn't have the right dimensions. 

3222 TypeError 

3223 If the cube isn't a Cube or CubeList. 

3224 """ 

3225 recipe_title = get_recipe_metadata().get("title", "Untitled") 

3226 

3227 cubes = iter_maybe(cubes) 

3228 

3229 # Internal plotting function. 

3230 plotting_func = _plot_and_save_histogram_series 

3231 

3232 num_models = _get_num_models(cubes) 

3233 

3234 _validate_cube_shape(cubes, num_models) 

3235 

3236 # If several histograms are plotted with time as sequence_coordinate for the 

3237 # time slider option. 

3238 for cube in cubes: 

3239 try: 

3240 cube.coord(sequence_coordinate) 

3241 except iris.exceptions.CoordinateNotFoundError as err: 

3242 raise ValueError( 

3243 f"Cube must have a {sequence_coordinate} coordinate." 

3244 ) from err 

3245 

3246 # Get minimum and maximum from levels information. 

3247 levels = None 

3248 for cube in cubes: 3248 ↛ 3264line 3248 didn't jump to line 3264 because the loop on line 3248 didn't complete

3249 # First check if user-specified "auto" range variable. 

3250 # This maintains the value of levels as None, so proceed. 

3251 _, levels, _ = _colorbar_map_levels(cube, axis="y") 

3252 if levels is None: 

3253 break 

3254 # If levels is changed, recheck to use the vmin,vmax or 

3255 # levels-based ranges for histogram plots. 

3256 _, levels, _ = _colorbar_map_levels(cube) 

3257 logging.debug("levels: %s", levels) 

3258 if levels is not None: 3258 ↛ 3248line 3258 didn't jump to line 3248 because the condition on line 3258 was always true

3259 vmin = min(levels) 

3260 vmax = max(levels) 

3261 logging.debug("Updated vmin, vmax: %s, %s", vmin, vmax) 

3262 break 

3263 

3264 if levels is None: 

3265 vmin = min(cb.data.min() for cb in cubes) 

3266 vmax = max(cb.data.max() for cb in cubes) 

3267 

3268 # Make postage stamp plots if stamp_coordinate exists and has more than a 

3269 # single point. If single_plot is True: 

3270 # -- all postage stamp plots will be plotted in a single plot instead of 

3271 # separate postage stamp plots. 

3272 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are 

3273 # produced per single model only 

3274 if num_models == 1: 3274 ↛ 3287line 3274 didn't jump to line 3287 because the condition on line 3274 was always true

3275 if ( 3275 ↛ 3279line 3275 didn't jump to line 3279 because the condition on line 3275 was never true

3276 stamp_coordinate in [c.name() for c in cubes[0].coords()] 

3277 and cubes[0].coord(stamp_coordinate).shape[0] > 1 

3278 ): 

3279 if single_plot: 

3280 plotting_func = ( 

3281 _plot_and_save_postage_stamps_in_single_plot_histogram_series 

3282 ) 

3283 else: 

3284 plotting_func = _plot_and_save_postage_stamp_histogram_series 

3285 cube_iterables = cubes[0].slices_over(sequence_coordinate) 

3286 else: 

3287 all_points = sorted( 

3288 set( 

3289 itertools.chain.from_iterable( 

3290 cb.coord(sequence_coordinate).points for cb in cubes 

3291 ) 

3292 ) 

3293 ) 

3294 all_slices = list( 

3295 itertools.chain.from_iterable( 

3296 cb.slices_over(sequence_coordinate) for cb in cubes 

3297 ) 

3298 ) 

3299 # Matched slices (matched by seq coord point; it may happen that 

3300 # evaluated models do not cover the same seq coord range, hence matching 

3301 # necessary) 

3302 cube_iterables = [ 

3303 iris.cube.CubeList( 

3304 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point 

3305 ) 

3306 for point in all_points 

3307 ] 

3308 

3309 plot_index = [] 

3310 nplot = np.size(cube.coord(sequence_coordinate).points) 

3311 # Create a plot for each value of the sequence coordinate. Allowing for 

3312 # multiple cubes in a CubeList to be plotted in the same plot for similar 

3313 # sequence values. Passing a CubeList into the internal plotting function 

3314 # for similar values of the sequence coordinate. cube_slice can be an 

3315 # iris.cube.Cube or an iris.cube.CubeList. 

3316 for cube_slice in cube_iterables: 

3317 single_cube = cube_slice 

3318 if isinstance(cube_slice, iris.cube.CubeList): 3318 ↛ 3319line 3318 didn't jump to line 3319 because the condition on line 3318 was never true

3319 single_cube = cube_slice[0] 

3320 

3321 # Ensure valid stamp coordinate in cube dimensions 

3322 if stamp_coordinate == "realization": 3322 ↛ 3325line 3322 didn't jump to line 3325 because the condition on line 3322 was always true

3323 stamp_coordinate = check_stamp_coordinate(single_cube) 

3324 # Set plot titles and filename, based on sequence coordinate 

3325 seq_coord = single_cube.coord(sequence_coordinate) 

3326 # Use time coordinate in title and filename if single histogram output. 

3327 if sequence_coordinate == "realization" and nplot == 1: 3327 ↛ 3328line 3327 didn't jump to line 3328 because the condition on line 3327 was never true

3328 seq_coord = single_cube.coord("time") 

3329 plot_title, plot_filename = _set_title_and_filename( 

3330 seq_coord, nplot, recipe_title, filename 

3331 ) 

3332 

3333 # Do the actual plotting. 

3334 plotting_func( 

3335 cube_slice, 

3336 filename=plot_filename, 

3337 stamp_coordinate=stamp_coordinate, 

3338 title=plot_title, 

3339 vmin=vmin, 

3340 vmax=vmax, 

3341 ) 

3342 plot_index.append(plot_filename) 

3343 

3344 # Add list of plots to plot metadata. 

3345 complete_plot_index = _append_to_plot_index(plot_index) 

3346 

3347 # Make a page to display the plots. 

3348 _make_plot_html_page(complete_plot_index) 

3349 

3350 return cubes 

3351 

3352 

3353def plot_power_spectrum_series( 

3354 cubes: iris.cube.Cube | iris.cube.CubeList, 

3355 filename: str = None, 

3356 sequence_coordinate: str = "time", 

3357 stamp_coordinate: str = "realization", 

3358 single_plot: bool = False, 

3359 **kwargs, 

3360) -> iris.cube.Cube | iris.cube.CubeList: 

3361 """Plot a power spectrum plot for each vertical level provided. 

3362 

3363 A power spectrum plot can be plotted, but if the sequence_coordinate (i.e. time) 

3364 is present then a sequence of plots will be produced using the time slider 

3365 functionality to scroll through power spectrum against time. If a 

3366 stamp_coordinate is present then postage stamp plots will be produced. If 

3367 stamp_coordinate and single_plot is True, all postage stamp plots will be 

3368 plotted in a single plot instead of separate postage stamp plots. 

3369 

3370 Parameters 

3371 ---------- 

3372 cubes: Cube | iris.cube.CubeList 

3373 Iris cube or CubeList of the data to plot. It should have a single dimension other 

3374 than the stamp coordinate. 

3375 The cubes should cover the same phenomenon i.e. all cubes contain temperature data. 

3376 We do not support different data such as temperature and humidity in the same CubeList for plotting. 

3377 filename: str, optional 

3378 Name of the plot to write, used as a prefix for plot sequences. Defaults 

3379 to the recipe name. 

3380 sequence_coordinate: str, optional 

3381 Coordinate about which to make a plot sequence. Defaults to ``"time"``. 

3382 This coordinate must exist in the cube and will be used for the time 

3383 slider. 

3384 stamp_coordinate: str, optional 

3385 Coordinate about which to plot postage stamp plots. Defaults to 

3386 ``"realization"``. 

3387 single_plot: bool, optional 

3388 If True, all postage stamp plots will be plotted in a single plot. If 

3389 False, each postage stamp plot will be plotted separately. Is only valid 

3390 if stamp_coordinate exists and has more than a single point. 

3391 

3392 Returns 

3393 ------- 

3394 iris.cube.Cube | iris.cube.CubeList 

3395 The original Cube or CubeList (so further operations can be applied). 

3396 Plotted data. 

3397 

3398 Raises 

3399 ------ 

3400 ValueError 

3401 If the cube doesn't have the right dimensions. 

3402 TypeError 

3403 If the cube isn't a Cube or CubeList. 

3404 """ 

3405 recipe_title = get_recipe_metadata().get("title", "Untitled") 

3406 

3407 cubes = iter_maybe(cubes) 

3408 

3409 # Internal plotting function. 

3410 plotting_func = _plot_and_save_power_spectrum_series 

3411 

3412 num_models = _get_num_models(cubes) 

3413 

3414 _validate_cube_shape(cubes, num_models) 

3415 

3416 # If several power spectra are plotted with time as sequence_coordinate for the 

3417 # time slider option. 

3418 for cube in cubes: 

3419 try: 

3420 cube.coord(sequence_coordinate) 

3421 except iris.exceptions.CoordinateNotFoundError as err: 

3422 raise ValueError( 

3423 f"Cube must have a {sequence_coordinate} coordinate." 

3424 ) from err 

3425 

3426 # Make postage stamp plots if stamp_coordinate exists and has more than a 

3427 # single point. If single_plot is True: 

3428 # -- all postage stamp plots will be plotted in a single plot instead of 

3429 # separate postage stamp plots. 

3430 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are 

3431 # produced per single model only 

3432 if num_models == 1: 3432 ↛ 3445line 3432 didn't jump to line 3445 because the condition on line 3432 was always true

3433 if ( 3433 ↛ 3437line 3433 didn't jump to line 3437 because the condition on line 3433 was never true

3434 stamp_coordinate in [c.name() for c in cubes[0].coords()] 

3435 and cubes[0].coord(stamp_coordinate).shape[0] > 1 

3436 ): 

3437 if single_plot: 

3438 plotting_func = ( 

3439 _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series 

3440 ) 

3441 else: 

3442 plotting_func = _plot_and_save_postage_stamp_power_spectrum_series 

3443 cube_iterables = cubes[0].slices_over(sequence_coordinate) 

3444 else: 

3445 all_points = sorted( 

3446 set( 

3447 itertools.chain.from_iterable( 

3448 cb.coord(sequence_coordinate).points for cb in cubes 

3449 ) 

3450 ) 

3451 ) 

3452 all_slices = list( 

3453 itertools.chain.from_iterable( 

3454 cb.slices_over(sequence_coordinate) for cb in cubes 

3455 ) 

3456 ) 

3457 # Matched slices (matched by seq coord point; it may happen that 

3458 # evaluated models do not cover the same seq coord range, hence matching 

3459 # necessary) 

3460 cube_iterables = [ 

3461 iris.cube.CubeList( 

3462 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point 

3463 ) 

3464 for point in all_points 

3465 ] 

3466 

3467 plot_index = [] 

3468 nplot = np.size(cube.coord(sequence_coordinate).points) 

3469 # Create a plot for each value of the sequence coordinate. Allowing for 

3470 # multiple cubes in a CubeList to be plotted in the same plot for similar 

3471 # sequence values. Passing a CubeList into the internal plotting function 

3472 # for similar values of the sequence coordinate. cube_slice can be an 

3473 # iris.cube.Cube or an iris.cube.CubeList. 

3474 for cube_slice in cube_iterables: 

3475 single_cube = cube_slice 

3476 if isinstance(cube_slice, iris.cube.CubeList): 3476 ↛ 3477line 3476 didn't jump to line 3477 because the condition on line 3476 was never true

3477 single_cube = cube_slice[0] 

3478 

3479 # Set stamp coordinate 

3480 if stamp_coordinate == "realization": 3480 ↛ 3483line 3480 didn't jump to line 3483 because the condition on line 3480 was always true

3481 stamp_coordinate = check_stamp_coordinate(single_cube) 

3482 # Set plot title and filenames based on sequence values 

3483 seq_coord = single_cube.coord(sequence_coordinate) 

3484 plot_title, plot_filename = _set_title_and_filename( 

3485 seq_coord, nplot, recipe_title, filename 

3486 ) 

3487 

3488 # Do the actual plotting. 

3489 plotting_func( 

3490 cube_slice, 

3491 filename=plot_filename, 

3492 stamp_coordinate=stamp_coordinate, 

3493 title=plot_title, 

3494 ) 

3495 plot_index.append(plot_filename) 

3496 

3497 # Add list of plots to plot metadata. 

3498 complete_plot_index = _append_to_plot_index(plot_index) 

3499 

3500 # Make a page to display the plots. 

3501 _make_plot_html_page(complete_plot_index) 

3502 

3503 return cubes 

3504 

3505 

3506def _DCT_ps(y_3d): 

3507 """Calculate power spectra for regional domains. 

3508 

3509 Parameters 

3510 ---------- 

3511 y_3d: 3D array 

3512 3 dimensional array to calculate spectrum for. 

3513 (2D field data with 3rd dimension of time) 

3514 

3515 Returns: ps_array 

3516 Array of power spectra values calculated for input field (for each time) 

3517 

3518 Method for regional domains: 

3519 Calculate power spectra over limited area domain using Discrete Cosine Transform (DCT) 

3520 as described in Denis et al 2002 [Denis_etal_2002]_. 

3521 

3522 References 

3523 ---------- 

3524 .. [Denis_etal_2002] Bertrand Denis, Jean Côté and René Laprise (2002) 

3525 "Spectral Decomposition of Two-Dimensional Atmospheric Fields on 

3526 Limited-Area Domains Using the Discrete Cosine Transform (DCT)" 

3527 Monthly Weather Review, Vol. 130, 1812-1828 

3528 doi: https://doi.org/10.1175/1520-0493(2002)130<1812:SDOTDA>2.0.CO;2 

3529 """ 

3530 Nt, Ny, Nx = y_3d.shape 

3531 

3532 # Max coefficient 

3533 Nmin = min(Nx - 1, Ny - 1) 

3534 

3535 # Create alpha matrix (of wavenumbers) 

3536 alpha_matrix = _create_alpha_matrix(Ny, Nx) 

3537 

3538 # Prepare output array 

3539 ps_array = np.zeros((Nt, Nmin)) 

3540 

3541 # Loop over time to get spectrum for each time. 

3542 for t in range(Nt): 

3543 y_2d = y_3d[t] 

3544 

3545 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space. 

3546 # fkk is a 2D array of DCT coefficients, representing the amplitudes of 

3547 # cosine basis functions at different spatial frequencies. 

3548 

3549 # normalise spectrum to allow comparison between models. 

3550 fkk = fft.dctn(y_2d, norm="ortho") 

3551 

3552 # Normalise fkk 

3553 fkk = fkk / np.sqrt(Ny * Nx) 

3554 

3555 # calculate variance of spectral coefficient 

3556 sigma_2 = fkk**2 / Nx / Ny 

3557 

3558 # Group ellipses of alphas into the same wavenumber k/Nmin 

3559 for k in range(1, Nmin + 1): 

3560 alpha = k / Nmin 

3561 alpha_p1 = (k + 1) / Nmin 

3562 

3563 # Sum up elements matching k 

3564 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1)) 

3565 ps_array[t, k - 1] = np.sum(sigma_2[mask_k]) 

3566 

3567 return ps_array 

3568 

3569 

3570def _create_alpha_matrix(Ny, Nx): 

3571 """Construct an array of 2D wavenumbers from 2D wavenumber pair. 

3572 

3573 Parameters 

3574 ---------- 

3575 Ny, Nx: 

3576 Dimensions of the 2D field for which the power spectra is calculated. Used to 

3577 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a 

3578 single-scale parameter. 

3579 

3580 Returns: alpha_matrix 

3581 normalisation of 2D wavenumber axes, transforming the spectral domain into 

3582 an elliptic coordinate system. 

3583 

3584 """ 

3585 # Create x_indices: each row is [1, 2, ..., Nx] 

3586 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1)) 

3587 

3588 # Create y_indices: each column is [1, 2, ..., Ny] 

3589 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx)) 

3590 

3591 # Compute alpha_matrix 

3592 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2) 

3593 

3594 return alpha_matrix