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

439 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-24 08:19 +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 

44logger = logging.getLogger(__name__) 

45 

46 

47class NoDataError(FileNotFoundError): 

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

49 

50 

51def read_cube( 

52 file_paths: list[str] | str, 

53 constraint: iris.Constraint = None, 

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

55 subarea_type: str | None = None, 

56 subarea_extent: list[float] | None = None, 

57 **kwargs, 

58) -> iris.cube.Cube: 

59 """Read a single cube from files. 

60 

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

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

63 directory, all the files contained within are loaded. 

64 

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

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

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

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

69 

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

71 processed in the same way as ensemble data. 

72 

73 Arguments 

74 --------- 

75 file_paths: str | list[str] 

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

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

78 Constraints to filter data by. Defaults to unconstrained. 

79 model_names: str | list[str], optional 

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

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

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

83 coordinates. 

84 subarea_extent: list, optional 

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

86 upper latitude, lower longitude, upper longitude. 

87 

88 Returns 

89 ------- 

90 cubes: iris.cube.Cube 

91 Cube loaded 

92 

93 Raises 

94 ------ 

95 FileNotFoundError 

96 If the provided path does not exist 

97 ValueError 

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

99 """ 

100 cubes = read_cubes( 

101 file_paths=file_paths, 

102 constraint=constraint, 

103 model_names=model_names, 

104 subarea_type=subarea_type, 

105 subarea_extent=subarea_extent, 

106 ) 

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

108 if len(cubes) == 1: 

109 return cubes[0] 

110 else: 

111 raise ValueError( 

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

113 ) 

114 

115 

116def read_cubes( 

117 file_paths: list[str] | str, 

118 constraint: iris.Constraint | None = None, 

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

120 subarea_type: str | None = None, 

121 subarea_extent: list | None = None, 

122 **kwargs, 

123) -> iris.cube.CubeList: 

124 """Read cubes from files. 

125 

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

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

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

129 

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

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

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

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

134 

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

136 processed in the same way as ensemble data. 

137 

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

139 that the cubes merge across files. 

140 

141 Arguments 

142 --------- 

143 file_paths: str | list[str] 

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

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

146 Constraints to filter data by. Defaults to unconstrained. 

147 model_names: str | list[str], optional 

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

149 subarea_type: str, optional 

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

151 coordinates. 

152 subarea_extent: list[float], optional 

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

154 upper latitude, lower longitude, upper longitude. 

155 

156 Returns 

157 ------- 

158 cubes: iris.cube.CubeList 

159 Cubes loaded after being merged and concatenated. 

160 

161 Raises 

162 ------ 

163 FileNotFoundError 

164 If the provided path does not exist 

165 """ 

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

167 paths = iter_maybe(file_paths) 

168 

169 model_names = iter_maybe(model_names) 

170 

171 # flattens model_names if needed into one dimensional list. 

172 if model_names != (None,): 

173 flat = [] 

174 for item in model_names: 

175 if isinstance(item, list): 175 ↛ 176line 175 didn't jump to line 176 because the condition on line 175 was never true

176 flat.extend(item) 

177 else: 

178 flat.append(item) 

179 model_names = flat 

180 

181 # Check we have appropriate number of model names. 

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

183 raise ValueError( 

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

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

186 ) 

187 

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

189 model_cubes = ( 

190 _load_model(path, name, constraint) 

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

192 ) 

193 

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

195 cubes = next(model_cubes) 

196 for cube in cubes: 

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

198 cube.attributes["cset_comparison_base"] = 1 

199 

200 # Load the rest of the models. 

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

202 

203 # Enable different point-based observation sources to be concatenated. 

204 cubes = _check_combine_point_observations(cubes) 

205 

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

207 iris.util.unify_time_units(cubes) 

208 

209 # Select sub region. 

210 cubes = _cutout_cubes(cubes, subarea_type, subarea_extent) 

211 

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

213 cubes = _merge_cubes_check_ensemble(cubes) 

214 cubes = cubes.concatenate() 

215 

216 # Squeeze single valued coordinates into scalar coordinates. 

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

218 

219 # Ensure dimension coordinates are bounded. 

220 for cube in cubes: 

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

222 if (dim_coord.standard_name == "time") and ( 

223 dim_coord.name() 

224 not in itertools.chain.from_iterable( 

225 m.coord_names for m in cube.cell_methods if m.method != "point" 

226 ) 

227 ): 

228 # Instantaneous time coordinate 

229 continue 

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

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

232 dim_coord.guess_bounds() 

233 

234 logger.info("Loaded cubes: %s", cubes) 

235 if len(cubes) == 0: 

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

237 return cubes 

238 

239 

240def _load_model( 

241 paths: str | list[str], 

242 model_name: str | None, 

243 constraint: iris.Constraint | None, 

244) -> iris.cube.CubeList: 

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

246 input_files = _check_input_files(paths) 

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

248 logger.debug("Constraint: %s", constraint) 

249 

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

251 # If required, compute wind_speed from components. 

252 cubes = _compute_winds(cubes) 

253 

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

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

256 if model_name is not None: 

257 for cube in cubes: 

258 cube.attributes["model_name"] = model_name 

259 return cubes 

260 

261 

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

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

264 

265 Arguments 

266 --------- 

267 input_paths: list[str] 

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

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

270 

271 Returns 

272 ------- 

273 list[Path] 

274 A list of files to load. 

275 

276 Raises 

277 ------ 

278 FileNotFoundError: 

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

280 """ 

281 files = [] 

282 for raw_filename in iter_maybe(input_paths): 

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

284 raw_path = Path(raw_filename) 

285 if raw_path.is_file(): 

286 files.append(raw_path) 

287 else: 

288 for input_path in glob.glob(raw_filename): 

289 # Convert string paths into Path objects. 

290 input_path = Path(input_path) 

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

292 if input_path.is_dir(): 

293 logger.debug("Checking directory '%s' for files", input_path) 

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

295 else: 

296 files.append(input_path) 

297 

298 files.sort() 

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

300 if len(files) == 0: 

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

302 return files 

303 

304 

305def _merge_cubes_check_ensemble(cubes: iris.cube.CubeList): 

306 """Merge CubeList, renumbering realizations of 0 if required. 

307 

308 An unsuccessful merge indicates common input cube attributes, most 

309 commonly from ensemble members missing an explicit realization 

310 coordinate. Therefore the members are renumbered before being merged 

311 again. 

312 """ 

313 try: 

314 cubes = cubes.merge() 

315 except iris.exceptions.MergeError: 

316 _log_once( 

317 "Attempt to merge input CubeList failed. Attempting to iterate realization coords to enable merge.", 

318 level=logging.WARNING, 

319 ) 

320 for ir, cube in enumerate(cubes): 

321 if cube.coord("realization").points == 0: 321 ↛ 320line 321 didn't jump to line 320 because the condition on line 321 was always true

322 cube.coord("realization").points = ir + 1 

323 cubes = cubes.merge() 

324 return cubes 

325 

326 

327def _cutout_cubes( 

328 cubes: iris.cube.CubeList, 

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

330 subarea_extent: list[float], 

331): 

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

333 if subarea_type is None: 

334 logger.debug("Subarea selection is disabled.") 

335 return cubes 

336 

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

338 cutout_cubes = iris.cube.CubeList() 

339 # Find spatial coordinates 

340 for cube in cubes: 

341 # Find dimension coordinates. 

342 lat_name, lon_name = get_cube_yxcoordname(cube) 

343 

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

345 if subarea_type == "gridcells": 

346 logger.debug( 

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

348 subarea_extent[0], 

349 subarea_extent[1], 

350 subarea_extent[2], 

351 subarea_extent[3], 

352 ) 

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

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

355 # Define cutout region using user provided cell points. 

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

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

358 

359 # Compute cutout based on specified coordinate values. 

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

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

362 logger.debug( 

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

364 subarea_extent[0], 

365 subarea_extent[1], 

366 subarea_extent[2], 

367 subarea_extent[3], 

368 ) 

369 # Define cutout region using user provided coordinates. 

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

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

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

373 while lons[0] < -180.0: 

374 lons += 360.0 

375 while lons[1] > 180.0: 

376 lons -= 360.0 

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

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

379 coord_system = cube.coord(lat_name).coord_system 

380 if subarea_type == "realworld" and isinstance( 

381 coord_system, iris.coord_systems.RotatedGeogCS 

382 ): 

383 lons, lats = rotate_pole( 

384 lons, 

385 lats, 

386 pole_lon=coord_system.grid_north_pole_longitude, 

387 pole_lat=coord_system.grid_north_pole_latitude, 

388 ) 

389 else: 

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

391 

392 # Do cutout and add to cutout_cubes. 

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

394 logger.debug("Cutting out coords: %s", intersection_args) 

395 try: 

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

397 except IndexError as err: 

398 raise ValueError( 

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

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

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

402 ) from err 

403 

404 return cutout_cubes 

405 

406 

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

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

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

410 _realization_callback(cube) 

411 _um_normalise_callback(cube) 

412 _lfric_normalise_callback(cube) 

413 _nimrod_normalise_callback(cube) 

414 cube = _lfric_time_coord_fix_callback(cube) 

415 _normalise_var0_varname(cube) 

416 cube = _fix_no_spatial_coords_callback(cube) 

417 _fix_spatial_coords_callback(cube) 

418 _fix_pressure_coord_callback(cube) 

419 _fix_um_radtime(cube) 

420 _fix_cell_methods(cube) 

421 cube = _convert_cube_units_callback(cube) 

422 cube = _grid_longitude_fix_callback(cube) 

423 _fix_lfric_cloud_base_altitude(cube) 

424 _proleptic_gregorian_fix(cube) 

425 _lfric_time_callback(cube) 

426 _lfric_forecast_period_callback(cube) 

427 cube = _fix_no_time_coords_callback(cube) 

428 _normalise_longname(cube) 

429 return cube 

430 

431 

432def _realization_callback(cube): 

433 """Add a realization coordinate initialised to 0 if missing. 

434 

435 This means deterministic and ensemble cubes can assume realization coordinate through the rest 

436 of the code. 

437 """ 

438 # Only add if realization coordinate does not exist. 

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

440 cube.add_aux_coord( 

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

442 ) 

443 

444 

445@functools.lru_cache(None) 

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

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

448 logger.log(level, msg) 

449 

450 

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

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

453 

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

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

456 """ 

457 # Convert STASH to LFRic variable name 

458 if "STASH" in cube.attributes: 

459 stash = cube.attributes["STASH"] 

460 try: 

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

462 cube.long_name = name 

463 except KeyError: 

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

465 _log_once( 

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

467 level=logging.WARNING, 

468 ) 

469 

470 

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

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

473 

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

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

476 

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

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

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

480 has been converted to look like UM data. 

481 """ 

482 # Remove unwanted attributes. 

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

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

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

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

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

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

489 

490 # Sort STASH code list. 

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

492 if stash_list: 

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

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

495 

496 

497def _nimrod_normalise_callback(cube: iris.cube.Cube): 

498 """Normalise attributes that prevents NIMROD radar cubes from merging.""" 

499 # Remove unwanted attributes. 

500 cube.attributes.pop("radar_sites", None) 

501 cube.attributes.pop("additional_radar_sites", None) 

502 cube.attributes.pop("recursive_filter_iterations", None) 

503 cube.attributes.pop("Probability methods", None) 

504 

505 

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

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

508 

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

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

511 Scalar time values are left as AuxCoords. 

512 """ 

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

514 # always ends up as an AuxCoord. 

515 if cube.coords("time"): 

516 time_coord = cube.coord("time") 

517 if ( 

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

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

520 ): 

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

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

523 time_coord.has_bounds() 

524 and (time_coord.bounds[-1][0] - time_coord.bounds[0][0]) < 1.0e-8 

525 ): 

526 time_coord.bounds = [ 

527 [ 

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

529 time_coord.bounds[i][1], 

530 ] 

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

532 ] 

533 iris.util.promote_aux_coord_to_dim_coord(cube, time_coord) 

534 return cube 

535 

536 

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

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

539 

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

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

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

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

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

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

546 """ 

547 try: 

548 y, x = get_cube_yxcoordname(cube) 

549 except ValueError: 

550 # Don't modify non-spatial cubes. 

551 return cube 

552 

553 long_coord = cube.coord(x) 

554 # Wrap longitudes if rotated pole coordinates 

555 coord_system = long_coord.coord_system 

556 if x == "grid_longitude" and isinstance( 

557 coord_system, iris.coord_systems.RotatedGeogCS 

558 ): 

559 long_points = long_coord.points.copy() 

560 long_centre = np.median(long_points) 

561 while long_centre < -175.0: 

562 long_centre += 360.0 

563 long_points += 360.0 

564 while long_centre >= 175.0: 

565 long_centre -= 360.0 

566 long_points -= 360.0 

567 long_coord.points = long_points 

568 

569 # Update coord bounds to be consistent with wrapping. 

570 if long_coord.has_bounds(): 

571 long_coord.bounds = None 

572 long_coord.guess_bounds() 

573 

574 return cube 

575 

576 

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

578 import CSET.operators._utils as utils 

579 

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

581 if utils.is_spatialdim(cube): 

582 return cube 

583 

584 else: 

585 # attempt to get lat/long from cube attributes 

586 try: 

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

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

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

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

591 

592 lon_val = (lon_min + lon_max) / 2.0 

593 lat_val = (lat_min + lat_max) / 2.0 

594 

595 lat_coord = iris.coords.DimCoord( 

596 lat_val, 

597 standard_name="latitude", 

598 units="degrees_north", 

599 var_name="latitude", 

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

601 circular=True, 

602 ) 

603 

604 lon_coord = iris.coords.DimCoord( 

605 lon_val, 

606 standard_name="longitude", 

607 units="degrees_east", 

608 var_name="longitude", 

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

610 circular=True, 

611 ) 

612 

613 cube.add_aux_coord(lat_coord) 

614 cube.add_aux_coord(lon_coord) 

615 return cube 

616 

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

618 except TypeError: 

619 return cube 

620 

621 

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

623 """Check latitude and longitude coordinates name. 

624 

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

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

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

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

629 particularly where comparing multiple input models with differing spatial 

630 coordinates. 

631 """ 

632 # Check if cube is spatial. 

633 if not is_spatialdim(cube): 

634 # Don't modify non-spatial cubes. 

635 return 

636 

637 # Get spatial coords and dimension index. 

638 y_name, x_name = get_cube_yxcoordname(cube) 

639 ny = get_cube_coordindex(cube, y_name) 

640 nx = get_cube_coordindex(cube, x_name) 

641 

642 # Remove spatial coords bounds if erroneous values detected. 

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

644 # invalid threshold of 10000.0 

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

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

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

648 if bx_max > 10000.0 or by_max > 10000.0: 

649 cube.coord(x_name).bounds = None 

650 cube.coord(y_name).bounds = None 

651 

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

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

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

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

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

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

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

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

660 

661 cube.remove_coord("grid_latitude") 

662 cube.add_dim_coord( 

663 iris.coords.DimCoord( 

664 lats, 

665 standard_name="latitude", 

666 var_name="latitude", 

667 units="degrees", 

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

669 circular=True, 

670 ), 

671 ny, 

672 ) 

673 y_name = "latitude" 

674 cube.remove_coord("grid_longitude") 

675 cube.add_dim_coord( 

676 iris.coords.DimCoord( 

677 lons, 

678 standard_name="longitude", 

679 var_name="longitude", 

680 units="degrees", 

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

682 circular=True, 

683 ), 

684 nx, 

685 ) 

686 x_name = "longitude" 

687 

688 # Create additional AuxCoord [grid_latitude, grid_longitude] with 

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

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

691 "degrees", 

692 "degrees_north", 

693 "degrees_south", 

694 ]: 

695 # Add grid_latitude AuxCoord 

696 if "grid_latitude" not in [ 

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

698 ]: 

699 cube.add_aux_coord( 

700 iris.coords.AuxCoord( 

701 cube.coord(y_name).points, 

702 var_name="grid_latitude", 

703 units="degrees", 

704 ), 

705 ny, 

706 ) 

707 # Ensure input latitude DimCoord has CoordSystem 

708 # This attribute is sometimes lost on iris.save 

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

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

711 

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

713 "degrees", 

714 "degrees_west", 

715 "degrees_east", 

716 ]: 

717 # Add grid_longitude AuxCoord 

718 if "grid_longitude" not in [ 

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

720 ]: 

721 cube.add_aux_coord( 

722 iris.coords.AuxCoord( 

723 cube.coord(x_name).points, 

724 var_name="grid_longitude", 

725 units="degrees", 

726 ), 

727 nx, 

728 ) 

729 

730 # Ensure input longitude DimCoord has CoordSystem 

731 # This attribute is sometimes lost on iris.save 

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

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

734 

735 

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

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

738 

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

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

741 than compliant CF coordinate names. 

742 

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

744 and approach the coordinates in a unified way. 

745 """ 

746 for coord in cube.dim_coords: 

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

748 coord.rename("pressure") 

749 

750 if coord.name() == "pressure" and str(cube.coord("pressure").units) != "hPa": 

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

752 

753 

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

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

756 

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

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

759 diagnostics are checked. 

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

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

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

763 """ 

764 try: 

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

766 "m01s01i207", 

767 "m01s01i208", 

768 "m01s02i205", 

769 "m01s02i201", 

770 "m01s01i207", 

771 "m01s02i207", 

772 "m01s01i235", 

773 ]: 

774 time_coord = cube.coord("time") 

775 

776 # Convert time points to datetime objects 

777 time_unit = time_coord.units 

778 time_points = time_unit.num2date(time_coord.points) 

779 # Skip if times don't need fixing. 

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

781 return 

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

783 return 

784 

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

786 n_minute = time_points[0].minute 

787 n_second = time_points[0].second 

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

789 if n_minute > 30: 

790 n_minute = n_minute - 60 

791 # Compute new diagnostic time stamp 

792 new_time_points = ( 

793 time_points 

794 - datetime.timedelta(minutes=n_minute) 

795 - datetime.timedelta(seconds=n_second) 

796 ) 

797 

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

799 new_time_values = time_unit.date2num(new_time_points) 

800 

801 # Replace the time coordinate with updated values. 

802 time_coord.points = new_time_values 

803 

804 # Recompute forecast_period with corrected values. 

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

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

807 new_fcst_points = ( 

808 time_unit.num2date(fcst_prd_points) 

809 - datetime.timedelta(minutes=n_minute) 

810 - datetime.timedelta(seconds=n_second) 

811 ) 

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

813 new_fcst_points 

814 ) 

815 except KeyError: 

816 pass 

817 

818 

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

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

821 

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

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

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

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

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

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

828 """ 

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

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

831 "m01s21i104", 

832 "m01s04i201", 

833 "m01s04i202", 

834 "m01s05i201", 

835 "m01s05i202", 

836 ] and {cm.method for cm in cube.cell_methods} == {"mean"}: 

837 # Retrieve interval and any comment information. 

838 for cell_method in cube.cell_methods: 

839 interval_str = cell_method.intervals 

840 comment_str = cell_method.comments 

841 

842 # Remove input aggregation method. 

843 cube.cell_methods = () 

844 

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

846 cube.add_cell_method( 

847 iris.coords.CellMethod( 

848 method="sum", 

849 coords="time", 

850 intervals=interval_str, 

851 comments=comment_str, 

852 ) 

853 ) 

854 

855 

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

857 """Adjust diagnostic units for specific variables. 

858 

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

860 converted here to mm hr-1. 

861 

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

863 formatting. 

864 """ 

865 # Convert precipitation diagnostic units if required. 

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

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

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

869 _log_once( 

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

871 level=logging.DEBUG, 

872 ) 

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

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

875 cube.units = "mm s-1" 

876 # Convert the units to per hour. 

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

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

879 _log_once( 

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

881 level=logging.DEBUG, 

882 ) 

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

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

885 cube.units = "mm" 

886 

887 # Convert visibility diagnostic units if required. 

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

889 if any("visibility" in name for name in varnames) and cube.units == "m": 

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

891 # Convert the units to km. 

892 cube.convert_units("km") 

893 

894 return cube 

895 

896 

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

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

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

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

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

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

903 

904 

905def _compute_winds(cubes: iris.cube.CubeList): 

906 """To compute wind_speed from vector components if not available as diagnostic. 

907 

908 Diagnostics of wind are also not always consistent between the UM 

909 and LFRic. Here, winds from the UM are adjusted to make them 

910 consistent with LFRic. 

911 """ 

912 # Check whether we have components of the wind identified by varname 

913 # but not the wind speed and calculate it if it is missing. Note that 

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

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

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

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

918 # 

919 # A check on UM STASH attributes is also conducted to adjust directions. 

920 u_constr = iris.Constraint("eastward_wind_at_10m") 

921 v_constr = iris.Constraint("northward_wind_at_10m") 

922 speed_constr = iris.Constraint("wind_speed_at_10m") 

923 try: 

924 if cubes.extract(u_constr) and cubes.extract(v_constr): 

925 if len(cubes) == 2: 

926 wind_only = True 

927 else: 

928 wind_only = False 

929 if len(cubes.extract(u_constr)) == 1 and not cubes.extract(speed_constr): 929 ↛ 932line 929 didn't jump to line 932 because the condition on line 929 was always true

930 _add_wind_speed_um(cubes) 

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

932 if cubes.extract(u_constr) and cubes.extract(v_constr): 932 ↛ 935line 932 didn't jump to line 935 because the condition on line 932 was always true

933 _convert_wind_true_dirn_um(cubes) 

934 # Return only wind_speed cube 

935 if wind_only: 

936 cubes = cubes.extract(speed_constr) 

937 except (KeyError, AttributeError): 

938 pass 

939 

940 return cubes 

941 

942 

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

944 """Add windspeeds to cubes from components.""" 

945 u_wind = cubes.extract_cube(iris.Constraint("eastward_wind_at_10m")) 

946 v_wind = cubes.extract_cube(iris.Constraint("northward_wind_at_10m")) 

947 wspd10 = (u_wind**2 + v_wind**2) ** 0.5 

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

949 wspd10.standard_name = "wind_speed" 

950 wspd10.long_name = "wind_speed_at_10m" 

951 wspd10.units = "ms-1" 

952 cubes.append(wspd10) 

953 

954 

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

956 """To convert winds to true directions. 

957 

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

959 This functionality only handles the simplest case. 

960 Constrains using STASH code only to ensure applied to UM outputs only. 

961 """ 

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

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

964 for u, v in zip(u_grids, v_grids, strict=True): 964 ↛ 965line 964 didn't jump to line 965 because the loop on line 964 never started

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

966 u.data = true_u.core_data() 

967 v.data = true_v.core_data() 

968 

969 

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

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

972 

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

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

975 with different attributes. This can be inconsistently managed in 

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

977 """ 

978 for coord in cube.coords(): 

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

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

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

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

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

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

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

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

987 

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

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

990 

991 

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

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

994 try: 

995 time_coord = cube.coord("time") 

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

997 logger.debug( 

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

999 repr(time_coord.units), 

1000 ) 

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

1002 except iris.exceptions.CoordinateNotFoundError: 

1003 pass 

1004 

1005 

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

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

1008 

1009 Some model data does not contain forecast_reference_time or forecast_period as 

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

1011 metadata. This callback fixes these issues. 

1012 

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

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

1015 

1016 Notes 

1017 ----- 

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

1019 """ 

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

1021 try: 

1022 tcoord = cube.coord("time") 

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

1024 try: 

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

1026 except ValueError: 

1027 logger.warning("Unrecognised base time unit: %s", tcoord.units) 

1028 

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

1030 try: 

1031 init_time = datetime.datetime.fromisoformat( 

1032 tcoord.attributes["time_origin"] 

1033 ) 

1034 frt_point = tcoord.units.date2num(init_time) 

1035 frt_coord = iris.coords.AuxCoord( 

1036 frt_point, 

1037 units=tcoord.units, 

1038 standard_name="forecast_reference_time", 

1039 long_name="forecast_reference_time", 

1040 ) 

1041 cube.add_aux_coord(frt_coord) 

1042 except KeyError: 

1043 logger.warning( 

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

1045 ) 

1046 

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

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

1049 

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

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

1052 try: 

1053 # Create array of forecast lead times. 

1054 init_coord = cube.coord("forecast_reference_time") 

1055 init_time_points_in_tcoord_units = tcoord.units.date2num( 

1056 init_coord.units.num2date(init_coord.points) 

1057 ) 

1058 lead_times = tcoord.points - init_time_points_in_tcoord_units 

1059 

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

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

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

1063 lead_times = lead_times / 3600.0 

1064 units = "hours" 

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

1066 units = "hours" 

1067 else: 

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

1069 

1070 # Create lead time coordinate. 

1071 lead_time_coord = iris.coords.AuxCoord( 

1072 lead_times, 

1073 standard_name="forecast_period", 

1074 long_name="forecast_period", 

1075 units=units, 

1076 ) 

1077 

1078 # Associate lead time coordinate with time dimension. 

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

1080 except iris.exceptions.CoordinateNotFoundError: 

1081 logger.warning( 

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

1083 ) 

1084 except iris.exceptions.CoordinateNotFoundError: 

1085 logger.warning("No time coordinate on cube.") 

1086 

1087 

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

1089 """Check forecast_period name and units.""" 

1090 try: 

1091 coord = cube.coord("forecast_period") 

1092 if coord.units != "hours": 

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

1094 if not coord.standard_name: 

1095 coord.standard_name = "forecast_period" 

1096 except iris.exceptions.CoordinateNotFoundError: 

1097 pass 

1098 

1099 

1100def _fix_no_time_coords_callback(cube: iris.cube.Cube): 

1101 """Add dummy time coord to process cubes that don't have sequence coord.""" 

1102 # Only add if time coordinate does not exist. 

1103 if not cube.coords("time"): 

1104 cube.add_aux_coord( 

1105 iris.coords.DimCoord( 

1106 0, standard_name="time", units="hours since 0001-01-01 00:00:00" 

1107 ) 

1108 ) 

1109 

1110 return cube 

1111 

1112 

1113def _normalise_longname(cube: iris.cube.Cube): 

1114 """Normalise long_name to the LFRic standard list.""" 

1115 if cube.coords("pressure"): 

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

1117 cube.long_name = "zonal_wind_at_pressure_levels" 

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

1119 cube.long_name = "meridional_wind_at_pressure_levels" 

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

1121 cube.long_name = "temperature_at_pressure_levels" 

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

1123 cube.long_name = ( 

1124 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging" 

1125 ) 

1126 else: 

1127 if cube.name() == "x_wind" and cube.var_name == "u_wind_at_10m": 1127 ↛ 1128line 1127 didn't jump to line 1128 because the condition on line 1127 was never true

1128 cube.long_name = "eastward_wind_at_10m" 

1129 if cube.name() == "y_wind" and cube.var_name == "v_wind_at_10m": 1129 ↛ 1130line 1129 didn't jump to line 1130 because the condition on line 1129 was never true

1130 cube.long_name = "northward_wind_at_10m" 

1131 if cube.name() == "air_pressure_at_sea_level": 

1132 cube.long_name = "air_pressure_at_mean_sea_level" 

1133 

1134 

1135def _check_combine_point_observations(cubes: iris.cube.CubeList): 

1136 """Enable cubes containing different point observation sources to be concatenated.""" 

1137 nstation = 0 

1138 for cube in cubes: 

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

1140 if "obs_source" in [coord.name() for coord in cube.coords()]: 

1141 cube.remove_coord("obs_source") 

1142 cube.coord("station").points = cube.coord("station").points + nstation 

1143 nstation = nstation + len(cube.coord("station").points) 

1144 

1145 return cubes