Coverage for src/CSET/operators/_utils.py: 87%

260 statements  

« prev     ^ index     » next       coverage.py v7.16.1, created at 2026-09-22 13:57 +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""" 

16Common operator functionality used across CSET. 

17 

18Functions below should only be added if it is not suitable as a standalone 

19operator, and will be used across multiple operators. 

20""" 

21 

22import logging 

23import os 

24import re 

25from datetime import timedelta 

26from pathlib import Path 

27 

28import iris 

29import iris.coords 

30import iris.cube 

31import iris.exceptions 

32import iris.util 

33import numpy as np 

34from iris.time import PartialDateTime 

35 

36from CSET._common import iter_maybe 

37 

38logger = logging.getLogger(__name__) 

39 

40 

41def pdt_fromisoformat( 

42 datestring, 

43) -> tuple[iris.time.PartialDateTime, timedelta | None]: 

44 """Generate PartialDateTime object. 

45 

46 Function that takes an ISO 8601 date string and returns a PartialDateTime object. 

47 

48 Arguments 

49 --------- 

50 datestring: str 

51 ISO 8601 date. 

52 

53 Returns 

54 ------- 

55 time_object: iris.time.PartialDateTime 

56 """ 

57 

58 def make_offset(sign, value) -> timedelta: 

59 if len(value) not in [2, 4, 5]: 59 ↛ 60line 59 didn't jump to line 60 because the condition on line 59 was never true

60 raise ValueError(f'expected "hh", "hhmm", or "hh:mm", got {value}') 

61 

62 hours = int(value[:2]) 

63 minutes = 0 

64 if len(value) in [4, 5]: 

65 minutes = int(value[-2:]) 

66 return timedelta(hours=sign * hours, minutes=sign * minutes) 

67 

68 # Remove the microseconds coord due to no support in PartialDateTime 

69 datestring = re.sub(r"\.\d+", "", datestring) 

70 

71 datetime_split = datestring.split("T") 

72 date = datetime_split[0] 

73 if len(datetime_split) == 1: 

74 time = "" 

75 elif len(datetime_split) == 2: 75 ↛ 78line 75 didn't jump to line 78 because the condition on line 75 was always true

76 time = datetime_split[1] 

77 else: 

78 raise ValueError("datesting in an unexpected format") 

79 

80 offset = None 

81 time_split = time.split("+") 

82 if len(time_split) == 2: 

83 time = time_split[0] 

84 offset = make_offset(1, time_split[1]) 

85 else: 

86 time_split = time.split("-") 

87 if len(time_split) == 2: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true

88 time = time_split[0] 

89 offset = make_offset(-1, time_split[1]) 

90 else: 

91 offset = None 

92 

93 if re.fullmatch(r"\d{8}", date): 

94 date = f"{date[0:4]}-{date[4:6]}-{date[6:8]}" 

95 elif re.fullmatch(r"\d{6}", date): 

96 date = f"{date[0:4]}-{date[4:6]}" 

97 

98 if len(date) < 7: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true

99 raise ValueError(f"Invalid datestring: {datestring}, must be at least YYYY-MM") 

100 

101 # Returning a PartialDateTime for the special case of string form "YYYY-MM" 

102 if re.fullmatch(r"\d{4}-\d{2}", date): 

103 pdt = PartialDateTime( 

104 year=int(date[0:4]), 

105 month=int(date[5:7]), 

106 day=None, 

107 hour=None, 

108 minute=None, 

109 second=None, 

110 ) 

111 return pdt, offset 

112 

113 year = int(date[0:4]) 

114 month = int(date[5:7]) 

115 day = int(date[8:10]) 

116 

117 kwargs = { 

118 "year": year, 

119 "month": month, 

120 "day": day, 

121 "hour": 0, 

122 "minute": 0, 

123 "second": 0, 

124 } 

125 

126 # Normalise the time parts into standard format 

127 if re.fullmatch(r"\d{4}", time): 127 ↛ 128line 127 didn't jump to line 128 because the condition on line 127 was never true

128 time = f"{time[0:2]}:{time[2:4]}" 

129 if re.fullmatch(r"\d{6}", time): 

130 time = f"{time[0:2]}:{time[2:4]}:{time[4:6]}" 

131 

132 if len(time) >= 2: 

133 kwargs["hour"] = int(time[0:2]) 

134 if len(time) >= 5: 

135 kwargs["minute"] = int(time[3:5]) 

136 if len(time) >= 8: 

137 kwargs["second"] = int(time[6:8]) 

138 

139 pdt = PartialDateTime(**kwargs) 

140 

141 return pdt, offset 

142 

143 

144def get_cube_yxcoordname(cube: iris.cube.Cube) -> tuple[str, str]: 

145 """ 

146 Return horizontal dimension coordinate name(s) from a given cube. 

147 

148 Arguments 

149 --------- 

150 

151 cube: iris.cube.Cube 

152 An iris cube which will be checked to see if it contains coordinate 

153 names that match a pre-defined list of acceptable horizontal 

154 dimension coordinate names. 

155 

156 Returns 

157 ------- 

158 (y_coord, x_coord) 

159 A tuple containing the horizontal coordinate name for latitude and longitude respectively 

160 found within the cube. 

161 

162 Raises 

163 ------ 

164 ValueError 

165 If a unique y/x horizontal coordinate cannot be found. 

166 """ 

167 # Acceptable horizontal coordinate names. 

168 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"] 

169 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"] 

170 

171 # Get a list of dimension coordinate names for the cube 

172 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)] 

173 coord_names = [coord.name() for coord in cube.coords()] 

174 

175 # Check which x-coordinate we have, if any 

176 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES] 

177 if len(x_coords) != 1: 

178 x_coords = [coord for coord in dim_coord_names if coord in X_COORD_NAMES] 

179 if len(x_coords) != 1: 

180 raise ValueError("Could not identify a unique x-coordinate in cube") 

181 

182 # Check which y-coordinate we have, if any 

183 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES] 

184 if len(y_coords) != 1: 

185 y_coords = [coord for coord in dim_coord_names if coord in Y_COORD_NAMES] 

186 if len(y_coords) != 1: 

187 raise ValueError("Could not identify a unique y-coordinate in cube") 

188 

189 return (y_coords[0], x_coords[0]) 

190 

191 

192def get_cube_coordindex(cube: iris.cube.Cube, coord_name) -> int: 

193 """ 

194 Return coordinate dimension for a named coordinate from a given cube. 

195 

196 Arguments 

197 --------- 

198 

199 cube: iris.cube.Cube 

200 An iris cube which will be checked to see if it contains coordinate 

201 names that match a pre-defined list of acceptable horizontal 

202 coordinate names. 

203 

204 coord_name: str 

205 A cube dimension name 

206 

207 Returns 

208 ------- 

209 coord_index 

210 An integer specifying where in the cube dimension list a specified coordinate name is found. 

211 

212 Raises 

213 ------ 

214 ValueError 

215 If a specified dimension coordinate cannot be found. 

216 """ 

217 # Get a list of dimension coordinate names for the cube 

218 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)] 

219 

220 # Check if requested dimension is found in cube and get index 

221 if coord_name in coord_names: 

222 coord_index = cube.coord_dims(coord_name)[0] 

223 else: 

224 raise ValueError("Could not find requested dimension %s", coord_name) 

225 

226 return coord_index 

227 

228 

229def is_spatialdim(cube: iris.cube.Cube) -> bool: 

230 """Determine whether a cube is has two spatial dimension coordinates. 

231 

232 If cube has both spatial dims, it will contain two unique coordinates 

233 that explain space (latitude and longitude). The coordinates have to 

234 be iterable/contain usable dimension data, as cubes may contain these 

235 coordinates as scalar dimensions after being collapsed. 

236 

237 Arguments 

238 --------- 

239 cube: iris.cube.Cube 

240 An iris cube which will be checked to see if it contains coordinate 

241 names that match a pre-defined list of acceptable coordinate names. 

242 

243 Returns 

244 ------- 

245 bool 

246 If true, then the cube has a spatial projection and thus can be plotted 

247 as a map. 

248 """ 

249 # Acceptable horizontal coordinate names. 

250 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"] 

251 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"] 

252 

253 # Get a list of coordinate names for the cube 

254 coord_names = [coord.name() for coord in cube.dim_coords] 

255 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES] 

256 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES] 

257 

258 # If there is one coordinate for both x and y direction return True. 

259 return len(x_coords) == 1 and len(y_coords) == 1 

260 

261 

262def is_coorddim(cube: iris.cube.Cube, coord_name) -> bool: 

263 """Determine whether a cube has specified dimension coordinates. 

264 

265 Arguments 

266 --------- 

267 cube: iris.cube.Cube 

268 An iris cube which will be checked to see if it contains coordinate 

269 names that match a pre-defined list of acceptable coordinate names. 

270 

271 coord_name: str 

272 A cube dimension name 

273 

274 Returns 

275 ------- 

276 bool 

277 If true, then the cube has a spatial projection and thus can be plotted 

278 as a map. 

279 """ 

280 # Get a list of dimension coordinate names for the cube 

281 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)] 

282 

283 # Check if requested dimension is found in cube and get index 

284 return coord_name in coord_names 

285 

286 

287def is_transect(cube: iris.cube.Cube) -> bool: 

288 """Determine whether a cube is a transect. 

289 

290 If cube is a transect, it will contain only one spatial (map) coordinate, 

291 and one vertical coordinate (either pressure or model level). 

292 

293 Arguments 

294 --------- 

295 cube: iris.cube.Cube 

296 An iris cube which will be checked to see if it contains coordinate 

297 names that match a pre-defined list of acceptable coordinate names. 

298 

299 Returns 

300 ------- 

301 bool 

302 If true, then the cube is a transect that contains one spatial (map) 

303 coordinate and one vertical coordinate. 

304 """ 

305 # Acceptable spatial (map) coordinate names. 

306 SPATIAL_MAP_COORD_NAMES = [ 

307 "longitude", 

308 "grid_longitude", 

309 "projection_x_coordinate", 

310 "x", 

311 "latitude", 

312 "grid_latitude", 

313 "projection_y_coordinate", 

314 "y", 

315 "distance", 

316 ] 

317 

318 # Acceptable vertical coordinate names 

319 VERTICAL_COORD_NAMES = ["pressure", "model_level_number", "level_height"] 

320 

321 # Get a list of coordinate names for the cube 

322 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)] 

323 

324 # Check which spatial coordinates we have. 

325 spatial_coords = [ 

326 coord for coord in coord_names if coord in SPATIAL_MAP_COORD_NAMES 

327 ] 

328 if len(spatial_coords) != 1: 

329 return False 

330 

331 # Check which vertical coordinates we have. 

332 vertical_coords = [coord for coord in coord_names if coord in VERTICAL_COORD_NAMES] 

333 if len(vertical_coords) != 1: # noqa: SIM103 Clearer to keep separate. 

334 return False 

335 

336 # Passed criteria so return True 

337 return True 

338 

339 

340def check_stamp_coordinate(cube: iris.cube.Cube) -> str: 

341 """ 

342 Return stamp dimension coordinate name from a given cube, if exists. 

343 

344 If cube contains a valid stamp coordinate as a dimension coordinate, 

345 function will return name of the stamp coordinate. 

346 

347 Arguments 

348 --------- 

349 cube: iris.cube.Cube 

350 An iris cube which will be checked to see if it contains coordinate 

351 names that match a pre-defined list of acceptable coordinate names. 

352 

353 Returns 

354 ------- 

355 str 

356 If available, then return name of stamp coordinate. 

357 Defaults to "realization" if alternative stamp coordinate not found. 

358 """ 

359 # Acceptable stamp coordinate names 

360 STAMP_COORD_NAMES = ["realization", "member", "sample", "pseudo_level"] 

361 

362 # Check which dimension coordinates we have. 

363 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)] 

364 

365 # Check if any acceptable stamp coordinates are cube dimensions. 

366 stamp_coords = [coord for coord in dim_coord_names if coord in STAMP_COORD_NAMES] 

367 if len(stamp_coords) == 1: 

368 stamp_coordinate = stamp_coords[0] 

369 else: 

370 stamp_coordinate = "realization" 

371 

372 return stamp_coordinate 

373 

374 

375def fully_equalise_attributes(cubes: iris.cube.CubeList): 

376 """Remove any unique attributes between cubes or coordinates in place.""" 

377 # Equalise cube attributes. 

378 removed = iris.util.equalise_attributes(cubes) 

379 logger.debug("Removed attributes from cube: %s", removed) 

380 

381 # Equalise coordinate attributes. 

382 coord_sets = [{coord.name() for coord in cube.coords()} for cube in cubes] 

383 

384 all_coords = set.union(*coord_sets) 

385 coords_to_equalise = set.intersection(*coord_sets) 

386 coords_to_remove = set.difference(all_coords, coords_to_equalise) 

387 

388 logger.debug("All coordinates: %s", all_coords) 

389 logger.debug("Coordinates to remove: %s", coords_to_remove) 

390 logger.debug("Coordinates to equalise: %s", coords_to_equalise) 

391 

392 for coord in coords_to_remove: 

393 for cube in cubes: 

394 try: 

395 cube.remove_coord(coord) 

396 logger.debug("Removed coordinate %s from %s cube.", coord, cube.name()) 

397 except iris.exceptions.CoordinateNotFoundError: 

398 pass 

399 

400 for coord in coords_to_equalise: 

401 removed = iris.util.equalise_attributes([cube.coord(coord) for cube in cubes]) 

402 logger.debug("Removed attributes from coordinate %s: %s", coord, removed) 

403 

404 return cubes 

405 

406 

407def slice_over_maybe(cube: iris.cube.Cube, coord_name, index): 

408 """Test slicing over cube if exists. 

409 

410 Return None if not existing. 

411 

412 Arguments 

413 --------- 

414 cube: iris.cube.Cube 

415 An iris cube which will be checked to see if it can be sliced over 

416 given coordinate. 

417 coord_name: coord 

418 An iris coordinate over which to slice cube. 

419 index: 

420 Coordinate index value to extract 

421 

422 Returns 

423 ------- 

424 cube_slice: iris.cube.Cube 

425 A slice of iris cube, if available to slice. 

426 """ 

427 if cube is None: 

428 return None 

429 

430 # Check if coord exists as dimension coordinate 

431 if not is_coorddim(cube, coord_name): 

432 return cube 

433 

434 # Use iris to find which axis the dimension coordinate corresponds to 

435 dim = cube.coord_dims(coord_name)[0] 

436 

437 # Create list of slices for each dimension 

438 slices = [slice(None)] * cube.ndim 

439 

440 # Only replace the slice for the dim to be extracted 

441 slices[dim] = index 

442 

443 return cube[tuple(slices)] 

444 

445 

446def is_time_aux_coord(cube: iris.cube.Cube) -> bool: 

447 """Determine whether a cube has time coordinates as auxiliary coordinates. 

448 

449 Checks if 'forecast_period' and 'forecast_reference_time' exist as 

450 auxiliary coordinates rather than dimension coordinates. 

451 

452 Arguments 

453 --------- 

454 cube: iris.cube.Cube 

455 An iris cube which will be checked to see if it contains 

456 'forecast_period' and 'forecast_reference_time' as auxiliary 

457 coordinates. 

458 

459 Returns 

460 ------- 

461 bool 

462 If True, then the cube has both 'forecast_period' and 

463 'forecast_reference_time' as auxiliary coordinates (not dimension 

464 coordinates). 

465 """ 

466 # Acceptable time coordinate names 

467 TEMPORAL_COORD_NAMES = ["forecast_period", "forecast_reference_time"] 

468 

469 # Get dimension coordinate names 

470 dim_coord_names = [coord.name() for coord in cube.dim_coords] 

471 

472 # Check which temporal coordinates are being used as dimension coordinates 

473 temporal_dim_coords = [ 

474 coord for coord in dim_coord_names if coord in TEMPORAL_COORD_NAMES 

475 ] 

476 

477 # If both coordinates are dimension coordinates, return False 

478 if len(temporal_dim_coords) == 2: 

479 return False 

480 

481 # Check if both temporal coordinates exist in the cube 

482 has_forecast_period = False 

483 has_forecast_reference_time = False 

484 

485 try: 

486 cube.coord("forecast_period") 

487 has_forecast_period = True 

488 except iris.exceptions.CoordinateNotFoundError: 

489 pass 

490 

491 try: 

492 cube.coord("forecast_reference_time") 

493 has_forecast_reference_time = True 

494 except iris.exceptions.CoordinateNotFoundError: 

495 pass 

496 

497 # Check that both exist and are not used as dimension coordinates 

498 if has_forecast_period and has_forecast_reference_time: 

499 return len(temporal_dim_coords) == 0 

500 

501 return False 

502 

503 

504def is_time_aggregatable(cube: iris.cube.Cube) -> bool: 

505 """Determine whether a cube can be aggregated in time. 

506 

507 If a cube is aggregatable it will contain both a 'forecast_reference_time' 

508 and 'forecast_period' coordinate as dimension or scalar coordinates. 

509 

510 Arguments 

511 --------- 

512 cube: iris.cube.Cube 

513 An iris cube which will be checked to see if it is aggregatable based 

514 on a set of pre-defined dimensional time coordinates: 

515 'forecast_period' and 'forecast_reference_time'. 

516 

517 Returns 

518 ------- 

519 bool 

520 If true, then the cube is aggregatable and contains dimensional 

521 coordinates including both 'forecast_reference_time' and 

522 'forecast_period'. 

523 """ 

524 # Acceptable time coordinate names for aggregatable cube. 

525 TEMPORAL_COORD_NAMES = ["forecast_period", "forecast_reference_time"] 

526 

527 def strictly_monotonic(coord: iris.coords.Coord) -> bool: 

528 """Return whether a coord is strictly monotonic, catching errors.""" 

529 try: 

530 return coord.is_monotonic() 

531 except iris.exceptions.CoordinateMultiDimError: 

532 return False 

533 

534 # Strictly monotonic coordinate names for the cube. 

535 coord_names = [coord.name() for coord in cube.coords() if strictly_monotonic(coord)] 

536 

537 # Check which temporal coordinates we have. 

538 temporal_coords = [coord for coord in coord_names if coord in TEMPORAL_COORD_NAMES] 

539 # Return whether both coordinates are in the temporal coordinates. 

540 return len(temporal_coords) == 2 

541 

542 

543def guess_bounds(cube): 

544 """ 

545 Guess bounds for x and y coordinates on a cube. 

546 

547 Arguments 

548 --------- 

549 cube: iris.cube.Cube 

550 Input cube whose x and y coordinate bounds will be guessed if missing. 

551 

552 Returns 

553 ------- 

554 iris.cube.Cube 

555 The same cube with bounds added to x and y coordinates where absent. 

556 

557 Raises 

558 ------ 

559 ValueError 

560 If the cube uses a variable resolution grid where bounds cannot be 

561 guessed reliably. 

562 """ 

563 # Loop over spatial coordinates 

564 for axis in ["x", "y"]: 

565 coord = cube.coord(axis=axis) 

566 try: 

567 _ = iris.util.regular_step(coord) 

568 except ValueError as e: 

569 logger.warning( 

570 "Cannot guess bounds for a variable resolution (non-regular) grid: %s", 

571 e, 

572 ) 

573 # Guess bounds if there aren't any 

574 if coord.bounds is None: 

575 coord.guess_bounds() 

576 return cube 

577 

578 

579def identify_unique_times(cubes, time_coord_name): 

580 """Identify unique time points across a Cube or CubeList. 

581 

582 Arguments 

583 --------- 

584 cubes: iris.cube.Cube | iris.cube.CubeList 

585 A single cube or CubeList to extract unique times from. 

586 time_coord_name: str 

587 Name of the time coordinate to extract (e.g., "time", "forecast_period"). 

588 

589 Returns 

590 ------- 

591 time_coord: iris.coords.DimCoord 

592 A dimension coordinate containing all unique time points sorted in order. 

593 """ 

594 # Handle single cube input 

595 if isinstance(cubes, iris.cube.Cube): 

596 cubes = iris.cube.CubeList([cubes]) 

597 

598 times = [] 

599 time_unit = None 

600 # Loop over cubes 

601 for cube in cubes: 

602 # Extract the desired time coordinate from the cube 

603 time_coord = cube.coord(time_coord_name) 

604 

605 # Get the units for the specified time coordinate 

606 if time_unit is None: 

607 time_unit = time_coord.units 

608 

609 # Store the time coordinate points 

610 times.extend(time_coord.points) 

611 

612 # Construct a list of unique times and store them in a new time coordinate 

613 times = sorted(set(times)) 

614 time_coord = iris.coords.DimCoord(times, units=time_unit) 

615 time_coord.rename(time_coord_name) 

616 

617 return time_coord 

618 

619 

620def remove_cell_method(cube, cell_method): 

621 cube.cell_methods = tuple(cm for cm in cube.cell_methods if cm != cell_method) 

622 return cube 

623 

624 

625def remove_duplicates(cubelist): 

626 # Nothing to do if the cubelist is empty 

627 if not cubelist: 627 ↛ 628line 627 didn't jump to line 628 because the condition on line 627 was never true

628 return cubelist 

629 # Build up a list of indices of the cubes to remove because they are 

630 # duplicated 

631 indices_to_remove = [] 

632 for i in range(len(cubelist) - 1): 

633 cube_i = cubelist[i] 

634 for j in range(i + 1, len(cubelist)): 

635 cube_j = cubelist[j] 

636 if cube_i == cube_j and j not in indices_to_remove: 

637 indices_to_remove.append(j) 

638 # Only keep unique cubes 

639 cubelist = iris.cube.CubeList( 

640 [cube for index, cube in enumerate(cubelist) if index not in indices_to_remove] 

641 ) 

642 return cubelist 

643 

644 

645def check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube: 

646 """Ensure a single cube is given. 

647 

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

649 otherwise an error is raised. 

650 

651 Parameters 

652 ---------- 

653 cube: Cube | CubeList 

654 The cube to check. 

655 

656 Returns 

657 ------- 

658 cube: Cube 

659 The checked cube. 

660 

661 Raises 

662 ------ 

663 TypeError 

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

665 """ 

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

667 return cube 

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

669 if len(cube) == 1: 

670 return cube[0] 

671 else: 

672 raise ValueError("CubeList did not contain a single cube.", cube) 

673 raise TypeError( 

674 "check_single_cube requires a Cube or CubeList of a single cube.", cube 

675 ) 

676 

677 

678def check_sequence_coordinate(cubes, sequence_coordinate): 

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

680 # time slider option. 

681 for cube in iter_maybe(cubes): 

682 try: 

683 cube.coord(sequence_coordinate) 

684 except iris.exceptions.CoordinateNotFoundError as err: 

685 raise ValueError( 

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

687 ) from err 

688 

689 

690def get_num_models(cube: iris.cube.Cube | iris.cube.CubeList) -> int: 

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

692 model_names = {cb.attributes.get("model_name") for cb in iter_maybe(cube)} 

693 

694 if not model_names: 694 ↛ 695line 694 didn't jump to line 695 because the condition on line 694 was never true

695 logger.debug("Missing model names. Will assume single model.") 

696 return 1 

697 else: 

698 return len(model_names) 

699 

700 

701def validate_cube_shape( 

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

703) -> None: 

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

705 if isinstance(cube, iris.cube.CubeList) and len(cube) != num_models: 

706 raise ValueError( 

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

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

709 ) 

710 

711 

712def validate_cubes_coords( 

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

714) -> None: 

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

716 if len(cubes) != len(coords): 

717 raise ValueError( 

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

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

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

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

722 ) 

723 

724 

725def check_if_cylc_workflow() -> Path | None: 

726 """Determine if we are running in a Cylc workflow. 

727 

728 If running in a Cylc workflow, the ROSE_DATAC environment variable 

729 will be set. 

730 

731 Returns 

732 ------- 

733 Path | None: 

734 If ROSE_DATAC is set, and the path exists, return a Path object 

735 containing the path. Otherwise, return None. 

736 """ 

737 # Standard location of ROSE_DATAC data dir in CSET. 

738 try: 

739 dataloc = Path(os.environ["ROSE_DATAC"]) 

740 if dataloc.exists(): 

741 return dataloc 

742 except KeyError: 

743 pass 

744 

745 # If ROSE_DATAC unset or its path does not exist, return None 

746 return None 

747 

748 

749def calc_array_stats( 

750 array: np.ndarray | np.ma.MaskedArray, 

751) -> tuple[float, float, float]: 

752 """Calculate the min, max, and mean of an array. 

753 

754 NaNs/Masked data is ignored. 

755 

756 Returns 

757 ------- 

758 stats: 

759 A tuple of (min, max, mean). 

760 """ 

761 if np.ma.isMaskedArray(array): 

762 array_min = array.min() 

763 array_max = array.max() 

764 array_mean = array.mean() 

765 else: 

766 array_min = np.nanmin(array) 

767 array_max = np.nanmax(array) 

768 array_mean = np.nanmean(array) 

769 return array_min, array_max, array_mean