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

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 fully_equalise_attributes, get_cube_yxcoordname

29from CSET.operators.regrid import regrid_onto_cube

32def noop(x, **kwargs):

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

35 Useful for constructing diagnostic chains.

37 Arguments

38 ---------

39 x: Any

40 Input to return.

42 Returns

43 -------

44 x: Any

45 The input that was given.

46 """

47 return x

50def remove_attribute(

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

52) -> CubeList:

53 """Remove a cube attribute.

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

57 Arguments

58 ---------

59 cubes: Cube | CubeList

60 One or more cubes to remove the attribute from.

61 attribute: str | Iterable

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

64 Returns

65 -------

66 cubes: CubeList

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

68 """

69 # Ensure cubes is a CubeList.

70 if not isinstance(cubes, CubeList):

71 cubes = CubeList(iter_maybe(cubes))

73 for cube in cubes:

74 for attr in iter_maybe(attribute):

75 cube.attributes.pop(attr, None)

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

78 # attributes.

79 cubes = cubes.merge()

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

81 cubes = cubes.concatenate()

82 return cubes

85def remove_scalar_coords(

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

87) -> CubeList:

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

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

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

92 include ``realization`` and ``forecast_reference_time`` on model

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

94 left unchanged.

96 Arguments

97 ---------

98 cubes: Cube | CubeList

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

100 coords: str | Iterable

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

102

103 Returns

104 -------

105 cubes: CubeList

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

107 removed where present.

108 """

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

110 cubes = CubeList(iter_maybe(cubes))

111

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

113 coords = [coords]

114

115 for cube in cubes:

116 for coord_name in iter_maybe(coords):

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

118 coord = cube.coord(coord_name)

119 # only remove if scalar

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

121 cube.remove_coord(coord)

122

123 return cubes

124

125

126def addition(addend_1, addend_2):

127 """Addition of two fields.

128

129 Parameters

130 ----------

131 addend_1: Cube

132 Any field to have another field added to it.

133 addend_2: Cube

134 Any field to be added to another field.

135

136 Returns

137 -------

138 Cube

139

140 Raises

141 ------

142 ValueError, iris.exceptions.NotYetImplementedError

143 When the cubes are not compatible.

144

145 Notes

146 -----

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

148 creating new diagnostics by using recipes.

149

150 Examples

151 --------

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

153

154 """

155 return addend_1 + addend_2

156

157

158def subtraction(

159 minuend: Cube | CubeList, subtrahend: Cube | CubeList

160) -> Cube | CubeList:

161 """Subtraction of two fields.

162

163 Parameters

164 ----------

165 minuend: Cube | CubeList

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

167 subtrahend: Cube | CubeList

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

169

170 Returns

171 -------

172 Cube | CubeList

173

174 Raises

175 ------

176 ValueError, iris.exceptions.NotYetImplementedError

177 When the cubes are not compatible.

178

179 Notes

180 -----

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

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

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

184 models or model configurations.

185

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

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

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

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

190

191 Examples

192 --------

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

194

195 """

196

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

198 result = cube_a - cube_b

199 result.attributes.update(cube_a.attributes)

200 return result

201

202 # Case where both inputs are single cubes

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

204 return subtract_preserve_attributes(minuend, subtrahend)

205

206 # Check if minuend is iterable

207 cubes_a = iter_maybe(minuend)

208

209 # Case: subtrahend also iterable

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

211 cubes_b = iter_maybe(subtrahend)

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

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

214 result = iris.cube.CubeList(

215 [

216 subtract_preserve_attributes(cube_a, cube_b)

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

218 ]

219 )

220 else:

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

222 result = iris.cube.CubeList(

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

224 )

225 else:

226 # Case: subtract single cube from each minuend

227 result = iris.cube.CubeList(

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

229 )

230

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

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

233

234

235def division(numerator, denominator):

236 """Division of two fields.

237

238 Parameters

239 ----------

240 numerator: Cube

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

242 denominator: Cube

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

244 value in a ratio.

245

246 Returns

247 -------

248 Cube

249

250 Raises

251 ------

252 ValueError

253 When the cubes are not compatible.

254

255 Notes

256 -----

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

258 creating new diagnostics by using recipes.

259

260 Examples

261 --------

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

263

264 """

265 return numerator / denominator

266

267

268def multiplication(

269 multiplicand: Cube | CubeList, multiplier: Cube | CubeList

270) -> Cube | CubeList:

271 """Multiplication of two fields.

272

273 Parameters

274 ----------

275 multiplicand: Cube | CubeList

276 Any field to be multiplied by another field.

277 multiplier: Cube | CubeList

278 Any field to be multiplied to another field.

279

280 Returns

281 -------

282 Cube | CubeList

283 The result of multiplicand x multiplier.

284

285 Raises

286 ------

287 ValueError

288 When the cubes are not compatible.

289

290 Notes

291 -----

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

293 creating new diagnostics by using recipes. CubeLists are multiplied

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

295

296 Examples

297 --------

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

299

300 """

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

302 for cube_a, cube_b in zip(

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

304 ):

305 multiplied_cube = cube_a * cube_b

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

307 new_cubelist.append(multiplied_cube)

308 if len(new_cubelist) == 1:

309 return new_cubelist[0]

310 else:

311 return new_cubelist

312

313

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

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

316

317 Arguments

318 ---------

319 first: Cube | CubeList

320 First cube or CubeList to merge into CubeList.

321 second: Cube | CubeList

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

323 argument.

324 third: Cube | CubeList

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

326 names.

327 ...

328

329 Returns

330 -------

331 combined_cubelist: CubeList

332 Combined CubeList containing all cubes/CubeLists.

333

334 Raises

335 ------

336 TypeError:

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

338 """

339 # Create empty CubeList to store cubes/CubeList.

340 all_cubes = CubeList()

341 # Combine all CubeLists into a single flat iterable.

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

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

344 if isinstance(item, Cube):

345 # Add cube to CubeList.

346 all_cubes.append(item)

347 else:

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

349 return all_cubes

350

351

352def difference(cubes: CubeList):

353 """Difference of two fields.

354

355 Parameters

356 ----------

357 cubes: CubeList

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

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

360

361 Returns

362 -------

363 Cube

364

365 Raises

366 ------

367 ValueError

368 When the cubes are not compatible.

369

370 Notes

371 -----

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

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

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

375 models or model configurations.

376

377 Examples

378 --------

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

380

381 """

382 if len(cubes) != 2:

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

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

385 other: Cube = cubes.extract_cube(

386 iris.Constraint(

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

388 )

389 )

390

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

392 for cube in cubes:

393 try:

394 if len(cube.coord("pressure").points) > 2: 394 ↛ 392line 394 didn't jump to line 392 because the condition on line 394 was always true

395 if not is_increasing(cube.coord("pressure").points):

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

397

398 except iris.exceptions.CoordinateNotFoundError:

399 pass

400

401 # Get spatial coord names.

402 base_lat_name, base_lon_name = get_cube_yxcoordname(base)

403 other_lat_name, other_lon_name = get_cube_yxcoordname(other)

404

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

406 # This is triggered if either

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

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

409 # errors.

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

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

412 # for UM and LFRic comparisons.

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

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

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

416 # given this dependency on regridding.

417 if (

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

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

420 ) or (

421 base.long_name

422 in [

423 "eastward_wind_at_10m",

424 "northward_wind_at_10m",

425 "northward_wind_at_cell_centres",

426 "eastward_wind_at_cell_centres",

427 "zonal_wind_at_pressure_levels",

428 "meridional_wind_at_pressure_levels",

429 "potential_vorticity_at_pressure_levels",

430 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",

431 ]

432 ):

433 logging.debug(

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

435 )

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

437

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

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

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

441 if base_lat_direction != other_lat_direction:

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

443

444 # Extract just common time points.

445 base, other = _extract_common_time_points(base, other)

446

447 # Equalise attributes so we can merge.

448 fully_equalise_attributes([base, other])

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

450

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

452 difference = base.copy()

453

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

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

456 # its presents.

457 difference.standard_name = None

458 difference.long_name = (

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

460 ) + "_difference"

461 if base.var_name:

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

463 elif base.standard_name:

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

465

466 difference.data = other.data - base.data

467 return difference

468

469

470def _extract_common_time_points(base: Cube, other: Cube) -> tuple[Cube, Cube]:

471 """Extract common time points from cubes to allow comparison."""

472 # Get the name of the first non-scalar time coordinate.

473 time_coord = next(

474 map(

475 lambda coord: coord.name(),

476 filter(

477 lambda coord: coord.shape > (1,) and coord.name() in ["time", "hour"],

478 base.coords(),

479 ),

480 ),

481 None,

482 )

483 if not time_coord:

484 logging.debug("No time coord, skipping equalisation.")

485 return (base, other)

486 base_time_coord = base.coord(time_coord)

487 other_time_coord = other.coord(time_coord)

488 logging.debug("Base: %s\nOther: %s", base_time_coord, other_time_coord)

489 if time_coord == "hour":

490 # We directly compare points when comparing coordinates with

491 # non-absolute units, such as hour. We can't just check the units are

492 # equal as iris automatically converts to datetime objects in the

493 # comparison for certain coordinate names.

494 base_times = base_time_coord.points

495 other_times = other_time_coord.points

496 shared_times = set.intersection(set(base_times), set(other_times))

497 else:

498 # Units don't match, so converting to datetimes for comparison.

499 base_times = base_time_coord.units.num2date(base_time_coord.points)

500 other_times = other_time_coord.units.num2date(other_time_coord.points)

501 shared_times = set.intersection(set(base_times), set(other_times))

502 logging.debug("Shared times: %s", shared_times)

503 time_constraint = iris.Constraint(

504 coord_values={

505 time_coord: lambda cell, shared_times=shared_times: (

506 cell.point in shared_times

507 )

508 }

509 )

510 # Extract points matching the shared times.

511 base = base.extract(time_constraint)

512 other = other.extract(time_constraint)

513 if base is None or other is None:

514 raise ValueError("No common time points found!")

515 return (base, other)

516

517

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

519 """Convert the units of a cube.

520

521 Arguments

522 ---------

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

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

525

526 units: str

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

528 CF compliant units.

529

530 Returns

531 -------

532 iris.cube.Cube | iris.cube.CubeList

533 The field converted into the specified units.

534

535 Examples

536 --------

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

538

539 """

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

541 for cube in iter_maybe(cubes):

542 # Copy cube to keep original data.

543 cube_a = cube.copy()

544 # Convert cube units.

545 cube_a.convert_units(units)

546 new_cubelist.append(cube_a)

547 if len(new_cubelist) == 1:

548 return new_cubelist[0]

549 else:

550 return new_cubelist

551

552

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

554 """Rename a cube.

555

556 Arguments

557 ---------

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

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

560

561 name: str

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

563

564 Returns

565 -------

566 iris.cube.Cube | iris.cube.CubeList

567 The renamed field.

568

569 Notes

570 -----

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

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

573 long_name. For example, if combining masks

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

575 rather than "mask_for_microphysical_precip_gt_0.0_x_mask_for_microphysical_precip_lt_2.0".

576

577 Examples

578 --------

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

580 """

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

582 for cube in iter_maybe(cubes):

583 cube.rename(name)

584 new_cubelist.append(cube)

585 if len(new_cubelist) == 1:

586 return new_cubelist[0]

587 else:

588 return new_cubelist

589

590

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

592 """

593 Extract levels from a cube for a given coordinate.

594

595 Arguments

596 ---------

597 cube: iris.cube.Cube

598 A Cube to be sliced.

599

600 coord_name: str

601 The coordinate name to be sliced

602

603 levels: list

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

605

606 Returns

607 -------

608 iris.cube.Cube

609 The sliced cube.

610 """

611 coord = cube.coord(coord_name)

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

613

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

615

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

617 slicer[dim_index] = mask

618

619 return cube[tuple(slicer)]

620

621

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

623 """

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

625

626 Parameters

627 ----------

628 cubes: iris.cube.CubeList

629 CubeList containing cubes for which to extract common points.

630

631 coordinate: str

632 The coordinate name to be checked for common points.

633

634 Returns

635 -------

636 iris.cube.CubeList

637 CubeList containing the two cubes sliced to common points

638 for the given coordinate.

639 """

640 # Check type of input

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

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

643

644 # Extract coordinate

645 try:

646 points_list = []

647 for cube in cubes:

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

649 except iris.exceptions.CoordinateNotFoundError as err:

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

651

652 # Find common points

653 common_points = reduce(np.intersect1d, points_list)

654

655 # Check that common points is more than zero.

656 if common_points.size == 0:

657 raise ValueError("No common levels found")

658

659 # Extract common points

660 common_cubes = iris.cube.CubeList()

661 for cube in cubes:

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

663

664 return common_cubes

665

666

667def differentiate(

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

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

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

671

672 Arguments

673 ---------

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

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

676

677 coordinate: str

678 The coordinate that is to be differentiated over.

679

680 Returns

681 -------

682 iris.cube.Cube | iris.cube.CubeList

683 The differential of the cube along the specified coordinate.

684

685 Notes

686 -----

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

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

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

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

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

692

693 Examples

694 --------

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

696 """

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

698 for cube in iter_maybe(cubes):

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

700 new_cubelist.append(dcube)

701 if len(new_cubelist) == 1:

702 return new_cubelist[0]

703 else:

704 return new_cubelist