Coverage for src/CSET/operators/misc.py: 95%

1# © Crown copyright, Met Office (2022-2025) and CSET contributors.

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

7# http://www.apache.org/licenses/LICENSE-2.0

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.

15"""Miscellaneous operators."""

17import itertools

18import logging

19from collections.abc import Iterable

20from functools import reduce

22import iris

23import iris.analysis.calculus

24import numpy as np

25from iris.cube import Cube, CubeList

27from CSET._common import is_increasing, iter_maybe

28from CSET.operators._utils import (

29 ensure_common_time_points,

30 fully_equalise_attributes,

31 get_cube_yxcoordname,

32 get_shared_times,

33 get_time_coord,

34)

35from CSET.operators.regrid import regrid_onto_cube

37logger = logging.getLogger(__name__)

40def noop(x, **kwargs):

41 """Return its input without doing anything to it.

43 Useful for constructing diagnostic chains.

45 Arguments

46 ---------

47 x: Any

48 Input to return.

50 Returns

51 -------

52 x: Any

53 The input that was given.

54 """

55 return x

58def remove_attribute(

59 cubes: Cube | CubeList, attribute: str | Iterable, **kwargs

60) -> CubeList:

61 """Remove a cube attribute.

63 If the attribute is not on the cube, the cube is passed through unchanged.

65 Arguments

66 ---------

67 cubes: Cube | CubeList

68 One or more cubes to remove the attribute from.

69 attribute: str | Iterable

70 Name of attribute (or Iterable of names) to remove.

72 Returns

73 -------

74 cubes: CubeList

75 CubeList of cube(s) with the attribute removed.

76 """

77 # Ensure cubes is a CubeList.

78 if not isinstance(cubes, CubeList):

79 cubes = CubeList(iter_maybe(cubes))

81 for cube in cubes:

82 for attr in iter_maybe(attribute):

83 cube.attributes.pop(attr, None)

85 # Combine things that can be merged due to remove removing the

86 # attributes.

87 cubes = cubes.merge()

88 # combine items that can be merged after removing unwanted attributes

89 cubes = cubes.concatenate()

90 return cubes

93def remove_scalar_coords(

94 cubes: Cube | CubeList, coords: str | Iterable[str]

95) -> CubeList:

96 """Remove scalar coordinates from one or more cubes.

98 Coordinates are only removed if they exist on the cube and are

99 scalar coordinates (i.e. have no associated dimensions). Examples

100 include ``realization`` and ``forecast_reference_time`` on model

101 data cubes. Dimensional and non-scalar auxiliary coordinates are

102 left unchanged.

103

104 Arguments

105 ---------

106 cubes: Cube | CubeList

107 One or more cubes from which scalar coordinates will be removed.

108 coords: str | Iterable

109 Name of a coordinate (or Iterable of coordinate names) to remove.

110

111 Returns

112 -------

113 cubes: CubeList

114 CubeList of cube(s) with the requested scalar coordinates

115 removed where present.

116 """

117 if not isinstance(cubes, CubeList): 117 ↛ 120line 117 didn't jump to line 120 because the condition on line 117 was always true

118 cubes = CubeList(iter_maybe(cubes))

119

120 if isinstance(coords, str): 120 ↛ 121line 120 didn't jump to line 121 because the condition on line 120 was never true

121 coords = [coords]

122

123 for cube in cubes:

124 for coord_name in iter_maybe(coords):

125 if cube.coords(coord_name): 125 ↛ 124line 125 didn't jump to line 124 because the condition on line 125 was always true

126 coord = cube.coord(coord_name)

127 # only remove if scalar

128 if cube.coord_dims(coord) == ():

129 cube.remove_coord(coord)

130

131 return cubes

132

133

134def addition(addend_1, addend_2):

135 """Addition of two fields.

136

137 Parameters

138 ----------

139 addend_1: Cube

140 Any field to have another field added to it.

141 addend_2: Cube

142 Any field to be added to another field.

143

144 Returns

145 -------

146 Cube

147

148 Raises

149 ------

150 ValueError, iris.exceptions.NotYetImplementedError

151 When the cubes are not compatible.

152

153 Notes

154 -----

155 This is a simple operator designed for combination of diagnostics or

156 creating new diagnostics by using recipes.

157

158 Examples

159 --------

160 >>> field_addition = misc.addition(kinetic_energy_u, kinetic_energy_v)

161

162 """

163 return addend_1 + addend_2

164

165

166def subtraction(

167 minuend: Cube | CubeList, subtrahend: Cube | CubeList

168) -> Cube | CubeList:

169 """Subtraction of two fields.

170

171 Parameters

172 ----------

173 minuend: Cube | CubeList

174 Any field(s) to have another field subtracted from it.

175 subtrahend: Cube | CubeList

176 Any field(s) to be subtracted from to another field.

177

178 Returns

179 -------

180 Cube | CubeList

181

182 Raises

183 ------

184 ValueError, iris.exceptions.NotYetImplementedError

185 When the cubes are not compatible.

186

187 Notes

188 -----

189 This is a simple operator designed for combination of diagnostics or

190 creating new diagnostics by using recipes. It can be used for model

191 differences to allow for comparisons between the same field in different

192 models or model configurations.

193

194 If called with 2 Cubes as input, will return a difference cube.

195 If called with 2 CubeLists as input, with return a CubeList of differences.

196 If called with CubeList as minuend and Cube as subtrahend, will return CubeList of differences subtracting Cube from each element of input CubeList.

197 If called with Cube as minuend and CubeList as subtrahend, will return CubeList of differences subtracting each element on CubeList from input Cube.

198

199 Examples

200 --------

201 >>> model_diff = misc.subtraction(temperature_model_A, temperature_model_B)

202

203 """

204

205 def subtract_preserve_attributes(cube_a: Cube, cube_b: Cube) -> Cube:

206 result = cube_a - cube_b

207 result.attributes.update(cube_a.attributes)

208 return result

209

210 # Case where both inputs are single cubes

211 if isinstance(minuend, iris.cube.Cube) and isinstance(subtrahend, iris.cube.Cube):

212 return subtract_preserve_attributes(minuend, subtrahend)

213

214 # Check if minuend is iterable

215 cubes_a = iter_maybe(minuend)

216

217 # Case: subtrahend also iterable

218 if isinstance(subtrahend, iris.cube.CubeList):

219 cubes_b = iter_maybe(subtrahend)

220 if isinstance(minuend, iris.cube.CubeList):

221 # Case: subtract cubelist from cubelist - assume both same sizes

222 result = iris.cube.CubeList(

223 [

224 subtract_preserve_attributes(cube_a, cube_b)

225 for cube_a, cube_b in zip(cubes_a, cubes_b, strict=True)

226 ]

227 )

228 else:

229 # Case: subtract each element of subtrahend from single cube

230 result = iris.cube.CubeList(

231 [subtract_preserve_attributes(minuend, cube_b) for cube_b in cubes_b]

232 )

233 else:

234 # Case: subtract single cube from each minuend

235 result = iris.cube.CubeList(

236 [subtract_preserve_attributes(cube_a, subtrahend) for cube_a in cubes_a]

237 )

238

239 # Return single cube if only one result, else return CubeList

240 return result[0] if len(result) == 1 else result

241

242

243def division(numerator, denominator):

244 """Division of two fields.

245

246 Parameters

247 ----------

248 numerator: Cube

249 Any field to have the ratio taken with respect to another field.

250 denominator: Cube

251 Any field used to divide another field or provide the reference

252 value in a ratio.

253

254 Returns

255 -------

256 Cube

257

258 Raises

259 ------

260 ValueError

261 When the cubes are not compatible.

262

263 Notes

264 -----

265 This is a simple operator designed for combination of diagnostics or

266 creating new diagnostics by using recipes.

267

268 Examples

269 --------

270 >>> bowen_ratio = misc.division(sensible_heat_flux, latent_heat_flux)

271

272 """

273 return numerator / denominator

274

275

276def multiplication(

277 multiplicand: Cube | CubeList, multiplier: Cube | CubeList

278) -> Cube | CubeList:

279 """Multiplication of two fields.

280

281 Parameters

282 ----------

283 multiplicand: Cube | CubeList

284 Any field to be multiplied by another field.

285 multiplier: Cube | CubeList

286 Any field to be multiplied to another field.

287

288 Returns

289 -------

290 Cube | CubeList

291 The result of multiplicand x multiplier.

292

293 Raises

294 ------

295 ValueError

296 When the cubes are not compatible.

297

298 Notes

299 -----

300 This is a simple operator designed for combination of diagnostics or

301 creating new diagnostics by using recipes. CubeLists are multiplied

302 on a strict ordering (e.g. first cube with first cube).

303

304 Examples

305 --------

306 >>> filtered_CAPE_ratio = misc.multiplication(CAPE_ratio, inflow_layer_properties)

307

308 """

309 new_cubelist = iris.cube.CubeList([])

310 for cube_a, cube_b in zip(

311 iter_maybe(multiplicand), iter_maybe(multiplier), strict=True

312 ):

313 multiplied_cube = cube_a * cube_b

314 multiplied_cube.rename(f"{cube_a.name()}_x_{cube_b.name()}")

315 new_cubelist.append(multiplied_cube)

316 if len(new_cubelist) == 1:

317 return new_cubelist[0]

318 else:

319 return new_cubelist

320

321

322def combine_cubes_into_cubelist(first: Cube | CubeList, **kwargs) -> CubeList:

323 """Operator that combines multiple cubes or CubeLists into one.

324

325 Arguments

326 ---------

327 first: Cube | CubeList

328 First cube or CubeList to merge into CubeList.

329 second: Cube | CubeList

330 Second cube or CubeList to merge into CubeList. This must be a named

331 argument.

332 third: Cube | CubeList

333 There can be any number of additional arguments, they just need unique

334 names.

335 ...

336

337 Returns

338 -------

339 combined_cubelist: CubeList

340 Combined CubeList containing all cubes/CubeLists.

341

342 Raises

343 ------

344 TypeError:

345 If the provided arguments are not either a Cube or CubeList.

346 """

347 # Create empty CubeList to store cubes/CubeList.

348 all_cubes = CubeList()

349 # Combine all CubeLists into a single flat iterable.

350 for item in itertools.chain(iter_maybe(first), *map(iter_maybe, kwargs.values())):

351 # Check each item is a Cube, erroring if not.

352 if isinstance(item, Cube):

353 # Add cube to CubeList.

354 all_cubes.append(item)

355 else:

356 raise TypeError("Not a Cube or CubeList!", item)

357 return all_cubes

358

359

360def difference(cubes: CubeList):

361 """Difference of two fields.

362

363 Parameters

364 ----------

365 cubes: CubeList

366 A list of exactly two cubes. One must have the cset_comparison_base

367 attribute set to 1, and will be used as the base of the comparison.

368

369 Returns

370 -------

371 Cube

372

373 Raises

374 ------

375 ValueError

376 When the cubes are not compatible.

377

378 Notes

379 -----

380 This is a simple operator designed for combination of diagnostics or

381 creating new diagnostics by using recipes. It can be used for model

382 differences to allow for comparisons between the same field in different

383 models or model configurations.

384

385 Examples

386 --------

387 >>> model_diff = misc.difference(temperature_model_A, temperature_model_B)

388

389 """

390 if len(cubes) != 2:

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

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

393 other: Cube = cubes.extract_cube(

394 iris.Constraint(

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

396 )

397 )

398

399 # If cubes contain a pressure coordinate, ensure it is increasing.

400 for cube in cubes:

401 try:

402 if len(cube.coord("pressure").points) > 2 and not is_increasing(

403 cube.coord("pressure").points

404 ):

405 cube.data = np.flip(cube.data, axis=cube.coord_dims("pressure")[0])

406

407 except iris.exceptions.CoordinateNotFoundError:

408 pass

409

410 # Get spatial coord names.

411 base_lat_name, base_lon_name = get_cube_yxcoordname(base)

412 other_lat_name, other_lon_name = get_cube_yxcoordname(other)

413

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

415 # This is triggered if either

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

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

418 # errors.

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

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

421 # for UM and LFRic comparisons.

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

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

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

425 # given this dependency on regridding.

426 if (

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

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

429 ) or (

430 base.long_name

431 in [

432 "eastward_wind_at_10m",

433 "northward_wind_at_10m",

434 "northward_wind_at_cell_centres",

435 "eastward_wind_at_cell_centres",

436 "zonal_wind_at_pressure_levels",

437 "meridional_wind_at_pressure_levels",

438 "potential_vorticity_at_pressure_levels",

439 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",

440 ]

441 ):

442 logger.debug("Linear regridding base cube to other grid to compute differences")

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

444

445 # Figure out if we are comparing between UM and LFRic; flip array if so.

446 base_lat_direction = is_increasing(base.coord(base_lat_name).points)

447 other_lat_direction = is_increasing(other.coord(other_lat_name).points)

448 if base_lat_direction != other_lat_direction:

449 other.data = np.flip(other.data, other.coord(other_lat_name).cube_dims(other))

450

451 # Equalise attributes so we can merge.

452 fully_equalise_attributes([base, other])

453 logger.debug("Base: %s\nOther: %s", base, other)

454

455 # This currently relies on the cubes having the same underlying data layout.

456 difference = base.copy()

457

458 # Differences don't have a standard name; long name gets a suffix. We are

459 # assuming we can rely on cubes having a long name, so we don't check for

460 # its presents.

461 difference.standard_name = None

462 difference.long_name = (

463 base.long_name if base.long_name else base.name()

464 ) + "_difference"

465 if base.var_name:

466 difference.var_name = base.var_name + "_difference"

467 elif base.standard_name:

468 difference.var_name = base.standard_name + "_difference"

469

470 # will raise if not common times

471 ensure_common_time_points(base, other)

472

473 difference.data = other.data - base.data

474 return difference

475

476

477def extract_common_time_points(cubes):

478 """Return cubes with common time points."""

479 coords = cubes[0].coords("forecast_reference_time")

480

481 if bool(coords) and len(coords[0].points) > 1:

482 try:

483 reference_frts = cubes[0].coord("forecast_reference_time").points

484 for cube in cubes[1:]:

485 cube_frts = cube.coord("forecast_reference_time").points

486

487 if not np.array_equal(reference_frts, cube_frts): 487 ↛ 488line 487 didn't jump to line 488 because the condition on line 487 was never true

488 raise ValueError(

489 "Cubes do not share the same forecast_reference_time values."

490 )

491

492 except iris.exceptions.CoordinateNotFoundError as err:

493 raise ValueError(

494 "All cubes must have a forecast_reference_time coordinate."

495 ) from err

496

497 return extract_common_points(cubes, "forecast_period")

498 else:

499 return _extract_common_time_points_one_frt(cubes)

500

501

502def _extract_common_time_points_one_frt(cubes: CubeList) -> CubeList:

503

504 obs: Cube | None

505 models: CubeList

506 obs, models = split_obs_and_models(cubes)

507

508 base = models[0]

509 others = models[1:]

510 list_shared_times: list[set[int]] = []

511

512 time_coord = get_time_coord(cubes)

513

514 if not time_coord:

515 logger.debug("No time coord, skipping equalisation.")

516 return cubes

517

518 for other in iter_maybe(others):

519 list_shared_times.append(get_shared_times(base, other, time_coord))

520

521 shared_times_all = reduce(lambda x, y: x.intersection(y), list_shared_times)

522

523 logger.debug("Shared times: %s", shared_times_all)

524

525 time_constraint = iris.Constraint(

526 coord_values={

527 time_coord: lambda cell, shared_times=shared_times_all: (

528 cell.point in shared_times

529 )

530 }

531 )

532

533 # Extract points matching the shared times.

534 cubes_common_time = CubeList()

535 for cube in cubes:

536 cubes_common_time.append(cube.extract(time_constraint))

537

538 return cubes_common_time

539

540

541def convert_units(cubes: iris.cube.Cube | iris.cube.CubeList, units: str):

542 """Convert the units of a cube.

543

544 Arguments

545 ---------

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

547 A Cube or CubeList of a field for its units to be converted.

548

549 units: str

550 The unit that the original field is to be converted to. It takes

551 CF compliant units.

552

553 Returns

554 -------

555 iris.cube.Cube | iris.cube.CubeList

556 The field converted into the specified units.

557

558 Examples

559 --------

560 >>> T_in_F = misc.convert_units(temperature_in_K, "Fahrenheit")

561

562 """

563 new_cubelist = iris.cube.CubeList([])

564 for cube in iter_maybe(cubes):

565 # Copy cube to keep original data.

566 cube_a = cube.copy()

567 # Convert cube units.

568 cube_a.convert_units(units)

569 new_cubelist.append(cube_a)

570 if len(new_cubelist) == 1:

571 return new_cubelist[0]

572 else:

573 return new_cubelist

574

575

576def rename_cube(cubes: iris.cube.Cube | iris.cube.CubeList, name: str):

577 """Rename a cube.

578

579 Arguments

580 ---------

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

582 A Cube or CubeList of a field to be renamed.

583

584 name: str

585 The new name of the cube. It should be CF compliant.

586

587 Returns

588 -------

589 iris.cube.Cube | iris.cube.CubeList

590 The renamed field.

591

592 Notes

593 -----

594 This operator is designed to be used when the output field name does not

595 match expectations or needs to be different to defaults in standard_name, var_name or

596 long_name. For example, if combining masks

597 to create light rain you would like the field to be named "mask_for_light_rain"

598 rather than "mask_for_microphysical_precip_gt_0.0_x_mask_for_microphysical_precip_lt_2.0".

599

600 Examples

601 --------

602 >>> light_rain_mask = misc.rename_cube(light_rain_mask,"mask_for_light_rainfall"

603 """

604 new_cubelist = iris.cube.CubeList([])

605 for cube in iter_maybe(cubes):

606 cube.rename(name)

607 new_cubelist.append(cube)

608 if len(new_cubelist) == 1:

609 return new_cubelist[0]

610 else:

611 return new_cubelist

612

613

614def _slice_cube_on_levels(cube: iris.cube.Cube, coord_name: str, levels: list):

615 """

616 Extract levels from a cube for a given coordinate.

617

618 Arguments

619 ---------

620 cube: iris.cube.Cube

621 A Cube to be sliced.

622

623 coord_name: str

624 The coordinate name to be sliced

625

626 levels: list

627 A list containing points to be extracted from the cube.

628

629 Returns

630 -------

631 iris.cube.Cube

632 The sliced cube.

633 """

634 coord = cube.coord(coord_name)

635 (dim_index,) = cube.coord_dims(coord)

636

637 mask = np.isin(coord.points, levels)

638

639 slicer = [slice(None)] * cube.ndim

640 slicer[dim_index] = mask

641

642 return cube[tuple(slicer)]

643

644

645def extract_common_points(cubes: iris.cube.CubeList, coordinate: str):

646 """

647 Extract common points for a given coordinate between cubes in a CubeList.

648

649 Parameters

650 ----------

651 cubes: iris.cube.CubeList

652 CubeList containing cubes for which to extract common points.

653

654 coordinate: str

655 The coordinate name to be checked for common points.

656

657 Returns

658 -------

659 iris.cube.CubeList

660 CubeList containing the two cubes sliced to common points

661 for the given coordinate.

662 """

663 # Check type of input

664 if type(cubes) is not iris.cube.CubeList:

665 raise TypeError(f"Not a CubeList, got type {type(cubes)}")

666

667 # Extract coordinate

668 try:

669 points_list = []

670 for cube in cubes:

671 points_list.append(cube.coord(coordinate).points)

672 except iris.exceptions.CoordinateNotFoundError as err:

673 raise ValueError(f"Both cubes must have an {coordinate} coordinate") from err

674

675 # Find common points

676 common_points = reduce(np.intersect1d, points_list)

677

678 # Check that common points is more than zero.

679 if common_points.size == 0:

680 raise ValueError("No common levels found")

681

682 # Extract common points

683 common_cubes = iris.cube.CubeList()

684 for cube in cubes:

685 common_cubes.append(_slice_cube_on_levels(cube, coordinate, common_points))

686

687 return common_cubes

688

689

690def differentiate(

691 cubes: iris.cube.Cube | iris.cube.CubeList, coordinate: str, **kwargs

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

693 """Differentiate a cube on a specified coordinate.

694

695 Arguments

696 ---------

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

698 A Cube or CubeList of a field that is to be differentiated.

699

700 coordinate: str

701 The coordinate that is to be differentiated over.

702

703 Returns

704 -------

705 iris.cube.Cube | iris.cube.CubeList

706 The differential of the cube along the specified coordinate.

707

708 Notes

709 -----

710 The differential is calculated based on a carteisan grid. This calculation

711 is then suitable for vertical and temporal derivatives. It is not sensible

712 for horizontal derivatives if they are based on spherical coordinates (e.g.

713 latitude and longitude). In essence this operator is a CSET wrapper around

714 `iris.analysis.calculus.differentiate <https://scitools-iris.readthedocs.io/en/stable/generated/api/iris.analysis.calculus.html#iris.analysis.calculus.differentiate>`_.

715

716 Examples

717 --------

718 >>> dT_dz = misc.differentiate(temperature, "altitude")

719 """

720 new_cubelist = iris.cube.CubeList([])

721 for cube in iter_maybe(cubes):

722 dcube = iris.analysis.calculus.differentiate(cube, coordinate)

723 new_cubelist.append(dcube)

724 if len(new_cubelist) == 1:

725 return new_cubelist[0]

726 else:

727 return new_cubelist

728

729

730def split_obs_and_models(cubes: CubeList) -> tuple[Cube | None, CubeList]:

731 """Return (obs_cube or None, CubeList of model cubes)."""

732 obs = [cb for cb in cubes if "observed" in (cb.long_name or "")]

733

734 if len(obs) > 1: 734 ↛ 735line 734 didn't jump to line 735 because the condition on line 734 was never true

735 raise ValueError(f"Expected at most one 'observed' cube, found {len(obs)}")

736 if not obs: 736 ↛ 739line 736 didn't jump to line 739 because the condition on line 736 was always true

737 return None, CubeList(cubes)

738

739 base = obs[0]

740 return base, CubeList(cb for cb in cubes if cb is not base)