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

929 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-11-28 16:33 +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 

25import sys 

26from typing import Literal 

27 

28import cartopy.crs as ccrs 

29import iris 

30import iris.coords 

31import iris.cube 

32import iris.exceptions 

33import iris.plot as iplt 

34import matplotlib as mpl 

35import matplotlib.colors as mcolors 

36import matplotlib.pyplot as plt 

37import numpy as np 

38import scipy.fft as fft 

39from iris.cube import Cube 

40from markdown_it import MarkdownIt 

41 

42from CSET._common import ( 

43 combine_dicts, 

44 get_recipe_metadata, 

45 iter_maybe, 

46 render_file, 

47 slugify, 

48) 

49from CSET.operators._utils import ( 

50 fully_equalise_attributes, 

51 get_cube_yxcoordname, 

52 is_transect, 

53) 

54from CSET.operators.collapse import collapse 

55from CSET.operators.misc import _extract_common_time_points 

56from CSET.operators.regrid import regrid_onto_cube 

57 

58# Use a non-interactive plotting backend. 

59mpl.use("agg") 

60 

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

62 

63############################ 

64# Private helper functions # 

65############################ 

66 

67 

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

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

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

71 fcntl.flock(fp, fcntl.LOCK_EX) 

72 fp.seek(0) 

73 meta = json.load(fp) 

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

75 complete_plot_index = complete_plot_index + plot_index 

76 meta["plots"] = complete_plot_index 

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

78 os.getenv("DO_CASE_AGGREGATION") 

79 ): 

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

81 fp.seek(0) 

82 fp.truncate() 

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

84 return complete_plot_index 

85 

86 

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

88 """Ensure a single cube is given. 

89 

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

91 otherwise an error is raised. 

92 

93 Parameters 

94 ---------- 

95 cube: Cube | CubeList 

96 The cube to check. 

97 

98 Returns 

99 ------- 

100 cube: Cube 

101 The checked cube. 

102 

103 Raises 

104 ------ 

105 TypeError 

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

107 """ 

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

109 return cube 

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

111 if len(cube) == 1: 

112 return cube[0] 

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

114 

115 

116def _py312_importlib_resources_files_shim(): 

117 """Importlib behaviour changed in 3.12 to avoid circular dependencies. 

118 

119 This shim is needed until python 3.12 is our oldest supported version, after 

120 which it can just be replaced by directly using importlib.resources.files. 

121 """ 

122 if sys.version_info.minor >= 12: 

123 files = importlib.resources.files() 

124 else: 

125 import CSET.operators 

126 

127 files = importlib.resources.files(CSET.operators) 

128 return files 

129 

130 

131def _make_plot_html_page(plots: list): 

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

133 # Debug check that plots actually contains some strings. 

134 assert isinstance(plots[0], str) 

135 

136 # Load HTML template file. 

137 operator_files = _py312_importlib_resources_files_shim() 

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

139 

140 # Get some metadata. 

141 meta = get_recipe_metadata() 

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

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

144 

145 # Prepare template variables. 

146 variables = { 

147 "title": title, 

148 "description": description, 

149 "initial_plot": plots[0], 

150 "plots": plots, 

151 "title_slug": slugify(title), 

152 } 

153 

154 # Render template. 

155 html = render_file(template_file, **variables) 

156 

157 # Save completed HTML. 

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

159 fp.write(html) 

160 

161 

162@functools.cache 

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

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

165 

166 This is a separate function to make it cacheable. 

167 """ 

168 colorbar_file = _py312_importlib_resources_files_shim().joinpath( 

169 "_colorbar_definition.json" 

170 ) 

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

172 colorbar = json.load(fp) 

173 

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

175 override_colorbar = {} 

176 if user_colorbar_file: 

177 try: 

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

179 override_colorbar = json.load(fp) 

180 except FileNotFoundError: 

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

182 

183 # Overwrite values with the user supplied colorbar definition. 

184 colorbar = combine_dicts(colorbar, override_colorbar) 

185 return colorbar 

186 

187 

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

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

190 

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

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

193 to model_name attribute. 

194 

195 Parameters 

196 ---------- 

197 cubes: CubeList or Cube 

198 Cubes with model_name attribute 

199 

200 Returns 

201 ------- 

202 model_colors_map: 

203 Dictionary mapping model_name attribute to colors 

204 """ 

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

206 colorbar = _load_colorbar_map(user_colorbar_file) 

207 model_names = sorted( 

208 filter( 

209 lambda x: x is not None, 

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

211 ) 

212 ) 

213 if not model_names: 

214 return {} 

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

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

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

218 

219 color_list = itertools.cycle(DEFAULT_DISCRETE_COLORS) 

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

221 

222 

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

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

225 

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

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

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

229 exist for specific pressure levels to account for variables with 

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

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

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

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

234 

235 Parameters 

236 ---------- 

237 cube: Cube 

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

239 axis: "x", "y", optional 

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

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

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

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

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

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

246 

247 Returns 

248 ------- 

249 cmap: 

250 Matplotlib colormap. 

251 levels: 

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

253 should be taken as the range. 

254 norm: 

255 BoundaryNorm information. 

256 """ 

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

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

259 colorbar = _load_colorbar_map(user_colorbar_file) 

260 cmap = None 

261 

262 try: 

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

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

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

266 pressure_level = str(int(pressure_level_raw)) 

267 except iris.exceptions.CoordinateNotFoundError: 

268 pressure_level = None 

269 

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

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

272 # consistent. 

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

274 for varname in varnames: 

275 # Get the colormap for this variable. 

276 try: 

277 var_colorbar = colorbar[varname] 

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

279 varname_key = varname 

280 break 

281 except KeyError: 

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

283 

284 # Get colormap if it is a mask. 

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

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

287 return cmap, levels, norm 

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

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

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

291 return cmap, levels, norm 

292 # If probability is plotted use custom colorbar and levels 

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

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

295 return cmap, levels, norm 

296 # If aviation colour state use custom colorbar and levels 

297 if any("aviation_colour_state" in name for name in varnames): 297 ↛ 298line 297 didn't jump to line 298 because the condition on line 297 was never true

298 cmap, levels, norm = _custom_colormap_aviation_colour_state(cube) 

299 return cmap, levels, norm 

300 

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

302 if not cmap: 

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

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

305 return cmap, levels, norm 

306 

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

308 if pressure_level: 

309 try: 

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

311 except KeyError: 

312 logging.debug( 

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

314 varname, 

315 pressure_level, 

316 ) 

317 

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

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

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

321 if axis: 

322 if axis == "x": 

323 try: 

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

325 except KeyError: 

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

327 if axis == "y": 

328 try: 

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

330 except KeyError: 

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

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

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

334 levels = None 

335 else: 

336 levels = [vmin, vmax] 

337 return None, levels, None 

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

339 else: 

340 try: 

341 levels = var_colorbar["levels"] 

342 # Use discrete bins when levels are specified, rather 

343 # than a smooth range. 

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

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

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

347 except KeyError: 

348 # Get the range for this variable. 

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

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

351 # Calculate levels from range. 

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

353 norm = None 

354 

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

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

357 # JSON file. 

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

359 cmap, levels, norm = _custom_colourmap_visibility_in_air( 

360 cube, cmap, levels, norm 

361 ) 

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

363 return cmap, levels, norm 

364 

365 

366def _setup_spatial_map( 

367 cube: iris.cube.Cube, 

368 figure, 

369 cmap, 

370 grid_size: int | None = None, 

371 subplot: int | None = None, 

372): 

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

374 

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

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

377 

378 Parameters 

379 ---------- 

380 cube: Cube 

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

382 figure: 

383 Matplotlib Figure object holding all plot elements. 

384 cmap: 

385 Matplotlib colormap. 

386 grid_size: int, optional 

387 Size of grid for subplots if multiple spatial subplots in figure. 

388 subplot: int, optional 

389 Subplot index if multiple spatial subplots in figure. 

390 

391 Returns 

392 ------- 

393 axes: 

394 Matplotlib GeoAxes definition. 

395 """ 

396 # Identify min/max plot bounds. 

397 try: 

398 lat_axis, lon_axis = get_cube_yxcoordname(cube) 

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

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

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

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

403 

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

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

406 x1 = x1 - 180.0 

407 x2 = x2 - 180.0 

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

409 

410 # Consider map projection orientation. 

411 # Adapting orientation enables plotting across international dateline. 

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

413 if x2 > 180.0: 

414 central_longitude = 180.0 

415 else: 

416 central_longitude = 0.0 

417 

418 # Define spatial map projection. 

419 coord_system = cube.coord(lat_axis).coord_system 

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

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

422 projection = ccrs.RotatedPole( 

423 pole_longitude=coord_system.grid_north_pole_longitude, 

424 pole_latitude=coord_system.grid_north_pole_latitude, 

425 central_rotated_longitude=0.0, 

426 ) 

427 crs = projection 

428 else: 

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

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

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

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

433 projection = ccrs.PlateCarree(central_longitude=central_longitude) 

434 crs = ccrs.PlateCarree() 

435 

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

437 if subplot is not None: 

438 axes = figure.add_subplot( 

439 grid_size, grid_size, subplot, projection=projection 

440 ) 

441 else: 

442 axes = figure.add_subplot(projection=projection) 

443 

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

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

446 coastcol = "magenta" 

447 else: 

448 coastcol = "black" 

449 logging.debug("Plotting coastlines in colour %s.", coastcol) 

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

451 

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

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

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

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

456 

457 except ValueError: 

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

459 axes = figure.gca() 

460 pass 

461 

462 return axes 

463 

464 

465def _get_plot_resolution() -> int: 

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

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

468 

469 

470def _plot_and_save_spatial_plot( 

471 cube: iris.cube.Cube, 

472 filename: str, 

473 title: str, 

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

475 **kwargs, 

476): 

477 """Plot and save a spatial plot. 

478 

479 Parameters 

480 ---------- 

481 cube: Cube 

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

483 filename: str 

484 Filename of the plot to write. 

485 title: str 

486 Plot title. 

487 method: "contourf" | "pcolormesh" 

488 The plotting method to use. 

489 """ 

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

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

492 

493 # Specify the color bar 

494 cmap, levels, norm = _colorbar_map_levels(cube) 

495 

496 # Setup plot map projection, extent and coastlines. 

497 axes = _setup_spatial_map(cube, fig, cmap) 

498 

499 # Plot the field. 

500 if method == "contourf": 

501 # Filled contour plot of the field. 

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

503 elif method == "pcolormesh": 

504 try: 

505 vmin = min(levels) 

506 vmax = max(levels) 

507 except TypeError: 

508 vmin, vmax = None, None 

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

510 # if levels are defined. 

511 if norm is not None: 

512 vmin = None 

513 vmax = None 

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

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

516 else: 

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

518 

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

520 if is_transect(cube): 

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

522 axes.invert_yaxis() 

523 axes.set_yscale("log") 

524 axes.set_ylim(1100, 100) 

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

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

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

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

529 ): 

530 axes.set_yscale("log") 

531 

532 axes.set_title( 

533 f"{title}\n" 

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

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

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

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

538 fontsize=16, 

539 ) 

540 

541 else: 

542 # Add title. 

543 axes.set_title(title, fontsize=16) 

544 

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

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

547 axes.annotate( 

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

549 xy=(1, -0.05), 

550 xycoords="axes fraction", 

551 xytext=(-5, 5), 

552 textcoords="offset points", 

553 ha="right", 

554 va="bottom", 

555 size=11, 

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

557 ) 

558 

559 # Add colour bar. 

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

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

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

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

564 cbar.set_ticks(levels) 

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

566 if "visibility" in cube.name(): 566 ↛ 567line 566 didn't jump to line 567 because the condition on line 566 was never true

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

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

569 

570 # Save plot. 

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

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

573 plt.close(fig) 

574 

575 

576def _plot_and_save_postage_stamp_spatial_plot( 

577 cube: iris.cube.Cube, 

578 filename: str, 

579 stamp_coordinate: str, 

580 title: str, 

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

582 **kwargs, 

583): 

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

585 

586 Parameters 

587 ---------- 

588 cube: Cube 

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

590 filename: str 

591 Filename of the plot to write. 

592 stamp_coordinate: str 

593 Coordinate that becomes different plots. 

594 method: "contourf" | "pcolormesh" 

595 The plotting method to use. 

596 

597 Raises 

598 ------ 

599 ValueError 

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

601 """ 

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

603 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points)))) 

604 

605 fig = plt.figure(figsize=(10, 10)) 

606 

607 # Specify the color bar 

608 cmap, levels, norm = _colorbar_map_levels(cube) 

609 

610 # Make a subplot for each member. 

611 for member, subplot in zip( 

612 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False 

613 ): 

614 # Setup subplot map projection, extent and coastlines. 

615 axes = _setup_spatial_map( 

616 member, fig, cmap, grid_size=grid_size, subplot=subplot 

617 ) 

618 if method == "contourf": 

619 # Filled contour plot of the field. 

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

621 elif method == "pcolormesh": 

622 if levels is not None: 

623 vmin = min(levels) 

624 vmax = max(levels) 

625 else: 

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

627 vmin, vmax = None, None 

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

629 # if levels are defined. 

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

631 vmin = None 

632 vmax = None 

633 # pcolormesh plot of the field. 

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

635 else: 

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

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

638 axes.set_axis_off() 

639 

640 # Put the shared colorbar in its own axes. 

641 colorbar_axes = fig.add_axes([0.15, 0.07, 0.7, 0.03]) 

642 colorbar = fig.colorbar( 

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

644 ) 

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

646 

647 # Overall figure title. 

648 fig.suptitle(title, fontsize=16) 

649 

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

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

652 plt.close(fig) 

653 

654 

655def _plot_and_save_line_series( 

656 cubes: iris.cube.CubeList, 

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

658 ensemble_coord: str, 

659 filename: str, 

660 title: str, 

661 **kwargs, 

662): 

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

664 

665 Parameters 

666 ---------- 

667 cubes: Cube or CubeList 

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

669 coords: list[Coord] 

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

671 ensemble_coord: str 

672 Ensemble coordinate in the cube. 

673 filename: str 

674 Filename of the plot to write. 

675 title: str 

676 Plot title. 

677 """ 

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

679 

680 model_colors_map = _get_model_colors_map(cubes) 

681 

682 # Store min/max ranges. 

683 y_levels = [] 

684 

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

686 _validate_cubes_coords(cubes, coords) 

687 

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

689 label = None 

690 color = "black" 

691 if model_colors_map: 

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

693 color = model_colors_map.get(label) 

694 for cube_slice in cube.slices_over(ensemble_coord): 

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

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

697 iplt.plot( 

698 coord, 

699 cube_slice, 

700 color=color, 

701 marker="o", 

702 ls="-", 

703 lw=3, 

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

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

706 else label, 

707 ) 

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

709 else: 

710 iplt.plot( 

711 coord, 

712 cube_slice, 

713 color=color, 

714 ls="-", 

715 lw=1.5, 

716 alpha=0.75, 

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

718 ) 

719 

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

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

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

723 y_levels.append(min(levels)) 

724 y_levels.append(max(levels)) 

725 

726 # Get the current axes. 

727 ax = plt.gca() 

728 

729 # Add some labels and tweak the style. 

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

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

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

733 ax.set_title(title, fontsize=16) 

734 

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

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

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

738 

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

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

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

742 # Add zero line. 

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

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

745 logging.debug( 

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

747 ) 

748 else: 

749 ax.autoscale() 

750 

751 # Add gridlines 

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

753 # Ientify unique labels for legend 

754 handles = list( 

755 { 

756 label: handle 

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

758 }.values() 

759 ) 

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

761 

762 # Save plot. 

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

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

765 plt.close(fig) 

766 

767 

768def _plot_and_save_vertical_line_series( 

769 cubes: iris.cube.CubeList, 

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

771 ensemble_coord: str, 

772 filename: str, 

773 series_coordinate: str, 

774 title: str, 

775 vmin: float, 

776 vmax: float, 

777 **kwargs, 

778): 

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

780 

781 Parameters 

782 ---------- 

783 cubes: CubeList 

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

785 coord: list[Coord] 

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

787 ensemble_coord: str 

788 Ensemble coordinate in the cube. 

789 filename: str 

790 Filename of the plot to write. 

791 series_coordinate: str 

792 Coordinate to use as vertical axis. 

793 title: str 

794 Plot title. 

795 vmin: float 

796 Minimum value for the x-axis. 

797 vmax: float 

798 Maximum value for the x-axis. 

799 """ 

800 # plot the vertical pressure axis using log scale 

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

802 

803 model_colors_map = _get_model_colors_map(cubes) 

804 

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

806 _validate_cubes_coords(cubes, coords) 

807 

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

809 label = None 

810 color = "black" 

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

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

813 color = model_colors_map.get(label) 

814 

815 for cube_slice in cube.slices_over(ensemble_coord): 

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

817 # unless single forecast. 

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

819 iplt.plot( 

820 cube_slice, 

821 coord, 

822 color=color, 

823 marker="o", 

824 ls="-", 

825 lw=3, 

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

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

828 else label, 

829 ) 

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

831 else: 

832 iplt.plot( 

833 cube_slice, 

834 coord, 

835 color=color, 

836 ls="-", 

837 lw=1.5, 

838 alpha=0.75, 

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

840 ) 

841 

842 # Get the current axis 

843 ax = plt.gca() 

844 

845 # Special handling for pressure level data. 

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

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

848 ax.invert_yaxis() 

849 ax.set_yscale("log") 

850 

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

852 y_tick_labels = [ 

853 "1000", 

854 "850", 

855 "700", 

856 "500", 

857 "300", 

858 "200", 

859 "100", 

860 ] 

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

862 

863 # Set y-axis limits and ticks. 

864 ax.set_ylim(1100, 100) 

865 

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

867 # model_level_number and lfric uses full_levels as coordinate. 

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

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

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

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

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

873 

874 ax.set_yticks(y_ticks) 

875 ax.set_yticklabels(y_tick_labels) 

876 

877 # Set x-axis limits. 

878 ax.set_xlim(vmin, vmax) 

879 # Mark y=0 if present in plot. 

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

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

882 

883 # Add some labels and tweak the style. 

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

885 ax.set_xlabel( 

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

887 ) 

888 ax.set_title(title, fontsize=16) 

889 ax.ticklabel_format(axis="x") 

890 ax.tick_params(axis="y") 

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

892 

893 # Add gridlines 

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

895 # Ientify unique labels for legend 

896 handles = list( 

897 { 

898 label: handle 

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

900 }.values() 

901 ) 

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

903 

904 # Save plot. 

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

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

907 plt.close(fig) 

908 

909 

910def _plot_and_save_scatter_plot( 

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

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

913 filename: str, 

914 title: str, 

915 one_to_one: bool, 

916 model_names: list[str] = None, 

917 **kwargs, 

918): 

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

920 

921 Parameters 

922 ---------- 

923 cube_x: Cube | CubeList 

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

925 cube_y: Cube | CubeList 

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

927 filename: str 

928 Filename of the plot to write. 

929 title: str 

930 Plot title. 

931 one_to_one: bool 

932 Whether a 1:1 line is plotted. 

933 """ 

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

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

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

937 # over the pairs simultaneously. 

938 

939 # Ensure cube_x and cube_y are iterable 

940 cube_x_iterable = iter_maybe(cube_x) 

941 cube_y_iterable = iter_maybe(cube_y) 

942 

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

944 iplt.scatter(cube_x_iter, cube_y_iter) 

945 if one_to_one is True: 

946 plt.plot( 

947 [ 

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

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

950 ], 

951 [ 

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

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

954 ], 

955 "k", 

956 linestyle="--", 

957 ) 

958 ax = plt.gca() 

959 

960 # Add some labels and tweak the style. 

961 if model_names is None: 

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

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

964 else: 

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

966 ax.set_xlabel( 

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

968 ) 

969 ax.set_ylabel( 

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

971 ) 

972 ax.set_title(title, fontsize=16) 

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

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

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

976 ax.autoscale() 

977 

978 # Save plot. 

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

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

981 plt.close(fig) 

982 

983 

984def _plot_and_save_vector_plot( 

985 cube_u: iris.cube.Cube, 

986 cube_v: iris.cube.Cube, 

987 filename: str, 

988 title: str, 

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

990 **kwargs, 

991): 

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

993 

994 Parameters 

995 ---------- 

996 cube_u: Cube 

997 2 dimensional Cube of u component of the data. 

998 cube_v: Cube 

999 2 dimensional Cube of v component of the data. 

1000 filename: str 

1001 Filename of the plot to write. 

1002 title: str 

1003 Plot title. 

1004 """ 

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

1006 

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

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

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

1010 

1011 # Specify the color bar 

1012 cmap, levels, norm = _colorbar_map_levels(cube_vec_mag) 

1013 

1014 # Setup plot map projection, extent and coastlines. 

1015 axes = _setup_spatial_map(cube_vec_mag, fig, cmap) 

1016 

1017 if method == "contourf": 1017 ↛ 1020line 1017 didn't jump to line 1020 because the condition on line 1017 was always true

1018 # Filled contour plot of the field. 

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

1020 elif method == "pcolormesh": 

1021 try: 

1022 vmin = min(levels) 

1023 vmax = max(levels) 

1024 except TypeError: 

1025 vmin, vmax = None, None 

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

1027 # if levels are defined. 

1028 if norm is not None: 

1029 vmin = None 

1030 vmax = None 

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

1032 else: 

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

1034 

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

1036 if is_transect(cube_vec_mag): 1036 ↛ 1037line 1036 didn't jump to line 1037 because the condition on line 1036 was never true

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

1038 axes.invert_yaxis() 

1039 axes.set_yscale("log") 

1040 axes.set_ylim(1100, 100) 

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

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

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

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

1045 ): 

1046 axes.set_yscale("log") 

1047 

1048 axes.set_title( 

1049 f"{title}\n" 

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

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

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

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

1054 fontsize=16, 

1055 ) 

1056 

1057 else: 

1058 # Add title. 

1059 axes.set_title(title, fontsize=16) 

1060 

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

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

1063 axes.annotate( 

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

1065 xy=(1, -0.05), 

1066 xycoords="axes fraction", 

1067 xytext=(-5, 5), 

1068 textcoords="offset points", 

1069 ha="right", 

1070 va="bottom", 

1071 size=11, 

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

1073 ) 

1074 

1075 # Add colour bar. 

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

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

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

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

1080 cbar.set_ticks(levels) 

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

1082 

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

1084 # with less than 30 points. 

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

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

1087 

1088 # Save plot. 

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

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

1091 plt.close(fig) 

1092 

1093 

1094def _plot_and_save_histogram_series( 

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

1096 filename: str, 

1097 title: str, 

1098 vmin: float, 

1099 vmax: float, 

1100 **kwargs, 

1101): 

1102 """Plot and save a histogram series. 

1103 

1104 Parameters 

1105 ---------- 

1106 cubes: Cube or CubeList 

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

1108 filename: str 

1109 Filename of the plot to write. 

1110 title: str 

1111 Plot title. 

1112 vmin: float 

1113 minimum for colorbar 

1114 vmax: float 

1115 maximum for colorbar 

1116 """ 

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

1118 ax = plt.gca() 

1119 

1120 model_colors_map = _get_model_colors_map(cubes) 

1121 

1122 # Set default that histograms will produce probability density function 

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

1124 density = True 

1125 

1126 for cube in iter_maybe(cubes): 

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

1128 # than seeing if long names exist etc. 

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

1130 if "surface_microphysical" in title: 

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

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

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

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

1135 density = False 

1136 else: 

1137 bins = 10.0 ** ( 

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

1139 ) # Suggestion from RMED toolbox. 

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

1141 ax.set_yscale("log") 

1142 vmin = bins[1] 

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

1144 ax.set_xscale("log") 

1145 elif "lightning" in title: 

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

1147 else: 

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

1149 logging.debug( 

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

1151 np.size(bins), 

1152 np.min(bins), 

1153 np.max(bins), 

1154 ) 

1155 

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

1157 # Otherwise we plot xdim histograms stacked. 

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

1159 

1160 label = None 

1161 color = "black" 

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

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

1164 color = model_colors_map[label] 

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

1166 

1167 # Compute area under curve. 

1168 if "surface_microphysical" in title and "amount" in title: 1168 ↛ 1169line 1168 didn't jump to line 1169 because the condition on line 1168 was never true

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

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

1171 x = x[1:] 

1172 y = y[1:] 

1173 

1174 ax.plot( 

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

1176 ) 

1177 

1178 # Add some labels and tweak the style. 

1179 ax.set_title(title, fontsize=16) 

1180 ax.set_xlabel( 

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

1182 ) 

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

1184 if "surface_microphysical" in title and "amount" in title: 1184 ↛ 1185line 1184 didn't jump to line 1185 because the condition on line 1184 was never true

1185 ax.set_ylabel( 

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

1187 ) 

1188 ax.set_xlim(vmin, vmax) 

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

1190 

1191 # Overlay grid-lines onto histogram plot. 

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

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

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

1195 

1196 # Save plot. 

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

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

1199 plt.close(fig) 

1200 

1201 

1202def _plot_and_save_postage_stamp_histogram_series( 

1203 cube: iris.cube.Cube, 

1204 filename: str, 

1205 title: str, 

1206 stamp_coordinate: str, 

1207 vmin: float, 

1208 vmax: float, 

1209 **kwargs, 

1210): 

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

1212 

1213 Parameters 

1214 ---------- 

1215 cube: Cube 

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

1217 filename: str 

1218 Filename of the plot to write. 

1219 title: str 

1220 Plot title. 

1221 stamp_coordinate: str 

1222 Coordinate that becomes different plots. 

1223 vmin: float 

1224 minimum for pdf x-axis 

1225 vmax: float 

1226 maximum for pdf x-axis 

1227 """ 

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

1229 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points)))) 

1230 

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

1232 # Make a subplot for each member. 

1233 for member, subplot in zip( 

1234 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False 

1235 ): 

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

1237 # cartopy GeoAxes generated. 

1238 plt.subplot(grid_size, grid_size, subplot) 

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

1240 # Otherwise we plot xdim histograms stacked. 

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

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

1243 ax = plt.gca() 

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

1245 ax.set_xlim(vmin, vmax) 

1246 

1247 # Overall figure title. 

1248 fig.suptitle(title, fontsize=16) 

1249 

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

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

1252 plt.close(fig) 

1253 

1254 

1255def _plot_and_save_postage_stamps_in_single_plot_histogram_series( 

1256 cube: iris.cube.Cube, 

1257 filename: str, 

1258 title: str, 

1259 stamp_coordinate: str, 

1260 vmin: float, 

1261 vmax: float, 

1262 **kwargs, 

1263): 

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

1265 ax.set_title(title, fontsize=16) 

1266 ax.set_xlim(vmin, vmax) 

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

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

1269 # Loop over all slices along the stamp_coordinate 

1270 for member in cube.slices_over(stamp_coordinate): 

1271 # Flatten the member data to 1D 

1272 member_data_1d = member.data.flatten() 

1273 # Plot the histogram using plt.hist 

1274 plt.hist( 

1275 member_data_1d, 

1276 density=True, 

1277 stacked=True, 

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

1279 ) 

1280 

1281 # Add a legend 

1282 ax.legend(fontsize=16) 

1283 

1284 # Save the figure to a file 

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

1286 

1287 # Close the figure 

1288 plt.close(fig) 

1289 

1290 

1291def _plot_and_save_power_spectrum_series( 

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

1293 filename: str, 

1294 title: str, 

1295 **kwargs, 

1296): 

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

1298 

1299 Parameters 

1300 ---------- 

1301 cubes: Cube or CubeList 

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

1303 filename: str 

1304 Filename of the plot to write. 

1305 title: str 

1306 Plot title. 

1307 """ 

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

1309 ax = plt.gca() 

1310 

1311 model_colors_map = _get_model_colors_map(cubes) 

1312 

1313 for cube in iter_maybe(cubes): 

1314 # Calculate power spectrum 

1315 

1316 # Extract time coordinate and convert to datetime 

1317 time_coord = cube.coord("time") 

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

1319 

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

1321 target_time = time_points[0] 

1322 

1323 # Bind target_time inside the lambda using a default argument 

1324 time_constraint = iris.Constraint( 

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

1326 ) 

1327 

1328 cube = cube.extract(time_constraint) 

1329 

1330 if cube.ndim == 2: 

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

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

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

1334 cube_3d = cube.data 

1335 else: 

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

1337 raise ValueError( 

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

1339 ) 

1340 

1341 # Calculate spectra 

1342 ps_array = _DCT_ps(cube_3d) 

1343 

1344 ps_cube = iris.cube.Cube( 

1345 ps_array, 

1346 long_name="power_spectra", 

1347 ) 

1348 

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

1350 

1351 # Create a frequency/wavelength array for coordinate 

1352 ps_len = ps_cube.data.shape[1] 

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

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

1355 

1356 # Convert datetime to numeric time using original units 

1357 numeric_time = time_coord.units.date2num(time_points) 

1358 # Create a new DimCoord with numeric time 

1359 new_time_coord = iris.coords.DimCoord( 

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

1361 ) 

1362 

1363 # Add time and frequency coordinate to spectra cube. 

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

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

1366 

1367 # Extract data from the cube 

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

1369 power_spectrum = ps_cube.data 

1370 

1371 label = None 

1372 color = "black" 

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

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

1375 color = model_colors_map[label] 

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

1377 

1378 # Add some labels and tweak the style. 

1379 ax.set_title(title, fontsize=16) 

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

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

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

1383 

1384 # Set log-log scale 

1385 ax.set_xscale("log") 

1386 ax.set_yscale("log") 

1387 

1388 # Overlay grid-lines onto power spectrum plot. 

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

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

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

1392 

1393 # Save plot. 

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

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

1396 plt.close(fig) 

1397 

1398 

1399def _plot_and_save_postage_stamp_power_spectrum_series( 

1400 cube: iris.cube.Cube, 

1401 filename: str, 

1402 title: str, 

1403 stamp_coordinate: str, 

1404 **kwargs, 

1405): 

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

1407 

1408 Parameters 

1409 ---------- 

1410 cube: Cube 

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

1412 filename: str 

1413 Filename of the plot to write. 

1414 title: str 

1415 Plot title. 

1416 stamp_coordinate: str 

1417 Coordinate that becomes different plots. 

1418 """ 

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

1420 grid_size = int(math.ceil(math.sqrt(len(cube.coord(stamp_coordinate).points)))) 

1421 

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

1423 # Make a subplot for each member. 

1424 for member, subplot in zip( 

1425 cube.slices_over(stamp_coordinate), range(1, grid_size**2 + 1), strict=False 

1426 ): 

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

1428 # cartopy GeoAxes generated. 

1429 plt.subplot(grid_size, grid_size, subplot) 

1430 

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

1432 power_spectrum = member.data 

1433 

1434 ax = plt.gca() 

1435 ax.plot(frequency, power_spectrum[0]) 

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

1437 

1438 # Overall figure title. 

1439 fig.suptitle(title, fontsize=16) 

1440 

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

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

1443 plt.close(fig) 

1444 

1445 

1446def _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series( 

1447 cube: iris.cube.Cube, 

1448 filename: str, 

1449 title: str, 

1450 stamp_coordinate: str, 

1451 **kwargs, 

1452): 

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

1454 ax.set_title(title, fontsize=16) 

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

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

1457 # Loop over all slices along the stamp_coordinate 

1458 for member in cube.slices_over(stamp_coordinate): 

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

1460 power_spectrum = member.data 

1461 ax.plot( 

1462 frequency, 

1463 power_spectrum[0], 

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

1465 ) 

1466 

1467 # Add a legend 

1468 ax.legend(fontsize=16) 

1469 

1470 # Save the figure to a file 

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

1472 

1473 # Close the figure 

1474 plt.close(fig) 

1475 

1476 

1477def _spatial_plot( 

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

1479 cube: iris.cube.Cube, 

1480 filename: str | None, 

1481 sequence_coordinate: str, 

1482 stamp_coordinate: str, 

1483): 

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

1485 

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

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

1488 is present then postage stamp plots will be produced. 

1489 

1490 Parameters 

1491 ---------- 

1492 method: "contourf" | "pcolormesh" 

1493 The plotting method to use. 

1494 cube: Cube 

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

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

1497 plotted sequentially and/or as postage stamp plots. 

1498 filename: str | None 

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

1500 uses the recipe name. 

1501 sequence_coordinate: str 

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

1503 This coordinate must exist in the cube. 

1504 stamp_coordinate: str 

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

1506 ``"realization"``. 

1507 

1508 Raises 

1509 ------ 

1510 ValueError 

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

1512 TypeError 

1513 If the cube isn't a single cube. 

1514 """ 

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

1516 

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

1518 if filename is None: 

1519 filename = slugify(recipe_title) 

1520 

1521 # Ensure we've got a single cube. 

1522 cube = _check_single_cube(cube) 

1523 

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

1525 # single point. 

1526 plotting_func = _plot_and_save_spatial_plot 

1527 try: 

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

1529 plotting_func = _plot_and_save_postage_stamp_spatial_plot 

1530 except iris.exceptions.CoordinateNotFoundError: 

1531 pass 

1532 

1533 # Must have a sequence coordinate. 

1534 try: 

1535 cube.coord(sequence_coordinate) 

1536 except iris.exceptions.CoordinateNotFoundError as err: 

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

1538 

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

1540 plot_index = [] 

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

1542 for cube_slice in cube.slices_over(sequence_coordinate): 

1543 # Use sequence value so multiple sequences can merge. 

1544 sequence_value = cube_slice.coord(sequence_coordinate).points[0] 

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

1546 coord = cube_slice.coord(sequence_coordinate) 

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

1548 title = f"{recipe_title}\n [{coord.units.title(coord.points[0])}]" 

1549 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs 

1550 if nplot == 1 and coord.has_bounds: 

1551 if np.size(coord.bounds) > 1: 

1552 title = f"{recipe_title}\n [{coord.units.title(coord.bounds[0][0])} to {coord.units.title(coord.bounds[0][1])}]" 

1553 # Do the actual plotting. 

1554 plotting_func( 

1555 cube_slice, 

1556 filename=plot_filename, 

1557 stamp_coordinate=stamp_coordinate, 

1558 title=title, 

1559 method=method, 

1560 ) 

1561 plot_index.append(plot_filename) 

1562 

1563 # Add list of plots to plot metadata. 

1564 complete_plot_index = _append_to_plot_index(plot_index) 

1565 

1566 # Make a page to display the plots. 

1567 _make_plot_html_page(complete_plot_index) 

1568 

1569 

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

1571 """Get colourmap for mask. 

1572 

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

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

1575 

1576 Parameters 

1577 ---------- 

1578 cube: Cube 

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

1580 axis: "x", "y", optional 

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

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

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

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

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

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

1587 

1588 Returns 

1589 ------- 

1590 cmap: 

1591 Matplotlib colormap. 

1592 levels: 

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

1594 should be taken as the range. 

1595 norm: 

1596 BoundaryNorm information. 

1597 """ 

1598 if "difference" not in cube.long_name: 

1599 if axis: 

1600 levels = [0, 1] 

1601 # Complete settings based on levels. 

1602 return None, levels, None 

1603 else: 

1604 # Define the levels and colors. 

1605 levels = [0, 1, 2] 

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

1607 # Create a custom color map. 

1608 cmap = mcolors.ListedColormap(colors) 

1609 # Normalize the levels. 

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

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

1612 return cmap, levels, norm 

1613 else: 

1614 if axis: 

1615 levels = [-1, 1] 

1616 return None, levels, None 

1617 else: 

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

1619 # not <=. 

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

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

1622 cmap = mcolors.ListedColormap(colors) 

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

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

1625 return cmap, levels, norm 

1626 

1627 

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

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

1630 

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

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

1633 

1634 Parameters 

1635 ---------- 

1636 cube: Cube 

1637 Cube of variable with Beaufort Scale in name. 

1638 axis: "x", "y", optional 

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

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

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

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

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

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

1645 

1646 Returns 

1647 ------- 

1648 cmap: 

1649 Matplotlib colormap. 

1650 levels: 

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

1652 should be taken as the range. 

1653 norm: 

1654 BoundaryNorm information. 

1655 """ 

1656 if "difference" not in cube.long_name: 

1657 if axis: 

1658 levels = [0, 12] 

1659 return None, levels, None 

1660 else: 

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

1662 colors = [ 

1663 "black", 

1664 (0, 0, 0.6), 

1665 "blue", 

1666 "cyan", 

1667 "green", 

1668 "yellow", 

1669 (1, 0.5, 0), 

1670 "red", 

1671 "pink", 

1672 "magenta", 

1673 "purple", 

1674 "maroon", 

1675 "white", 

1676 ] 

1677 cmap = mcolors.ListedColormap(colors) 

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

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

1680 return cmap, levels, norm 

1681 else: 

1682 if axis: 

1683 levels = [-4, 4] 

1684 return None, levels, None 

1685 else: 

1686 levels = [ 

1687 -3.5, 

1688 -2.5, 

1689 -1.5, 

1690 -0.5, 

1691 0.5, 

1692 1.5, 

1693 2.5, 

1694 3.5, 

1695 ] 

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

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

1698 return cmap, levels, norm 

1699 

1700 

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

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

1703 

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

1705 

1706 Parameters 

1707 ---------- 

1708 cube: Cube 

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

1710 cmap: Matplotlib colormap. 

1711 levels: List 

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

1713 should be taken as the range. 

1714 norm: BoundaryNorm. 

1715 

1716 Returns 

1717 ------- 

1718 cmap: Matplotlib colormap. 

1719 levels: List 

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

1721 should be taken as the range. 

1722 norm: BoundaryNorm. 

1723 """ 

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

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

1726 levels = np.array(levels) 

1727 levels -= 273 

1728 levels = levels.tolist() 

1729 else: 

1730 # Do nothing keep the existing colourbar attributes 

1731 levels = levels 

1732 cmap = cmap 

1733 norm = norm 

1734 return cmap, levels, norm 

1735 

1736 

1737def _custom_colormap_probability( 

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

1739): 

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

1741 

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

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

1744 

1745 Parameters 

1746 ---------- 

1747 cube: Cube 

1748 Cube of variable with probability in name. 

1749 axis: "x", "y", optional 

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

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

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

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

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

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

1756 

1757 Returns 

1758 ------- 

1759 cmap: 

1760 Matplotlib colormap. 

1761 levels: 

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

1763 should be taken as the range. 

1764 norm: 

1765 BoundaryNorm information. 

1766 """ 

1767 if axis: 

1768 levels = [0, 1] 

1769 return None, levels, None 

1770 else: 

1771 cmap = mcolors.ListedColormap( 

1772 [ 

1773 "#FFFFFF", 

1774 "#636363", 

1775 "#e1dada", 

1776 "#B5CAFF", 

1777 "#8FB3FF", 

1778 "#7F97FF", 

1779 "#ABCF63", 

1780 "#E8F59E", 

1781 "#FFFA14", 

1782 "#FFD121", 

1783 "#FFA30A", 

1784 ] 

1785 ) 

1786 levels = [0.0, 0.01, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] 

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

1788 return cmap, levels, norm 

1789 

1790 

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

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

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

1794 if ( 

1795 any("surface_microphysical" in name for name in varnames) 

1796 and "difference" not in cube.long_name 

1797 and "mask" not in cube.long_name 

1798 ): 

1799 # Define the levels and colors 

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

1801 colors = [ 

1802 "w", 

1803 (0, 0, 0.6), 

1804 "b", 

1805 "c", 

1806 "g", 

1807 "y", 

1808 (1, 0.5, 0), 

1809 "r", 

1810 "pink", 

1811 "m", 

1812 "purple", 

1813 "maroon", 

1814 "gray", 

1815 ] 

1816 # Create a custom colormap 

1817 cmap = mcolors.ListedColormap(colors) 

1818 # Normalize the levels 

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

1820 logging.info("change colormap for surface_microphysical variable colorbar.") 

1821 else: 

1822 # do nothing and keep existing colorbar attributes 

1823 cmap = cmap 

1824 levels = levels 

1825 norm = norm 

1826 return cmap, levels, norm 

1827 

1828 

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

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

1831 

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

1833 this function will be called. 

1834 

1835 Parameters 

1836 ---------- 

1837 cube: Cube 

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

1839 

1840 Returns 

1841 ------- 

1842 cmap: Matplotlib colormap. 

1843 levels: List 

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

1845 should be taken as the range. 

1846 norm: BoundaryNorm. 

1847 """ 

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

1849 colors = [ 

1850 "#87ceeb", 

1851 "#ffffff", 

1852 "#8ced69", 

1853 "#ffff00", 

1854 "#ffd700", 

1855 "#ffa500", 

1856 "#fe3620", 

1857 ] 

1858 # Create a custom colormap 

1859 cmap = mcolors.ListedColormap(colors) 

1860 # Normalise the levels 

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

1862 return cmap, levels, norm 

1863 

1864 

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

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

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

1868 if ( 

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

1870 and "difference" not in cube.long_name 

1871 and "mask" not in cube.long_name 

1872 ): 

1873 # Define the levels and colors (in km) 

1874 levels = [0, 0.05, 0.1, 0.2, 1.0, 2.0, 5.0, 10.0, 20.0, 30.0, 50.0, 70.0, 100.0] 

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

1876 colours = [ 

1877 "#8f00d6", 

1878 "#d10000", 

1879 "#ff9700", 

1880 "#ffff00", 

1881 "#00007f", 

1882 "#6c9ccd", 

1883 "#aae8ff", 

1884 "#37a648", 

1885 "#8edc64", 

1886 "#c5ffc5", 

1887 "#dcdcdc", 

1888 "#ffffff", 

1889 ] 

1890 # Create a custom colormap 

1891 cmap = mcolors.ListedColormap(colours) 

1892 # Normalize the levels 

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

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

1895 else: 

1896 # do nothing and keep existing colorbar attributes 

1897 cmap = cmap 

1898 levels = levels 

1899 norm = norm 

1900 return cmap, levels, norm 

1901 

1902 

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

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

1905 model_names = list( 

1906 filter( 

1907 lambda x: x is not None, 

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

1909 ) 

1910 ) 

1911 if not model_names: 

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

1913 return 1 

1914 else: 

1915 return len(model_names) 

1916 

1917 

1918def _validate_cube_shape( 

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

1920) -> None: 

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

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

1923 raise ValueError( 

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

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

1926 ) 

1927 

1928 

1929def _validate_cubes_coords( 

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

1931) -> None: 

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

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

1934 raise ValueError( 

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

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

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

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

1939 ) 

1940 

1941 

1942#################### 

1943# Public functions # 

1944#################### 

1945 

1946 

1947def spatial_contour_plot( 

1948 cube: iris.cube.Cube, 

1949 filename: str = None, 

1950 sequence_coordinate: str = "time", 

1951 stamp_coordinate: str = "realization", 

1952 **kwargs, 

1953) -> iris.cube.Cube: 

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

1955 

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

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

1958 is present then postage stamp plots will be produced. 

1959 

1960 Parameters 

1961 ---------- 

1962 cube: Cube 

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

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

1965 plotted sequentially and/or as postage stamp plots. 

1966 filename: str, optional 

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

1968 to the recipe name. 

1969 sequence_coordinate: str, optional 

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

1971 This coordinate must exist in the cube. 

1972 stamp_coordinate: str, optional 

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

1974 ``"realization"``. 

1975 

1976 Returns 

1977 ------- 

1978 Cube 

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

1980 

1981 Raises 

1982 ------ 

1983 ValueError 

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

1985 TypeError 

1986 If the cube isn't a single cube. 

1987 """ 

1988 _spatial_plot("contourf", cube, filename, sequence_coordinate, stamp_coordinate) 

1989 return cube 

1990 

1991 

1992def spatial_pcolormesh_plot( 

1993 cube: iris.cube.Cube, 

1994 filename: str = None, 

1995 sequence_coordinate: str = "time", 

1996 stamp_coordinate: str = "realization", 

1997 **kwargs, 

1998) -> iris.cube.Cube: 

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

2000 

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

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

2003 is present then postage stamp plots will be produced. 

2004 

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

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

2007 contour areas are important. 

2008 

2009 Parameters 

2010 ---------- 

2011 cube: Cube 

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

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

2014 plotted sequentially and/or as postage stamp plots. 

2015 filename: str, optional 

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

2017 to the recipe name. 

2018 sequence_coordinate: str, optional 

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

2020 This coordinate must exist in the cube. 

2021 stamp_coordinate: str, optional 

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

2023 ``"realization"``. 

2024 

2025 Returns 

2026 ------- 

2027 Cube 

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

2029 

2030 Raises 

2031 ------ 

2032 ValueError 

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

2034 TypeError 

2035 If the cube isn't a single cube. 

2036 """ 

2037 _spatial_plot("pcolormesh", cube, filename, sequence_coordinate, stamp_coordinate) 

2038 return cube 

2039 

2040 

2041# TODO: Expand function to handle ensemble data. 

2042# line_coordinate: str, optional 

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

2044# ``"realization"``. 

2045def plot_line_series( 

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

2047 filename: str = None, 

2048 series_coordinate: str = "time", 

2049 # line_coordinate: str = "realization", 

2050 **kwargs, 

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

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

2053 

2054 The Cube or CubeList must be 1D. 

2055 

2056 Parameters 

2057 ---------- 

2058 iris.cube | iris.cube.CubeList 

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

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

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

2062 filename: str, optional 

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

2064 to the recipe name. 

2065 series_coordinate: str, optional 

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

2067 coordinate must exist in the cube. 

2068 

2069 Returns 

2070 ------- 

2071 iris.cube.Cube | iris.cube.CubeList 

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

2073 plotted data. 

2074 

2075 Raises 

2076 ------ 

2077 ValueError 

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

2079 TypeError 

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

2081 """ 

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

2083 title = get_recipe_metadata().get("title", "Untitled") 

2084 

2085 if filename is None: 

2086 filename = slugify(title) 

2087 

2088 # Add file extension. 

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

2090 

2091 num_models = _get_num_models(cube) 

2092 

2093 _validate_cube_shape(cube, num_models) 

2094 

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

2096 cubes = iter_maybe(cube) 

2097 coords = [] 

2098 for cube in cubes: 

2099 try: 

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

2101 except iris.exceptions.CoordinateNotFoundError as err: 

2102 raise ValueError( 

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

2104 ) from err 

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

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

2107 

2108 # Do the actual plotting. 

2109 _plot_and_save_line_series(cubes, coords, "realization", plot_filename, title) 

2110 

2111 # Add list of plots to plot metadata. 

2112 plot_index = _append_to_plot_index([plot_filename]) 

2113 

2114 # Make a page to display the plots. 

2115 _make_plot_html_page(plot_index) 

2116 

2117 return cube 

2118 

2119 

2120def plot_vertical_line_series( 

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

2122 filename: str = None, 

2123 series_coordinate: str = "model_level_number", 

2124 sequence_coordinate: str = "time", 

2125 # line_coordinate: str = "realization", 

2126 **kwargs, 

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

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

2129 

2130 The Cube or CubeList must be 1D. 

2131 

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

2133 then a sequence of plots will be produced. 

2134 

2135 Parameters 

2136 ---------- 

2137 iris.cube | iris.cube.CubeList 

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

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

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

2141 filename: str, optional 

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

2143 to the recipe name. 

2144 series_coordinate: str, optional 

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

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

2147 for LFRic. Defaults to ``model_level_number``. 

2148 This coordinate must exist in the cube. 

2149 sequence_coordinate: str, optional 

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

2151 This coordinate must exist in the cube. 

2152 

2153 Returns 

2154 ------- 

2155 iris.cube.Cube | iris.cube.CubeList 

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

2157 Plotted data. 

2158 

2159 Raises 

2160 ------ 

2161 ValueError 

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

2163 TypeError 

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

2165 """ 

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

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

2168 

2169 if filename is None: 

2170 filename = slugify(recipe_title) 

2171 

2172 cubes = iter_maybe(cubes) 

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

2174 all_data = [] 

2175 

2176 # Store min/max ranges for x range. 

2177 x_levels = [] 

2178 

2179 num_models = _get_num_models(cubes) 

2180 

2181 _validate_cube_shape(cubes, num_models) 

2182 

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

2184 coords = [] 

2185 for cube in cubes: 

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

2187 try: 

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

2189 except iris.exceptions.CoordinateNotFoundError as err: 

2190 raise ValueError( 

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

2192 ) from err 

2193 

2194 try: 

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

2196 cube.coord(sequence_coordinate) 

2197 except iris.exceptions.CoordinateNotFoundError as err: 

2198 raise ValueError( 

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

2200 ) from err 

2201 

2202 # Get minimum and maximum from levels information. 

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

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

2205 x_levels.append(min(levels)) 

2206 x_levels.append(max(levels)) 

2207 else: 

2208 all_data.append(cube.data) 

2209 

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

2211 # Combine all data into a single NumPy array 

2212 combined_data = np.concatenate(all_data) 

2213 

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

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

2216 # sequence and if applicable postage stamp coordinate. 

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

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

2219 else: 

2220 vmin = min(x_levels) 

2221 vmax = max(x_levels) 

2222 

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

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

2225 # necessary) 

2226 def filter_cube_iterables(cube_iterables) -> bool: 

2227 return len(cube_iterables) == len(coords) 

2228 

2229 cube_iterables = filter( 

2230 filter_cube_iterables, 

2231 ( 

2232 iris.cube.CubeList( 

2233 s 

2234 for s in itertools.chain.from_iterable( 

2235 cb.slices_over(sequence_coordinate) for cb in cubes 

2236 ) 

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

2238 ) 

2239 for point in sorted( 

2240 set( 

2241 itertools.chain.from_iterable( 

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

2243 ) 

2244 ) 

2245 ) 

2246 ), 

2247 ) 

2248 

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

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

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

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

2253 # or an iris.cube.CubeList. 

2254 plot_index = [] 

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

2256 for cubes_slice in cube_iterables: 

2257 # Use sequence value so multiple sequences can merge. 

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

2259 sequence_value = seq_coord.points[0] 

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

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

2262 title = f"{recipe_title}\n [{seq_coord.units.title(sequence_value)}]" 

2263 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs 

2264 if nplot == 1 and seq_coord.has_bounds: 2264 ↛ 2265line 2264 didn't jump to line 2265 because the condition on line 2264 was never true

2265 if np.size(seq_coord.bounds) > 1: 

2266 title = f"{recipe_title}\n [{seq_coord.units.title(seq_coord.bounds[0][0])} to {seq_coord.units.title(seq_coord.bounds[0][1])}]" 

2267 # Do the actual plotting. 

2268 _plot_and_save_vertical_line_series( 

2269 cubes_slice, 

2270 coords, 

2271 "realization", 

2272 plot_filename, 

2273 series_coordinate, 

2274 title=title, 

2275 vmin=vmin, 

2276 vmax=vmax, 

2277 ) 

2278 plot_index.append(plot_filename) 

2279 

2280 # Add list of plots to plot metadata. 

2281 complete_plot_index = _append_to_plot_index(plot_index) 

2282 

2283 # Make a page to display the plots. 

2284 _make_plot_html_page(complete_plot_index) 

2285 

2286 return cubes 

2287 

2288 

2289def qq_plot( 

2290 cubes: iris.cube.CubeList, 

2291 coordinates: list[str], 

2292 percentiles: list[float], 

2293 model_names: list[str], 

2294 filename: str = None, 

2295 one_to_one: bool = True, 

2296 **kwargs, 

2297) -> iris.cube.CubeList: 

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

2299 

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

2301 collapsed within the operator over all specified coordinates such as 

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

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

2304 

2305 Parameters 

2306 ---------- 

2307 cubes: iris.cube.CubeList 

2308 Two cubes of the same variable with different models. 

2309 coordinate: list[str] 

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

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

2312 the percentile coordinate. 

2313 percent: list[float] 

2314 A list of percentiles to appear in the plot. 

2315 model_names: list[str] 

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

2317 filename: str, optional 

2318 Filename of the plot to write. 

2319 one_to_one: bool, optional 

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

2321 

2322 Raises 

2323 ------ 

2324 ValueError 

2325 When the cubes are not compatible. 

2326 

2327 Notes 

2328 ----- 

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

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

2331 This plot does not use a theoretical distribution to compare against, but 

2332 compares percentiles of two datasets. This plot does 

2333 not use all raw data points, but plots the selected percentiles (quantiles) of 

2334 each variable instead for the two datasets, thereby normalising the data for a 

2335 direct comparison between the selected percentiles of the two dataset distributions. 

2336 

2337 Quantile-quantile plots are valuable for comparing against 

2338 observations and other models. Identical percentiles between the variables 

2339 will lie on the one-to-one line implying the values correspond well to each 

2340 other. Where there is a deviation from the one-to-one line a range of 

2341 possibilities exist depending on how and where the data is shifted (e.g., 

2342 Wilks 2011 [Wilks2011]_). 

2343 

2344 For distributions above the one-to-one line the distribution is left-skewed; 

2345 below is right-skewed. A distinct break implies a bimodal distribution, and 

2346 closer values/values further apart at the tails imply poor representation of 

2347 the extremes. 

2348 

2349 References 

2350 ---------- 

2351 .. [Wilks2011] Wilks, D.S., (2011) "Statistical Methods in the Atmospheric 

2352 Sciences" Third Edition, vol. 100, Academic Press, Oxford, UK, 676 pp. 

2353 """ 

2354 # Check cubes using same functionality as the difference operator. 

2355 if len(cubes) != 2: 

2356 raise ValueError("cubes should contain exactly 2 cubes.") 

2357 base: Cube = cubes.extract_cube(iris.AttributeConstraint(cset_comparison_base=1)) 

2358 other: Cube = cubes.extract_cube( 

2359 iris.Constraint( 

2360 cube_func=lambda cube: "cset_comparison_base" not in cube.attributes 

2361 ) 

2362 ) 

2363 

2364 # Get spatial coord names. 

2365 base_lat_name, base_lon_name = get_cube_yxcoordname(base) 

2366 other_lat_name, other_lon_name = get_cube_yxcoordname(other) 

2367 

2368 # Ensure cubes to compare are on common differencing grid. 

2369 # This is triggered if either 

2370 # i) latitude and longitude shapes are not the same. Note grid points 

2371 # are not compared directly as these can differ through rounding 

2372 # errors. 

2373 # ii) or variables are known to often sit on different grid staggering 

2374 # in different models (e.g. cell center vs cell edge), as is the case 

2375 # for UM and LFRic comparisons. 

2376 # In future greater choice of regridding method might be applied depending 

2377 # on variable type. Linear regridding can in general be appropriate for smooth 

2378 # variables. Care should be taken with interpretation of differences 

2379 # given this dependency on regridding. 

2380 if ( 

2381 base.coord(base_lat_name).shape != other.coord(other_lat_name).shape 

2382 or base.coord(base_lon_name).shape != other.coord(other_lon_name).shape 

2383 ) or ( 

2384 base.long_name 

2385 in [ 

2386 "eastward_wind_at_10m", 

2387 "northward_wind_at_10m", 

2388 "northward_wind_at_cell_centres", 

2389 "eastward_wind_at_cell_centres", 

2390 "zonal_wind_at_pressure_levels", 

2391 "meridional_wind_at_pressure_levels", 

2392 "potential_vorticity_at_pressure_levels", 

2393 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging", 

2394 ] 

2395 ): 

2396 logging.debug( 

2397 "Linear regridding base cube to other grid to compute differences" 

2398 ) 

2399 base = regrid_onto_cube(base, other, method="Linear") 

2400 

2401 # Extract just common time points. 

2402 base, other = _extract_common_time_points(base, other) 

2403 

2404 # Equalise attributes so we can merge. 

2405 fully_equalise_attributes([base, other]) 

2406 logging.debug("Base: %s\nOther: %s", base, other) 

2407 

2408 # Collapse cubes. 

2409 base = collapse( 

2410 base, 

2411 coordinate=coordinates, 

2412 method="PERCENTILE", 

2413 additional_percent=percentiles, 

2414 ) 

2415 other = collapse( 

2416 other, 

2417 coordinate=coordinates, 

2418 method="PERCENTILE", 

2419 additional_percent=percentiles, 

2420 ) 

2421 

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

2423 title = get_recipe_metadata().get("title", "Untitled") 

2424 

2425 if filename is None: 

2426 filename = slugify(title) 

2427 

2428 # Add file extension. 

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

2430 

2431 # Do the actual plotting on a scatter plot 

2432 _plot_and_save_scatter_plot( 

2433 base, other, plot_filename, title, one_to_one, model_names 

2434 ) 

2435 

2436 # Add list of plots to plot metadata. 

2437 plot_index = _append_to_plot_index([plot_filename]) 

2438 

2439 # Make a page to display the plots. 

2440 _make_plot_html_page(plot_index) 

2441 

2442 return iris.cube.CubeList([base, other]) 

2443 

2444 

2445def scatter_plot( 

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

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

2448 filename: str = None, 

2449 one_to_one: bool = True, 

2450 **kwargs, 

2451) -> iris.cube.CubeList: 

2452 """Plot a scatter plot between two variables. 

2453 

2454 Both cubes must be 1D. 

2455 

2456 Parameters 

2457 ---------- 

2458 cube_x: Cube | CubeList 

2459 1 dimensional Cube of the data to plot on y-axis. 

2460 cube_y: Cube | CubeList 

2461 1 dimensional Cube of the data to plot on x-axis. 

2462 filename: str, optional 

2463 Filename of the plot to write. 

2464 one_to_one: bool, optional 

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

2466 

2467 Returns 

2468 ------- 

2469 cubes: CubeList 

2470 CubeList of the original x and y cubes for further processing. 

2471 

2472 Raises 

2473 ------ 

2474 ValueError 

2475 If the cube doesn't have the right dimensions and cubes not the same 

2476 size. 

2477 TypeError 

2478 If the cube isn't a single cube. 

2479 

2480 Notes 

2481 ----- 

2482 Scatter plots are used for determining if there is a relationship between 

2483 two variables. Positive relations have a slope going from bottom left to top 

2484 right; Negative relations have a slope going from top left to bottom right. 

2485 """ 

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

2487 for cube_iter in iter_maybe(cube_x): 

2488 # Check cubes are correct shape. 

2489 cube_iter = _check_single_cube(cube_iter) 

2490 if cube_iter.ndim > 1: 

2491 raise ValueError("cube_x must be 1D.") 

2492 

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

2494 for cube_iter in iter_maybe(cube_y): 

2495 # Check cubes are correct shape. 

2496 cube_iter = _check_single_cube(cube_iter) 

2497 if cube_iter.ndim > 1: 

2498 raise ValueError("cube_y must be 1D.") 

2499 

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

2501 title = get_recipe_metadata().get("title", "Untitled") 

2502 

2503 if filename is None: 

2504 filename = slugify(title) 

2505 

2506 # Add file extension. 

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

2508 

2509 # Do the actual plotting. 

2510 _plot_and_save_scatter_plot(cube_x, cube_y, plot_filename, title, one_to_one) 

2511 

2512 # Add list of plots to plot metadata. 

2513 plot_index = _append_to_plot_index([plot_filename]) 

2514 

2515 # Make a page to display the plots. 

2516 _make_plot_html_page(plot_index) 

2517 

2518 return iris.cube.CubeList([cube_x, cube_y]) 

2519 

2520 

2521def vector_plot( 

2522 cube_u: iris.cube.Cube, 

2523 cube_v: iris.cube.Cube, 

2524 filename: str = None, 

2525 sequence_coordinate: str = "time", 

2526 **kwargs, 

2527) -> iris.cube.CubeList: 

2528 """Plot a vector plot based on the input u and v components.""" 

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

2530 

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

2532 if filename is None: 2532 ↛ 2533line 2532 didn't jump to line 2533 because the condition on line 2532 was never true

2533 filename = slugify(recipe_title) 

2534 

2535 # Cubes must have a matching sequence coordinate. 

2536 try: 

2537 # Check that the u and v cubes have the same sequence coordinate. 

2538 if cube_u.coord(sequence_coordinate) != cube_v.coord(sequence_coordinate): 2538 ↛ 2539line 2538 didn't jump to line 2539 because the condition on line 2538 was never true

2539 raise ValueError("Coordinates do not match.") 

2540 except (iris.exceptions.CoordinateNotFoundError, ValueError) as err: 

2541 raise ValueError( 

2542 f"Cubes should have matching {sequence_coordinate} coordinate:\n{cube_u}\n{cube_v}" 

2543 ) from err 

2544 

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

2546 plot_index = [] 

2547 for cube_u_slice, cube_v_slice in zip( 

2548 cube_u.slices_over(sequence_coordinate), 

2549 cube_v.slices_over(sequence_coordinate), 

2550 strict=True, 

2551 ): 

2552 # Use sequence value so multiple sequences can merge. 

2553 sequence_value = cube_u_slice.coord(sequence_coordinate).points[0] 

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

2555 coord = cube_u_slice.coord(sequence_coordinate) 

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

2557 title = f"{recipe_title}\n{coord.units.title(coord.points[0])}" 

2558 # Do the actual plotting. 

2559 _plot_and_save_vector_plot( 

2560 cube_u_slice, 

2561 cube_v_slice, 

2562 filename=plot_filename, 

2563 title=title, 

2564 method="contourf", 

2565 ) 

2566 plot_index.append(plot_filename) 

2567 

2568 # Add list of plots to plot metadata. 

2569 complete_plot_index = _append_to_plot_index(plot_index) 

2570 

2571 # Make a page to display the plots. 

2572 _make_plot_html_page(complete_plot_index) 

2573 

2574 return iris.cube.CubeList([cube_u, cube_v]) 

2575 

2576 

2577def plot_histogram_series( 

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

2579 filename: str = None, 

2580 sequence_coordinate: str = "time", 

2581 stamp_coordinate: str = "realization", 

2582 single_plot: bool = False, 

2583 **kwargs, 

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

2585 """Plot a histogram plot for each vertical level provided. 

2586 

2587 A histogram plot can be plotted, but if the sequence_coordinate (i.e. time) 

2588 is present then a sequence of plots will be produced using the time slider 

2589 functionality to scroll through histograms against time. If a 

2590 stamp_coordinate is present then postage stamp plots will be produced. If 

2591 stamp_coordinate and single_plot is True, all postage stamp plots will be 

2592 plotted in a single plot instead of separate postage stamp plots. 

2593 

2594 Parameters 

2595 ---------- 

2596 cubes: Cube | iris.cube.CubeList 

2597 Iris cube or CubeList of the data to plot. It should have a single dimension other 

2598 than the stamp coordinate. 

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

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

2601 filename: str, optional 

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

2603 to the recipe name. 

2604 sequence_coordinate: str, optional 

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

2606 This coordinate must exist in the cube and will be used for the time 

2607 slider. 

2608 stamp_coordinate: str, optional 

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

2610 ``"realization"``. 

2611 single_plot: bool, optional 

2612 If True, all postage stamp plots will be plotted in a single plot. If 

2613 False, each postage stamp plot will be plotted separately. Is only valid 

2614 if stamp_coordinate exists and has more than a single point. 

2615 

2616 Returns 

2617 ------- 

2618 iris.cube.Cube | iris.cube.CubeList 

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

2620 Plotted data. 

2621 

2622 Raises 

2623 ------ 

2624 ValueError 

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

2626 TypeError 

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

2628 """ 

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

2630 

2631 cubes = iter_maybe(cubes) 

2632 

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

2634 if filename is None: 

2635 filename = slugify(recipe_title) 

2636 

2637 # Internal plotting function. 

2638 plotting_func = _plot_and_save_histogram_series 

2639 

2640 num_models = _get_num_models(cubes) 

2641 

2642 _validate_cube_shape(cubes, num_models) 

2643 

2644 # If several histograms are plotted with time as sequence_coordinate for the 

2645 # time slider option. 

2646 for cube in cubes: 

2647 try: 

2648 cube.coord(sequence_coordinate) 

2649 except iris.exceptions.CoordinateNotFoundError as err: 

2650 raise ValueError( 

2651 f"Cube must have a {sequence_coordinate} coordinate." 

2652 ) from err 

2653 

2654 # Get minimum and maximum from levels information. 

2655 levels = None 

2656 for cube in cubes: 2656 ↛ 2672line 2656 didn't jump to line 2672 because the loop on line 2656 didn't complete

2657 # First check if user-specified "auto" range variable. 

2658 # This maintains the value of levels as None, so proceed. 

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

2660 if levels is None: 

2661 break 

2662 # If levels is changed, recheck to use the vmin,vmax or 

2663 # levels-based ranges for histogram plots. 

2664 _, levels, _ = _colorbar_map_levels(cube) 

2665 logging.debug("levels: %s", levels) 

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

2667 vmin = min(levels) 

2668 vmax = max(levels) 

2669 logging.debug("Updated vmin, vmax: %s, %s", vmin, vmax) 

2670 break 

2671 

2672 if levels is None: 

2673 vmin = min(cb.data.min() for cb in cubes) 

2674 vmax = max(cb.data.max() for cb in cubes) 

2675 

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

2677 # single point. If single_plot is True: 

2678 # -- all postage stamp plots will be plotted in a single plot instead of 

2679 # separate postage stamp plots. 

2680 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are 

2681 # produced per single model only 

2682 if num_models == 1: 2682 ↛ 2695line 2682 didn't jump to line 2695 because the condition on line 2682 was always true

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

2684 stamp_coordinate in [c.name() for c in cubes[0].coords()] 

2685 and cubes[0].coord(stamp_coordinate).shape[0] > 1 

2686 ): 

2687 if single_plot: 

2688 plotting_func = ( 

2689 _plot_and_save_postage_stamps_in_single_plot_histogram_series 

2690 ) 

2691 else: 

2692 plotting_func = _plot_and_save_postage_stamp_histogram_series 

2693 cube_iterables = cubes[0].slices_over(sequence_coordinate) 

2694 else: 

2695 all_points = sorted( 

2696 set( 

2697 itertools.chain.from_iterable( 

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

2699 ) 

2700 ) 

2701 ) 

2702 all_slices = list( 

2703 itertools.chain.from_iterable( 

2704 cb.slices_over(sequence_coordinate) for cb in cubes 

2705 ) 

2706 ) 

2707 # Matched slices (matched by seq coord point; it may happen that 

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

2709 # necessary) 

2710 cube_iterables = [ 

2711 iris.cube.CubeList( 

2712 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point 

2713 ) 

2714 for point in all_points 

2715 ] 

2716 

2717 plot_index = [] 

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

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

2720 # multiple cubes in a CubeList to be plotted in the same plot for similar 

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

2722 # for similar values of the sequence coordinate. cube_slice can be an 

2723 # iris.cube.Cube or an iris.cube.CubeList. 

2724 for cube_slice in cube_iterables: 

2725 single_cube = cube_slice 

2726 if isinstance(cube_slice, iris.cube.CubeList): 2726 ↛ 2727line 2726 didn't jump to line 2727 because the condition on line 2726 was never true

2727 single_cube = cube_slice[0] 

2728 

2729 # Use sequence value so multiple sequences can merge. 

2730 sequence_value = single_cube.coord(sequence_coordinate).points[0] 

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

2732 coord = single_cube.coord(sequence_coordinate) 

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

2734 title = f"{recipe_title}\n [{coord.units.title(coord.points[0])}]" 

2735 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs 

2736 if nplot == 1 and coord.has_bounds: 2736 ↛ 2737line 2736 didn't jump to line 2737 because the condition on line 2736 was never true

2737 if np.size(coord.bounds) > 1: 

2738 title = f"{recipe_title}\n [{coord.units.title(coord.bounds[0][0])} to {coord.units.title(coord.bounds[0][1])}]" 

2739 # Do the actual plotting. 

2740 plotting_func( 

2741 cube_slice, 

2742 filename=plot_filename, 

2743 stamp_coordinate=stamp_coordinate, 

2744 title=title, 

2745 vmin=vmin, 

2746 vmax=vmax, 

2747 ) 

2748 plot_index.append(plot_filename) 

2749 

2750 # Add list of plots to plot metadata. 

2751 complete_plot_index = _append_to_plot_index(plot_index) 

2752 

2753 # Make a page to display the plots. 

2754 _make_plot_html_page(complete_plot_index) 

2755 

2756 return cubes 

2757 

2758 

2759def plot_power_spectrum_series( 

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

2761 filename: str = None, 

2762 sequence_coordinate: str = "time", 

2763 stamp_coordinate: str = "realization", 

2764 single_plot: bool = False, 

2765 **kwargs, 

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

2767 """Plot a power spectrum plot for each vertical level provided. 

2768 

2769 A power spectrum plot can be plotted, but if the sequence_coordinate (i.e. time) 

2770 is present then a sequence of plots will be produced using the time slider 

2771 functionality to scroll through power spectrum against time. If a 

2772 stamp_coordinate is present then postage stamp plots will be produced. If 

2773 stamp_coordinate and single_plot is True, all postage stamp plots will be 

2774 plotted in a single plot instead of separate postage stamp plots. 

2775 

2776 Parameters 

2777 ---------- 

2778 cubes: Cube | iris.cube.CubeList 

2779 Iris cube or CubeList of the data to plot. It should have a single dimension other 

2780 than the stamp coordinate. 

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

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

2783 filename: str, optional 

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

2785 to the recipe name. 

2786 sequence_coordinate: str, optional 

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

2788 This coordinate must exist in the cube and will be used for the time 

2789 slider. 

2790 stamp_coordinate: str, optional 

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

2792 ``"realization"``. 

2793 single_plot: bool, optional 

2794 If True, all postage stamp plots will be plotted in a single plot. If 

2795 False, each postage stamp plot will be plotted separately. Is only valid 

2796 if stamp_coordinate exists and has more than a single point. 

2797 

2798 Returns 

2799 ------- 

2800 iris.cube.Cube | iris.cube.CubeList 

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

2802 Plotted data. 

2803 

2804 Raises 

2805 ------ 

2806 ValueError 

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

2808 TypeError 

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

2810 """ 

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

2812 

2813 cubes = iter_maybe(cubes) 

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

2815 if filename is None: 

2816 filename = slugify(recipe_title) 

2817 

2818 # Internal plotting function. 

2819 plotting_func = _plot_and_save_power_spectrum_series 

2820 

2821 num_models = _get_num_models(cubes) 

2822 

2823 _validate_cube_shape(cubes, num_models) 

2824 

2825 # If several power spectra are plotted with time as sequence_coordinate for the 

2826 # time slider option. 

2827 for cube in cubes: 

2828 try: 

2829 cube.coord(sequence_coordinate) 

2830 except iris.exceptions.CoordinateNotFoundError as err: 

2831 raise ValueError( 

2832 f"Cube must have a {sequence_coordinate} coordinate." 

2833 ) from err 

2834 

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

2836 # single point. If single_plot is True: 

2837 # -- all postage stamp plots will be plotted in a single plot instead of 

2838 # separate postage stamp plots. 

2839 # -- model names (hidden in cube attrs) are ignored, that is stamp plots are 

2840 # produced per single model only 

2841 if num_models == 1: 2841 ↛ 2854line 2841 didn't jump to line 2854 because the condition on line 2841 was always true

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

2843 stamp_coordinate in [c.name() for c in cubes[0].coords()] 

2844 and cubes[0].coord(stamp_coordinate).shape[0] > 1 

2845 ): 

2846 if single_plot: 

2847 plotting_func = ( 

2848 _plot_and_save_postage_stamps_in_single_plot_power_spectrum_series 

2849 ) 

2850 else: 

2851 plotting_func = _plot_and_save_postage_stamp_power_spectrum_series 

2852 cube_iterables = cubes[0].slices_over(sequence_coordinate) 

2853 else: 

2854 all_points = sorted( 

2855 set( 

2856 itertools.chain.from_iterable( 

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

2858 ) 

2859 ) 

2860 ) 

2861 all_slices = list( 

2862 itertools.chain.from_iterable( 

2863 cb.slices_over(sequence_coordinate) for cb in cubes 

2864 ) 

2865 ) 

2866 # Matched slices (matched by seq coord point; it may happen that 

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

2868 # necessary) 

2869 cube_iterables = [ 

2870 iris.cube.CubeList( 

2871 s for s in all_slices if s.coord(sequence_coordinate).points[0] == point 

2872 ) 

2873 for point in all_points 

2874 ] 

2875 

2876 plot_index = [] 

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

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

2879 # multiple cubes in a CubeList to be plotted in the same plot for similar 

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

2881 # for similar values of the sequence coordinate. cube_slice can be an 

2882 # iris.cube.Cube or an iris.cube.CubeList. 

2883 for cube_slice in cube_iterables: 

2884 single_cube = cube_slice 

2885 if isinstance(cube_slice, iris.cube.CubeList): 2885 ↛ 2886line 2885 didn't jump to line 2886 because the condition on line 2885 was never true

2886 single_cube = cube_slice[0] 

2887 

2888 # Use sequence value so multiple sequences can merge. 

2889 sequence_value = single_cube.coord(sequence_coordinate).points[0] 

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

2891 coord = single_cube.coord(sequence_coordinate) 

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

2893 title = f"{recipe_title}\n [{coord.units.title(coord.points[0])}]" 

2894 # Use sequence (e.g. time) bounds if plotting single non-sequence outputs 

2895 if nplot == 1 and coord.has_bounds: 2895 ↛ 2899line 2895 didn't jump to line 2899 because the condition on line 2895 was always true

2896 if np.size(coord.bounds) > 1: 

2897 title = f"{recipe_title}\n [{coord.units.title(coord.bounds[0][0])} to {coord.units.title(coord.bounds[0][1])}]" 

2898 # Do the actual plotting. 

2899 plotting_func( 

2900 cube_slice, 

2901 filename=plot_filename, 

2902 stamp_coordinate=stamp_coordinate, 

2903 title=title, 

2904 ) 

2905 plot_index.append(plot_filename) 

2906 

2907 # Add list of plots to plot metadata. 

2908 complete_plot_index = _append_to_plot_index(plot_index) 

2909 

2910 # Make a page to display the plots. 

2911 _make_plot_html_page(complete_plot_index) 

2912 

2913 return cubes 

2914 

2915 

2916def _DCT_ps(y_3d): 

2917 """Calculate power spectra for regional domains. 

2918 

2919 Parameters 

2920 ---------- 

2921 y_3d: 3D array 

2922 3 dimensional array to calculate spectrum for. 

2923 (2D field data with 3rd dimension of time) 

2924 

2925 Returns: ps_array 

2926 Array of power spectra values calculated for input field (for each time) 

2927 

2928 Method for regional domains: 

2929 Calculate power spectra over limited area domain using Discrete Cosine Transform (DCT) 

2930 as described in Denis et al 2002 [Denis_etal_2002]_. 

2931 

2932 References 

2933 ---------- 

2934 .. [Denis_etal_2002] Bertrand Denis, Jean Côté and René Laprise (2002) 

2935 "Spectral Decomposition of Two-Dimensional Atmospheric Fields on 

2936 Limited-Area Domains Using the Discrete Cosine Transform (DCT)" 

2937 Monthly Weather Review, Vol. 130, 1812-1828 

2938 doi: https://doi.org/10.1175/1520-0493(2002)130<1812:SDOTDA>2.0.CO;2 

2939 """ 

2940 Nt, Ny, Nx = y_3d.shape 

2941 

2942 # Max coefficient 

2943 Nmin = min(Nx - 1, Ny - 1) 

2944 

2945 # Create alpha matrix (of wavenumbers) 

2946 alpha_matrix = _create_alpha_matrix(Ny, Nx) 

2947 

2948 # Prepare output array 

2949 ps_array = np.zeros((Nt, Nmin)) 

2950 

2951 # Loop over time to get spectrum for each time. 

2952 for t in range(Nt): 

2953 y_2d = y_3d[t] 

2954 

2955 # Apply 2D DCT to transform y_3d[t] from physical space to spectral space. 

2956 # fkk is a 2D array of DCT coefficients, representing the amplitudes of 

2957 # cosine basis functions at different spatial frequencies. 

2958 

2959 # normalise spectrum to allow comparison between models. 

2960 fkk = fft.dctn(y_2d, norm="ortho") 

2961 

2962 # Normalise fkk 

2963 fkk = fkk / np.sqrt(Ny * Nx) 

2964 

2965 # calculate variance of spectral coefficient 

2966 sigma_2 = fkk**2 / Nx / Ny 

2967 

2968 # Group ellipses of alphas into the same wavenumber k/Nmin 

2969 for k in range(1, Nmin + 1): 

2970 alpha = k / Nmin 

2971 alpha_p1 = (k + 1) / Nmin 

2972 

2973 # Sum up elements matching k 

2974 mask_k = np.where((alpha_matrix >= alpha) & (alpha_matrix < alpha_p1)) 

2975 ps_array[t, k - 1] = np.sum(sigma_2[mask_k]) 

2976 

2977 return ps_array 

2978 

2979 

2980def _create_alpha_matrix(Ny, Nx): 

2981 """Construct an array of 2D wavenumbers from 2D wavenumber pair. 

2982 

2983 Parameters 

2984 ---------- 

2985 Ny, Nx: 

2986 Dimensions of the 2D field for which the power spectra is calculated. Used to 

2987 create the array of 2D wavenumbers. Each Ny, Nx pair is associated with a 

2988 single-scale parameter. 

2989 

2990 Returns: alpha_matrix 

2991 normalisation of 2D wavenumber axes, transforming the spectral domain into 

2992 an elliptic coordinate system. 

2993 

2994 """ 

2995 # Create x_indices: each row is [1, 2, ..., Nx] 

2996 x_indices = np.tile(np.arange(1, Nx + 1), (Ny, 1)) 

2997 

2998 # Create y_indices: each column is [1, 2, ..., Ny] 

2999 y_indices = np.tile(np.arange(1, Ny + 1).reshape(Ny, 1), (1, Nx)) 

3000 

3001 # Compute alpha_matrix 

3002 alpha_matrix = np.sqrt((x_indices**2) / Nx**2 + (y_indices**2) / Ny**2) 

3003 

3004 return alpha_matrix