Coverage for src / CSET / operators / read.py: 90%

386 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-04-20 13:52 +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 for reading various types of files from disk.""" 

16 

17import ast 

18import datetime 

19import functools 

20import glob 

21import itertools 

22import logging 

23from pathlib import Path 

24from typing import Literal 

25 

26import dask 

27import iris 

28import iris.coord_systems 

29import iris.coords 

30import iris.cube 

31import iris.exceptions 

32import iris.util 

33import numpy as np 

34from iris.analysis.cartography import rotate_pole, rotate_winds 

35 

36from CSET._common import iter_maybe 

37from CSET.operators._stash_to_lfric import STASH_TO_LFRIC 

38from CSET.operators._utils import ( 

39 get_cube_coordindex, 

40 get_cube_yxcoordname, 

41 is_spatialdim, 

42) 

43 

44 

45class NoDataError(FileNotFoundError): 

46 """Error that no data has been loaded.""" 

47 

48 

49def read_cube( 

50 file_paths: list[str] | str, 

51 constraint: iris.Constraint = None, 

52 model_names: list[str] | str | None = None, 

53 subarea_type: str = None, 

54 subarea_extent: list[float] = None, 

55 **kwargs, 

56) -> iris.cube.Cube: 

57 """Read a single cube from files. 

58 

59 Read operator that takes a path string (can include shell-style glob 

60 patterns), and loads the cube matching the constraint. If any paths point to 

61 directory, all the files contained within are loaded. 

62 

63 Ensemble data can also be loaded. If it has a realization coordinate 

64 already, it will be directly used. If not, it will have its member number 

65 guessed from the filename, based on one of several common patterns. For 

66 example the pattern *emXX*, where XX is the realization. 

67 

68 Deterministic data will be loaded with a realization of 0, allowing it to be 

69 processed in the same way as ensemble data. 

70 

71 Arguments 

72 --------- 

73 file_paths: str | list[str] 

74 Path or paths to where .pp/.nc files are located 

75 constraint: iris.Constraint | iris.ConstraintCombination, optional 

76 Constraints to filter data by. Defaults to unconstrained. 

77 model_names: str | list[str], optional 

78 Names of the models that correspond to respective paths in file_paths. 

79 subarea_type: "gridcells" | "modelrelative" | "realworld", optional 

80 Whether to constrain data by model relative coordinates or real world 

81 coordinates. 

82 subarea_extent: list, optional 

83 List of coordinates to constraint data by, in order lower latitude, 

84 upper latitude, lower longitude, upper longitude. 

85 

86 Returns 

87 ------- 

88 cubes: iris.cube.Cube 

89 Cube loaded 

90 

91 Raises 

92 ------ 

93 FileNotFoundError 

94 If the provided path does not exist 

95 ValueError 

96 If the constraint doesn't produce a single cube. 

97 """ 

98 cubes = read_cubes( 

99 file_paths=file_paths, 

100 constraint=constraint, 

101 model_names=model_names, 

102 subarea_type=subarea_type, 

103 subarea_extent=subarea_extent, 

104 ) 

105 # Check filtered cubes is a CubeList containing one cube. 

106 if len(cubes) == 1: 

107 return cubes[0] 

108 else: 

109 raise ValueError( 

110 f"Constraint doesn't produce single cube: {constraint}\n{cubes}" 

111 ) 

112 

113 

114def read_cubes( 

115 file_paths: list[str] | str, 

116 constraint: iris.Constraint | None = None, 

117 model_names: str | list[str] | None = None, 

118 subarea_type: str = None, 

119 subarea_extent: list = None, 

120 **kwargs, 

121) -> iris.cube.CubeList: 

122 """Read cubes from files. 

123 

124 Read operator that takes a path string (can include shell-style glob 

125 patterns), and loads the cubes matching the constraint. If any paths point 

126 to directory, all the files contained within are loaded. 

127 

128 Ensemble data can also be loaded. If it has a realization coordinate 

129 already, it will be directly used. If not, it will have its member number 

130 guessed from the filename, based on one of several common patterns. For 

131 example the pattern *emXX*, where XX is the realization. 

132 

133 Deterministic data will be loaded with a realization of 0, allowing it to be 

134 processed in the same way as ensemble data. 

135 

136 Data output by XIOS (such as LFRic) has its per-file metadata removed so 

137 that the cubes merge across files. 

138 

139 Arguments 

140 --------- 

141 file_paths: str | list[str] 

142 Path or paths to where .pp/.nc files are located. Can include globs. 

143 constraint: iris.Constraint | iris.ConstraintCombination, optional 

144 Constraints to filter data by. Defaults to unconstrained. 

145 model_names: str | list[str], optional 

146 Names of the models that correspond to respective paths in file_paths. 

147 subarea_type: str, optional 

148 Whether to constrain data by model relative coordinates or real world 

149 coordinates. 

150 subarea_extent: list[float], optional 

151 List of coordinates to constraint data by, in order lower latitude, 

152 upper latitude, lower longitude, upper longitude. 

153 

154 Returns 

155 ------- 

156 cubes: iris.cube.CubeList 

157 Cubes loaded after being merged and concatenated. 

158 

159 Raises 

160 ------ 

161 FileNotFoundError 

162 If the provided path does not exist 

163 """ 

164 # Get iterable of paths. Each path corresponds to 1 model. 

165 paths = iter_maybe(file_paths) 

166 model_names = iter_maybe(model_names) 

167 

168 # Check we have appropriate number of model names. 

169 if model_names != (None,) and len(model_names) != len(paths): 

170 raise ValueError( 

171 f"The number of model names ({len(model_names)}) should equal " 

172 f"the number of paths given ({len(paths)})." 

173 ) 

174 

175 # Load the data for each model into a CubeList per model. 

176 model_cubes = ( 

177 _load_model(path, name, constraint) 

178 for path, name in itertools.zip_longest(paths, model_names, fillvalue=None) 

179 ) 

180 

181 # Split out first model's cubes and mark it as the base for comparisons. 

182 cubes = next(model_cubes) 

183 for cube in cubes: 

184 # Use 1 to indicate True, as booleans can't be saved in NetCDF attributes. 

185 cube.attributes["cset_comparison_base"] = 1 

186 

187 # Load the rest of the models. 

188 cubes.extend(itertools.chain.from_iterable(model_cubes)) 

189 

190 # Unify time units so different case studies can merge. 

191 iris.util.unify_time_units(cubes) 

192 

193 # Select sub region. 

194 cubes = _cutout_cubes(cubes, subarea_type, subarea_extent) 

195 

196 # Merge and concatenate cubes now metadata has been fixed. 

197 cubes = cubes.merge() 

198 cubes = cubes.concatenate() 

199 

200 # Squeeze single valued coordinates into scalar coordinates. 

201 cubes = iris.cube.CubeList(iris.util.squeeze(cube) for cube in cubes) 

202 

203 # Ensure dimension coordinates are bounded. 

204 for cube in cubes: 

205 for dim_coord in cube.coords(dim_coords=True): 

206 # Iris can't guess the bounds of a scalar coordinate. 

207 if not dim_coord.has_bounds() and dim_coord.shape[0] > 1: 

208 dim_coord.guess_bounds() 

209 

210 logging.info("Loaded cubes: %s", cubes) 

211 if len(cubes) == 0: 

212 raise NoDataError("No cubes loaded, check your constraints!") 

213 return cubes 

214 

215 

216def _load_model( 

217 paths: str | list[str], 

218 model_name: str | None, 

219 constraint: iris.Constraint | None, 

220) -> iris.cube.CubeList: 

221 """Load a single model's data into a CubeList.""" 

222 input_files = _check_input_files(paths) 

223 # If unset, a constraint of None lets everything be loaded. 

224 logging.debug("Constraint: %s", constraint) 

225 cubes = iris.load(input_files, constraint, callback=_loading_callback) 

226 # Make the UM's winds consistent with LFRic. 

227 _fix_um_winds(cubes) 

228 

229 # Add model_name attribute to each cube to make it available at any further 

230 # step without needing to pass it as function parameter. 

231 if model_name is not None: 

232 for cube in cubes: 

233 cube.attributes["model_name"] = model_name 

234 return cubes 

235 

236 

237def _check_input_files(input_paths: str | list[str]) -> list[Path]: 

238 """Get an iterable of files to load, and check that they all exist. 

239 

240 Arguments 

241 --------- 

242 input_paths: list[str] 

243 List of paths to input files or directories. The path may itself contain 

244 glob patterns, but unlike in shells it will match directly first. 

245 

246 Returns 

247 ------- 

248 list[Path] 

249 A list of files to load. 

250 

251 Raises 

252 ------ 

253 FileNotFoundError: 

254 If the provided arguments don't resolve to at least one existing file. 

255 """ 

256 files = [] 

257 for raw_filename in iter_maybe(input_paths): 

258 # Match glob-like files first, if they exist. 

259 raw_path = Path(raw_filename) 

260 if raw_path.is_file(): 

261 files.append(raw_path) 

262 else: 

263 for input_path in glob.glob(raw_filename): 

264 # Convert string paths into Path objects. 

265 input_path = Path(input_path) 

266 # Get the list of files in the directory, or use it directly. 

267 if input_path.is_dir(): 

268 logging.debug("Checking directory '%s' for files", input_path) 

269 files.extend(p for p in input_path.iterdir() if p.is_file()) 

270 else: 

271 files.append(input_path) 

272 

273 files.sort() 

274 logging.info("Loading files:\n%s", "\n".join(str(path) for path in files)) 

275 if len(files) == 0: 

276 raise FileNotFoundError(f"No files found for {input_paths}") 

277 return files 

278 

279 

280def _cutout_cubes( 

281 cubes: iris.cube.CubeList, 

282 subarea_type: Literal["gridcells", "realworld", "modelrelative"] | None, 

283 subarea_extent: list[float], 

284): 

285 """Cut out a subarea from a CubeList.""" 

286 if subarea_type is None: 

287 logging.debug("Subarea selection is disabled.") 

288 return cubes 

289 

290 # If selected, cutout according to number of grid cells to trim from each edge. 

291 cutout_cubes = iris.cube.CubeList() 

292 # Find spatial coordinates 

293 for cube in cubes: 

294 # Find dimension coordinates. 

295 lat_name, lon_name = get_cube_yxcoordname(cube) 

296 

297 # Compute cutout based on number of cells to trim from edges. 

298 if subarea_type == "gridcells": 

299 logging.debug( 

300 "User requested LowerTrim: %s LeftTrim: %s UpperTrim: %s RightTrim: %s", 

301 subarea_extent[0], 

302 subarea_extent[1], 

303 subarea_extent[2], 

304 subarea_extent[3], 

305 ) 

306 lat_points = np.sort(cube.coord(lat_name).points) 

307 lon_points = np.sort(cube.coord(lon_name).points) 

308 # Define cutout region using user provided cell points. 

309 lats = [lat_points[subarea_extent[0]], lat_points[-subarea_extent[2] - 1]] 

310 lons = [lon_points[subarea_extent[1]], lon_points[-subarea_extent[3] - 1]] 

311 

312 # Compute cutout based on specified coordinate values. 

313 elif subarea_type == "realworld" or subarea_type == "modelrelative": 

314 # If not gridcells, cutout by requested geographic area, 

315 logging.debug( 

316 "User requested LLat: %s ULat: %s LLon: %s ULon: %s", 

317 subarea_extent[0], 

318 subarea_extent[1], 

319 subarea_extent[2], 

320 subarea_extent[3], 

321 ) 

322 # Define cutout region using user provided coordinates. 

323 lats = np.array(subarea_extent[0:2]) 

324 lons = np.array(subarea_extent[2:4]) 

325 # Ensure cutout longitudes are within +/- 180.0 bounds. 

326 while lons[0] < -180.0: 

327 lons += 360.0 

328 while lons[1] > 180.0: 

329 lons -= 360.0 

330 # If the coordinate system is rotated we convert coordinates into 

331 # model-relative coordinates to extract the appropriate cutout. 

332 coord_system = cube.coord(lat_name).coord_system 

333 if subarea_type == "realworld" and isinstance( 

334 coord_system, iris.coord_systems.RotatedGeogCS 

335 ): 

336 lons, lats = rotate_pole( 

337 lons, 

338 lats, 

339 pole_lon=coord_system.grid_north_pole_longitude, 

340 pole_lat=coord_system.grid_north_pole_latitude, 

341 ) 

342 else: 

343 raise ValueError("Unknown subarea_type:", subarea_type) 

344 

345 # Do cutout and add to cutout_cubes. 

346 intersection_args = {lat_name: lats, lon_name: lons} 

347 logging.debug("Cutting out coords: %s", intersection_args) 

348 try: 

349 cutout_cubes.append(cube.intersection(**intersection_args)) 

350 except IndexError as err: 

351 raise ValueError( 

352 "Region cutout error. Check and update SUBAREA_EXTENT." 

353 "Cutout region requested should be contained within data area. " 

354 "Also check if cutout region requested is smaller than input grid spacing." 

355 ) from err 

356 

357 return cutout_cubes 

358 

359 

360def _loading_callback(cube: iris.cube.Cube, field, filename: str) -> iris.cube.Cube: 

361 """Compose together the needed callbacks into a single function.""" 

362 # Most callbacks operate in-place, but save the cube when returned! 

363 _realization_callback(cube) 

364 _um_normalise_callback(cube) 

365 _lfric_normalise_callback(cube) 

366 cube = _lfric_time_coord_fix_callback(cube) 

367 _normalise_var0_varname(cube) 

368 cube = _fix_no_spatial_coords_callback(cube) 

369 _fix_spatial_coords_callback(cube) 

370 _fix_pressure_coord_callback(cube) 

371 _fix_um_radtime(cube) 

372 _fix_cell_methods(cube) 

373 cube = _convert_cube_units_callback(cube) 

374 cube = _grid_longitude_fix_callback(cube) 

375 _fix_lfric_cloud_base_altitude(cube) 

376 _proleptic_gregorian_fix(cube) 

377 _lfric_time_callback(cube) 

378 _lfric_forecast_period_callback(cube) 

379 _normalise_ML_varname(cube) 

380 return cube 

381 

382 

383def _realization_callback(cube): 

384 """Give deterministic cubes a realization of 0. 

385 

386 This means they can be handled in the same way as ensembles through the rest 

387 of the code. 

388 """ 

389 # Only add if realization coordinate does not exist. 

390 if not cube.coords("realization"): 

391 cube.add_aux_coord( 

392 iris.coords.DimCoord(0, standard_name="realization", units="1") 

393 ) 

394 

395 

396@functools.lru_cache(None) 

397def _log_once(msg, level=logging.WARNING): 

398 """Print a warning message, skipping recent duplicates.""" 

399 logging.log(level, msg) 

400 

401 

402def _um_normalise_callback(cube: iris.cube.Cube): 

403 """Normalise UM STASH variable long names to LFRic variable names. 

404 

405 Note standard names will remain associated with cubes where different. 

406 Long name will be used consistently in output filename and titles. 

407 """ 

408 # Convert STASH to LFRic variable name 

409 if "STASH" in cube.attributes: 

410 stash = cube.attributes["STASH"] 

411 try: 

412 (name, grid) = STASH_TO_LFRIC[str(stash)] 

413 cube.long_name = name 

414 except KeyError: 

415 # Don't change cubes with unknown stash codes. 

416 _log_once( 

417 f"Unknown STASH code: {stash}. Please check file stash_to_lfric.py to update.", 

418 level=logging.WARNING, 

419 ) 

420 

421 

422def _lfric_normalise_callback(cube: iris.cube.Cube): 

423 """Normalise attributes that prevents LFRic cube from merging. 

424 

425 The uuid and timeStamp relate to the output file, as saved by XIOS, and has 

426 no relation to the data contained. These attributes are removed. 

427 

428 The um_stash_source is a list of STASH codes for when an LFRic field maps to 

429 multiple UM fields, however it can be encoded in any order. This attribute 

430 is sorted to prevent this. This attribute is only present in LFRic data that 

431 has been converted to look like UM data. 

432 """ 

433 # Remove unwanted attributes. 

434 cube.attributes.pop("timeStamp", None) 

435 cube.attributes.pop("uuid", None) 

436 cube.attributes.pop("name", None) 

437 cube.attributes.pop("source", None) 

438 cube.attributes.pop("analysis_source", None) 

439 cube.attributes.pop("history", None) 

440 

441 # Sort STASH code list. 

442 stash_list = cube.attributes.get("um_stash_source") 

443 if stash_list: 

444 # Parse the string as a list, sort, then re-encode as a string. 

445 cube.attributes["um_stash_source"] = str(sorted(ast.literal_eval(stash_list))) 

446 

447 

448def _lfric_time_coord_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube: 

449 """Ensure the time coordinate is a DimCoord rather than an AuxCoord. 

450 

451 The coordinate is converted and replaced if not. SLAMed LFRic data has this 

452 issue, though the coordinate satisfies all the properties for a DimCoord. 

453 Scalar time values are left as AuxCoords. 

454 """ 

455 # This issue seems to come from iris's handling of NetCDF files where time 

456 # always ends up as an AuxCoord. 

457 if cube.coords("time"): 

458 time_coord = cube.coord("time") 

459 if ( 

460 not isinstance(time_coord, iris.coords.DimCoord) 

461 and len(cube.coord_dims(time_coord)) == 1 

462 ): 

463 # Fudge the bounds to foil checking for strict monotonicity. 

464 if time_coord.has_bounds(): 464 ↛ 465line 464 didn't jump to line 465 because the condition on line 464 was never true

465 if (time_coord.bounds[-1][0] - time_coord.bounds[0][0]) < 1.0e-8: 

466 time_coord.bounds = [ 

467 [ 

468 time_coord.bounds[i][0] + 1.0e-8 * float(i), 

469 time_coord.bounds[i][1], 

470 ] 

471 for i in range(len(time_coord.bounds)) 

472 ] 

473 iris.util.promote_aux_coord_to_dim_coord(cube, time_coord) 

474 return cube 

475 

476 

477def _grid_longitude_fix_callback(cube: iris.cube.Cube) -> iris.cube.Cube: 

478 """Check grid_longitude coordinates are in the range -180 deg to 180 deg. 

479 

480 This is necessary if comparing two models with different conventions -- 

481 for example, models where the prime meridian is defined as 0 deg or 

482 360 deg. If not in the range -180 deg to 180 deg, we wrap the grid_longitude 

483 so that it falls in this range. Checks are for near-180 bounds given 

484 model data bounds may not extend exactly to 0. or 360. 

485 Input cubes on non-rotated grid coordinates are not impacted. 

486 """ 

487 try: 

488 y, x = get_cube_yxcoordname(cube) 

489 except ValueError: 

490 # Don't modify non-spatial cubes. 

491 return cube 

492 

493 long_coord = cube.coord(x) 

494 # Wrap longitudes if rotated pole coordinates 

495 coord_system = long_coord.coord_system 

496 if x == "grid_longitude" and isinstance( 

497 coord_system, iris.coord_systems.RotatedGeogCS 

498 ): 

499 long_points = long_coord.points.copy() 

500 long_centre = np.median(long_points) 

501 while long_centre < -175.0: 

502 long_centre += 360.0 

503 long_points += 360.0 

504 while long_centre >= 175.0: 

505 long_centre -= 360.0 

506 long_points -= 360.0 

507 long_coord.points = long_points 

508 

509 # Update coord bounds to be consistent with wrapping. 

510 if long_coord.has_bounds(): 

511 long_coord.bounds = None 

512 long_coord.guess_bounds() 

513 

514 return cube 

515 

516 

517def _fix_no_spatial_coords_callback(cube: iris.cube.Cube): 

518 import CSET.operators._utils as utils 

519 

520 # Don't modify spatial cubes that already have spatial dimensions 

521 if utils.is_spatialdim(cube): 

522 return cube 

523 

524 else: 

525 # attempt to get lat/long from cube attributes 

526 try: 

527 lat_min = cube.attributes.get("geospatial_lat_min") 

528 lat_max = cube.attributes.get("geospatial_lat_max") 

529 lon_min = cube.attributes.get("geospatial_lon_min") 

530 lon_max = cube.attributes.get("geospatial_lon_max") 

531 

532 lon_val = (lon_min + lon_max) / 2.0 

533 lat_val = (lat_min + lat_max) / 2.0 

534 

535 lat_coord = iris.coords.DimCoord( 

536 lat_val, 

537 standard_name="latitude", 

538 units="degrees_north", 

539 var_name="latitude", 

540 coord_system=iris.coord_systems.GeogCS(6371229.0), 

541 circular=True, 

542 ) 

543 

544 lon_coord = iris.coords.DimCoord( 

545 lon_val, 

546 standard_name="longitude", 

547 units="degrees_east", 

548 var_name="longitude", 

549 coord_system=iris.coord_systems.GeogCS(6371229.0), 

550 circular=True, 

551 ) 

552 

553 cube.add_aux_coord(lat_coord) 

554 cube.add_aux_coord(lon_coord) 

555 return cube 

556 

557 # if lat/long are not in attributes, then return cube unchanged: 

558 except TypeError: 

559 return cube 

560 

561 

562def _fix_spatial_coords_callback(cube: iris.cube.Cube): 

563 """Check latitude and longitude coordinates name. 

564 

565 This is necessary as some models define their grid as on rotated 

566 'grid_latitude' and 'grid_longitude' coordinates while others define 

567 the grid on non-rotated 'latitude' and 'longitude'. 

568 Cube dimensions need to be made consistent to avoid recipe failures, 

569 particularly where comparing multiple input models with differing spatial 

570 coordinates. 

571 """ 

572 # Check if cube is spatial. 

573 if not is_spatialdim(cube): 

574 # Don't modify non-spatial cubes. 

575 return 

576 

577 # Get spatial coords and dimension index. 

578 y_name, x_name = get_cube_yxcoordname(cube) 

579 ny = get_cube_coordindex(cube, y_name) 

580 nx = get_cube_coordindex(cube, x_name) 

581 

582 # Remove spatial coords bounds if erroneous values detected. 

583 # Aims to catch some errors in input coord bounds by setting 

584 # invalid threshold of 10000.0 

585 if cube.coord(x_name).has_bounds() and cube.coord(y_name).has_bounds(): 

586 bx_max = np.max(np.abs(cube.coord(x_name).bounds)) 

587 by_max = np.max(np.abs(cube.coord(y_name).bounds)) 

588 if bx_max > 10000.0 or by_max > 10000.0: 

589 cube.coord(x_name).bounds = None 

590 cube.coord(y_name).bounds = None 

591 

592 # Translate [grid_latitude, grid_longitude] to an unrotated 1-d DimCoord 

593 # [latitude, longitude] for instances where rotated_pole=90.0 

594 if "grid_latitude" in [coord.name() for coord in cube.coords(dim_coords=True)]: 

595 coord_system = cube.coord("grid_latitude").coord_system 

596 pole_lat = getattr(coord_system, "grid_north_pole_latitude", None) 

597 if pole_lat == 90.0: 597 ↛ 598line 597 didn't jump to line 598 because the condition on line 597 was never true

598 lats = cube.coord("grid_latitude").points 

599 lons = cube.coord("grid_longitude").points 

600 

601 cube.remove_coord("grid_latitude") 

602 cube.add_dim_coord( 

603 iris.coords.DimCoord( 

604 lats, 

605 standard_name="latitude", 

606 var_name="latitude", 

607 units="degrees", 

608 coord_system=iris.coord_systems.GeogCS(6371229.0), 

609 circular=True, 

610 ), 

611 ny, 

612 ) 

613 y_name = "latitude" 

614 cube.remove_coord("grid_longitude") 

615 cube.add_dim_coord( 

616 iris.coords.DimCoord( 

617 lons, 

618 standard_name="longitude", 

619 var_name="longitude", 

620 units="degrees", 

621 coord_system=iris.coord_systems.GeogCS(6371229.0), 

622 circular=True, 

623 ), 

624 nx, 

625 ) 

626 x_name = "longitude" 

627 

628 # Create additional AuxCoord [grid_latitude, grid_longitude] with 

629 # rotated pole attributes for cases with [lat, lon] inputs 

630 if y_name in ["latitude"] and cube.coord(y_name).units in [ 

631 "degrees", 

632 "degrees_north", 

633 "degrees_south", 

634 ]: 

635 # Add grid_latitude AuxCoord 

636 if "grid_latitude" not in [ 636 ↛ 649line 636 didn't jump to line 649 because the condition on line 636 was always true

637 coord.name() for coord in cube.coords(dim_coords=False) 

638 ]: 

639 cube.add_aux_coord( 

640 iris.coords.AuxCoord( 

641 cube.coord(y_name).points, 

642 var_name="grid_latitude", 

643 units="degrees", 

644 ), 

645 ny, 

646 ) 

647 # Ensure input latitude DimCoord has CoordSystem 

648 # This attribute is sometimes lost on iris.save 

649 if not cube.coord(y_name).coord_system: 

650 cube.coord(y_name).coord_system = iris.coord_systems.GeogCS(6371229.0) 

651 

652 if x_name in ["longitude"] and cube.coord(x_name).units in [ 

653 "degrees", 

654 "degrees_west", 

655 "degrees_east", 

656 ]: 

657 # Add grid_longitude AuxCoord 

658 if "grid_longitude" not in [ 658 ↛ 672line 658 didn't jump to line 672 because the condition on line 658 was always true

659 coord.name() for coord in cube.coords(dim_coords=False) 

660 ]: 

661 cube.add_aux_coord( 

662 iris.coords.AuxCoord( 

663 cube.coord(x_name).points, 

664 var_name="grid_longitude", 

665 units="degrees", 

666 ), 

667 nx, 

668 ) 

669 

670 # Ensure input longitude DimCoord has CoordSystem 

671 # This attribute is sometimes lost on iris.save 

672 if not cube.coord(x_name).coord_system: 

673 cube.coord(x_name).coord_system = iris.coord_systems.GeogCS(6371229.0) 

674 

675 

676def _fix_pressure_coord_callback(cube: iris.cube.Cube): 

677 """Rename pressure coordinate to "pressure" if it exists and ensure hPa units. 

678 

679 This problem was raised because the AIFS model data from ECMWF 

680 defines the pressure coordinate with the name "pressure_level" rather 

681 than compliant CF coordinate names. 

682 

683 Additionally, set the units of pressure to be hPa to be consistent with the UM, 

684 and approach the coordinates in a unified way. 

685 """ 

686 for coord in cube.dim_coords: 

687 if coord.name() in ["pressure_level", "pressure_levels"]: 

688 coord.rename("pressure") 

689 

690 if coord.name() == "pressure": 

691 if str(cube.coord("pressure").units) != "hPa": 

692 cube.coord("pressure").convert_units("hPa") 

693 

694 

695def _fix_um_radtime(cube: iris.cube.Cube): 

696 """Move radiation diagnostics from timestamps which are output N minutes or seconds past every hour. 

697 

698 This callback does not have any effect for output diagnostics with 

699 timestamps exactly 00 or 30 minutes past the hour. Only radiation 

700 diagnostics are checked. 

701 Note this callback does not interpolate the data in time, only adjust 

702 timestamps to sit on the hour to enable time-to-time difference plotting 

703 with models which may output radiation data on the hour. 

704 """ 

705 try: 

706 if cube.attributes["STASH"] in [ 

707 "m01s01i207", 

708 "m01s01i208", 

709 "m01s02i205", 

710 "m01s02i201", 

711 "m01s01i207", 

712 "m01s02i207", 

713 "m01s01i235", 

714 ]: 

715 time_coord = cube.coord("time") 

716 

717 # Convert time points to datetime objects 

718 time_unit = time_coord.units 

719 time_points = time_unit.num2date(time_coord.points) 

720 # Skip if times don't need fixing. 

721 if time_points[0].minute == 0 and time_points[0].second == 0: 

722 return 

723 if time_points[0].minute == 30 and time_points[0].second == 0: 723 ↛ 724line 723 didn't jump to line 724 because the condition on line 723 was never true

724 return 

725 

726 # Subtract time difference from the hour from each time point 

727 n_minute = time_points[0].minute 

728 n_second = time_points[0].second 

729 # If times closer to next hour, compute difference to add on to following hour 

730 if n_minute > 30: 

731 n_minute = n_minute - 60 

732 # Compute new diagnostic time stamp 

733 new_time_points = ( 

734 time_points 

735 - datetime.timedelta(minutes=n_minute) 

736 - datetime.timedelta(seconds=n_second) 

737 ) 

738 

739 # Convert back to numeric values using the original time unit. 

740 new_time_values = time_unit.date2num(new_time_points) 

741 

742 # Replace the time coordinate with updated values. 

743 time_coord.points = new_time_values 

744 

745 # Recompute forecast_period with corrected values. 

746 if cube.coord("forecast_period"): 746 ↛ exitline 746 didn't return from function '_fix_um_radtime' because the condition on line 746 was always true

747 fcst_prd_points = cube.coord("forecast_period").points 

748 new_fcst_points = ( 

749 time_unit.num2date(fcst_prd_points) 

750 - datetime.timedelta(minutes=n_minute) 

751 - datetime.timedelta(seconds=n_second) 

752 ) 

753 cube.coord("forecast_period").points = time_unit.date2num( 

754 new_fcst_points 

755 ) 

756 except KeyError: 

757 pass 

758 

759 

760def _fix_cell_methods(cube: iris.cube.Cube): 

761 """To fix the assumed cell_methods in accumulation STASH from UM. 

762 

763 Lightning (m01s21i104), rainfall amount (m01s04i201, m01s05i201) and snowfall amount 

764 (m01s04i202, m01s05i202) in UM is being output as a time accumulation, 

765 over each hour (TAcc1hr), but input cubes show cell_methods as "mean". 

766 For UM and LFRic inputs to be compatible, we assume accumulated cell_methods are 

767 "sum". This callback changes "mean" cube attribute cell_method to "sum", 

768 enabling the cell_method constraint on reading to select correct input. 

769 """ 

770 # Shift "mean" cell_method to "sum" for selected UM inputs. 

771 if cube.attributes.get("STASH") in [ 

772 "m01s21i104", 

773 "m01s04i201", 

774 "m01s04i202", 

775 "m01s05i201", 

776 "m01s05i202", 

777 ]: 

778 # Check if input cell_method contains "mean" time-processing. 

779 if set(cm.method for cm in cube.cell_methods) == {"mean"}: 779 ↛ exitline 779 didn't return from function '_fix_cell_methods' because the condition on line 779 was always true

780 # Retrieve interval and any comment information. 

781 for cell_method in cube.cell_methods: 

782 interval_str = cell_method.intervals 

783 comment_str = cell_method.comments 

784 

785 # Remove input aggregation method. 

786 cube.cell_methods = () 

787 

788 # Replace "mean" with "sum" cell_method to indicate aggregation. 

789 cube.add_cell_method( 

790 iris.coords.CellMethod( 

791 method="sum", 

792 coords="time", 

793 intervals=interval_str, 

794 comments=comment_str, 

795 ) 

796 ) 

797 

798 

799def _convert_cube_units_callback(cube: iris.cube.Cube): 

800 """Adjust diagnostic units for specific variables. 

801 

802 Some precipitation diagnostics are output with unit kg m-2 s-1 and are 

803 converted here to mm hr-1. 

804 

805 Visibility diagnostics are converted here from m to km to improve output 

806 formatting. 

807 """ 

808 # Convert precipitation diagnostic units if required. 

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

810 if any("surface_microphysical" in name for name in varnames): 

811 if cube.units == "kg m-2 s-1": 

812 _log_once( 

813 "Converting precipitation rate units from kg m-2 s-1 to mm hr-1", 

814 level=logging.DEBUG, 

815 ) 

816 # Convert from kg m-2 s-1 to mm s-1 assuming 1kg water = 1l water = 1dm^3 water. 

817 # This is a 1:1 conversion, so we just change the units. 

818 cube.units = "mm s-1" 

819 # Convert the units to per hour. 

820 cube.convert_units("mm hr-1") 

821 elif cube.units == "kg m-2": 821 ↛ 831line 821 didn't jump to line 831 because the condition on line 821 was always true

822 _log_once( 

823 "Converting precipitation amount units from kg m-2 to mm", 

824 level=logging.DEBUG, 

825 ) 

826 # Convert from kg m-2 to mm assuming 1kg water = 1l water = 1dm^3 water. 

827 # This is a 1:1 conversion, so we just change the units. 

828 cube.units = "mm" 

829 

830 # Convert visibility diagnostic units if required. 

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

832 if any("visibility" in name for name in varnames): 

833 if cube.units == "m": 833 ↛ 838line 833 didn't jump to line 838 because the condition on line 833 was always true

834 _log_once("Converting visibility units m to km.", level=logging.DEBUG) 

835 # Convert the units to km. 

836 cube.convert_units("km") 

837 

838 return cube 

839 

840 

841def _fix_lfric_cloud_base_altitude(cube: iris.cube.Cube): 

842 """Mask cloud_base_altitude diagnostic in regions with no cloud.""" 

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

844 if any("cloud_base_altitude" in name for name in varnames): 

845 # Mask cube where set > 144kft to catch default 144.35695538058164 

846 cube.data = dask.array.ma.masked_greater(cube.core_data(), 144.0) 

847 

848 

849def _fix_um_winds(cubes: iris.cube.CubeList): 

850 """To make winds from the UM consistent with those from LFRic. 

851 

852 Diagnostics of wind are not always consistent between the UM 

853 and LFric. Here, winds from the UM are adjusted to make them i 

854 consistent with LFRic. 

855 """ 

856 # Check whether we have components of the wind identified by STASH, 

857 # (so this will apply only to cubes from the UM), but not the 

858 # wind speed and calculate it if it is missing. Note that 

859 # this will be biased low in general because the components will mostly 

860 # be time averages. For simplicity, we do this only if there is just one 

861 # cube of a component. A more complicated approach would be to consider 

862 # the cell methods, but it may not be warranted. 

863 u_constr = iris.AttributeConstraint(STASH="m01s03i225") 

864 v_constr = iris.AttributeConstraint(STASH="m01s03i226") 

865 speed_constr = iris.AttributeConstraint(STASH="m01s03i227") 

866 try: 

867 if cubes.extract(u_constr) and cubes.extract(v_constr): 867 ↛ 868line 867 didn't jump to line 868 because the condition on line 867 was never true

868 if len(cubes.extract(u_constr)) == 1 and not cubes.extract(speed_constr): 

869 _add_wind_speed_um(cubes) 

870 # Convert winds in the UM to be relative to true east and true north. 

871 _convert_wind_true_dirn_um(cubes) 

872 except (KeyError, AttributeError): 

873 pass 

874 

875 

876def _add_wind_speed_um(cubes: iris.cube.CubeList): 

877 """Add windspeeds to cubes from the UM.""" 

878 wspd10 = ( 

879 cubes.extract_cube(iris.AttributeConstraint(STASH="m01s03i225"))[0] ** 2 

880 + cubes.extract_cube(iris.AttributeConstraint(STASH="m01s03i226"))[0] ** 2 

881 ) ** 0.5 

882 wspd10.attributes["STASH"] = "m01s03i227" 

883 wspd10.standard_name = "wind_speed" 

884 wspd10.long_name = "wind_speed_at_10m" 

885 cubes.append(wspd10) 

886 

887 

888def _convert_wind_true_dirn_um(cubes: iris.cube.CubeList): 

889 """To convert winds to true directions. 

890 

891 Convert from the components relative to the grid to true directions. 

892 This functionality only handles the simplest case. 

893 """ 

894 u_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i225")) 

895 v_grids = cubes.extract(iris.AttributeConstraint(STASH="m01s03i226")) 

896 for u, v in zip(u_grids, v_grids, strict=True): 

897 true_u, true_v = rotate_winds(u, v, iris.coord_systems.GeogCS(6371229.0)) 

898 u.data = true_u.core_data() 

899 v.data = true_v.core_data() 

900 

901 

902def _normalise_var0_varname(cube: iris.cube.Cube): 

903 """Fix varnames for consistency to allow merging. 

904 

905 Some model data netCDF sometimes have a coordinate name end in 

906 "_0" etc, where duplicate coordinates of same name are defined but 

907 with different attributes. This can be inconsistently managed in 

908 different model inputs and can cause cubes to fail to merge. 

909 """ 

910 for coord in cube.coords(): 

911 if coord.var_name and coord.var_name.endswith("_0"): 

912 coord.var_name = coord.var_name.removesuffix("_0") 

913 if coord.var_name and coord.var_name.endswith("_1"): 

914 coord.var_name = coord.var_name.removesuffix("_1") 

915 if coord.var_name and coord.var_name.endswith("_2"): 915 ↛ 916line 915 didn't jump to line 916 because the condition on line 915 was never true

916 coord.var_name = coord.var_name.removesuffix("_2") 

917 if coord.var_name and coord.var_name.endswith("_3"): 917 ↛ 918line 917 didn't jump to line 918 because the condition on line 917 was never true

918 coord.var_name = coord.var_name.removesuffix("_3") 

919 

920 if cube.var_name and cube.var_name.endswith("_0"): 

921 cube.var_name = cube.var_name.removesuffix("_0") 

922 

923 

924def _proleptic_gregorian_fix(cube: iris.cube.Cube): 

925 """Convert the calendars of time units to use a standard calendar.""" 

926 try: 

927 time_coord = cube.coord("time") 

928 if time_coord.units.calendar == "proleptic_gregorian": 

929 logging.debug( 

930 "Changing proleptic Gregorian calendar to standard calendar for %s", 

931 repr(time_coord.units), 

932 ) 

933 time_coord.units = time_coord.units.change_calendar("standard") 

934 except iris.exceptions.CoordinateNotFoundError: 

935 pass 

936 

937 

938def _lfric_time_callback(cube: iris.cube.Cube): 

939 """Fix time coordinate metadata if missing dimensions. 

940 

941 Some model data does not contain forecast_reference_time or forecast_period as 

942 expected coordinates, and so we cannot aggregate over case studies without this 

943 metadata. This callback fixes these issues. 

944 

945 This callback also ensures all time coordinates are referenced as hours since 

946 1970-01-01 00:00:00 for consistency across different model inputs. 

947 

948 Notes 

949 ----- 

950 Some parts of the code have been adapted from Paul Earnshaw's scripts. 

951 """ 

952 # Construct forecast_reference time if it doesn't exist. 

953 try: 

954 tcoord = cube.coord("time") 

955 # Set time coordinate to common basis "hours since 1970" 

956 try: 

957 tcoord.convert_units("hours since 1970-01-01 00:00:00") 

958 except ValueError: 

959 logging.warning("Unrecognised base time unit: %s", tcoord.units) 

960 

961 if not cube.coords("forecast_reference_time"): 

962 try: 

963 init_time = datetime.datetime.fromisoformat( 

964 tcoord.attributes["time_origin"] 

965 ) 

966 frt_point = tcoord.units.date2num(init_time) 

967 frt_coord = iris.coords.AuxCoord( 

968 frt_point, 

969 units=tcoord.units, 

970 standard_name="forecast_reference_time", 

971 long_name="forecast_reference_time", 

972 ) 

973 cube.add_aux_coord(frt_coord) 

974 except KeyError: 

975 logging.warning( 

976 "Cannot find forecast_reference_time, but no `time_origin` attribute to construct it from." 

977 ) 

978 

979 # Remove time_origin to allow multiple case studies to merge. 

980 tcoord.attributes.pop("time_origin", None) 

981 

982 # Construct forecast_period axis (forecast lead time) if it doesn't exist. 

983 if not cube.coords("forecast_period"): 

984 try: 

985 # Create array of forecast lead times. 

986 init_coord = cube.coord("forecast_reference_time") 

987 init_time_points_in_tcoord_units = tcoord.units.date2num( 

988 init_coord.units.num2date(init_coord.points) 

989 ) 

990 lead_times = tcoord.points - init_time_points_in_tcoord_units 

991 

992 # Get unit for lead time from time coordinate's unit. 

993 # Convert all lead time to hours for consistency between models. 

994 if "seconds" in str(tcoord.units): 994 ↛ 995line 994 didn't jump to line 995 because the condition on line 994 was never true

995 lead_times = lead_times / 3600.0 

996 units = "hours" 

997 elif "hours" in str(tcoord.units): 997 ↛ 1000line 997 didn't jump to line 1000 because the condition on line 997 was always true

998 units = "hours" 

999 else: 

1000 raise ValueError(f"Unrecognised base time unit: {tcoord.units}") 

1001 

1002 # Create lead time coordinate. 

1003 lead_time_coord = iris.coords.AuxCoord( 

1004 lead_times, 

1005 standard_name="forecast_period", 

1006 long_name="forecast_period", 

1007 units=units, 

1008 ) 

1009 

1010 # Associate lead time coordinate with time dimension. 

1011 cube.add_aux_coord(lead_time_coord, cube.coord_dims("time")) 

1012 except iris.exceptions.CoordinateNotFoundError: 

1013 logging.warning( 

1014 "Cube does not have both time and forecast_reference_time coordinate, so cannot construct forecast_period" 

1015 ) 

1016 except iris.exceptions.CoordinateNotFoundError: 

1017 logging.warning("No time coordinate on cube.") 

1018 

1019 

1020def _lfric_forecast_period_callback(cube: iris.cube.Cube): 

1021 """Check forecast_period name and units.""" 

1022 try: 

1023 coord = cube.coord("forecast_period") 

1024 if coord.units != "hours": 

1025 cube.coord("forecast_period").convert_units("hours") 

1026 if not coord.standard_name: 

1027 coord.standard_name = "forecast_period" 

1028 except iris.exceptions.CoordinateNotFoundError: 

1029 pass 

1030 

1031 

1032def _normalise_ML_varname(cube: iris.cube.Cube): 

1033 """Fix plev variable names to standard names.""" 

1034 if cube.coords("pressure"): 

1035 if cube.name() == "x_wind": 

1036 cube.long_name = "zonal_wind_at_pressure_levels" 

1037 if cube.name() == "y_wind": 

1038 cube.long_name = "meridional_wind_at_pressure_levels" 

1039 if cube.name() == "air_temperature": 

1040 cube.long_name = "temperature_at_pressure_levels" 

1041 if cube.name() == "specific_humidity": 1041 ↛ 1042line 1041 didn't jump to line 1042 because the condition on line 1041 was never true

1042 cube.long_name = ( 

1043 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging" 

1044 )