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

31logger = logging.getLogger(__name__)

34def noop(x, **kwargs):

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

37 Useful for constructing diagnostic chains.

39 Arguments

40 ---------

41 x: Any

42 Input to return.

44 Returns

45 -------

46 x: Any

47 The input that was given.

48 """

49 return x

52def remove_attribute(

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

54) -> CubeList:

55 """Remove a cube attribute.

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

59 Arguments

60 ---------

61 cubes: Cube | CubeList

62 One or more cubes to remove the attribute from.

63 attribute: str | Iterable

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

66 Returns

67 -------

68 cubes: CubeList

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

70 """

71 # Ensure cubes is a CubeList.

72 if not isinstance(cubes, CubeList):

73 cubes = CubeList(iter_maybe(cubes))

75 for cube in cubes:

76 for attr in iter_maybe(attribute):

77 cube.attributes.pop(attr, None)

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

80 # attributes.

81 cubes = cubes.merge()

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

83 cubes = cubes.concatenate()

84 return cubes

87def remove_scalar_coords(

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

89) -> CubeList:

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

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

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

94 include ``realization`` and ``forecast_reference_time`` on model

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

96 left unchanged.

98 Arguments

99 ---------

100 cubes: Cube | CubeList

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

102 coords: str | Iterable

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

104

105 Returns

106 -------

107 cubes: CubeList

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

109 removed where present.

110 """

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

112 cubes = CubeList(iter_maybe(cubes))

113

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

115 coords = [coords]

116

117 for cube in cubes:

118 for coord_name in iter_maybe(coords):

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

120 coord = cube.coord(coord_name)

121 # only remove if scalar

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

123 cube.remove_coord(coord)

124

125 return cubes

126

127

128def addition(addend_1, addend_2):

129 """Addition of two fields.

130

131 Parameters

132 ----------

133 addend_1: Cube

134 Any field to have another field added to it.

135 addend_2: Cube

136 Any field to be added to another field.

137

138 Returns

139 -------

140 Cube

141

142 Raises

143 ------

144 ValueError, iris.exceptions.NotYetImplementedError

145 When the cubes are not compatible.

146

147 Notes

148 -----

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

150 creating new diagnostics by using recipes.

151

152 Examples

153 --------

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

155

156 """

157 return addend_1 + addend_2

158

159

160def subtraction(

161 minuend: Cube | CubeList, subtrahend: Cube | CubeList

162) -> Cube | CubeList:

163 """Subtraction of two fields.

164

165 Parameters

166 ----------

167 minuend: Cube | CubeList

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

169 subtrahend: Cube | CubeList

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

171

172 Returns

173 -------

174 Cube | CubeList

175

176 Raises

177 ------

178 ValueError, iris.exceptions.NotYetImplementedError

179 When the cubes are not compatible.

180

181 Notes

182 -----

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

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

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

186 models or model configurations.

187

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

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

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

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

192

193 Examples

194 --------

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

196

197 """

198

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

200 result = cube_a - cube_b

201 result.attributes.update(cube_a.attributes)

202 return result

203

204 # Case where both inputs are single cubes

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

206 return subtract_preserve_attributes(minuend, subtrahend)

207

208 # Check if minuend is iterable

209 cubes_a = iter_maybe(minuend)

210

211 # Case: subtrahend also iterable

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

213 cubes_b = iter_maybe(subtrahend)

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

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

216 result = iris.cube.CubeList(

217 [

218 subtract_preserve_attributes(cube_a, cube_b)

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

220 ]

221 )

222 else:

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

224 result = iris.cube.CubeList(

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

226 )

227 else:

228 # Case: subtract single cube from each minuend

229 result = iris.cube.CubeList(

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

231 )

232

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

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

235

236

237def division(numerator, denominator):

238 """Division of two fields.

239

240 Parameters

241 ----------

242 numerator: Cube

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

244 denominator: Cube

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

246 value in a ratio.

247

248 Returns

249 -------

250 Cube

251

252 Raises

253 ------

254 ValueError

255 When the cubes are not compatible.

256

257 Notes

258 -----

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

260 creating new diagnostics by using recipes.

261

262 Examples

263 --------

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

265

266 """

267 return numerator / denominator

268

269

270def multiplication(

271 multiplicand: Cube | CubeList, multiplier: Cube | CubeList

272) -> Cube | CubeList:

273 """Multiplication of two fields.

274

275 Parameters

276 ----------

277 multiplicand: Cube | CubeList

278 Any field to be multiplied by another field.

279 multiplier: Cube | CubeList

280 Any field to be multiplied to another field.

281

282 Returns

283 -------

284 Cube | CubeList

285 The result of multiplicand x multiplier.

286

287 Raises

288 ------

289 ValueError

290 When the cubes are not compatible.

291

292 Notes

293 -----

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

295 creating new diagnostics by using recipes. CubeLists are multiplied

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

297

298 Examples

299 --------

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

301

302 """

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

304 for cube_a, cube_b in zip(

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

306 ):

307 multiplied_cube = cube_a * cube_b

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

309 new_cubelist.append(multiplied_cube)

310 if len(new_cubelist) == 1:

311 return new_cubelist[0]

312 else:

313 return new_cubelist

314

315

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

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

318

319 Arguments

320 ---------

321 first: Cube | CubeList

322 First cube or CubeList to merge into CubeList.

323 second: Cube | CubeList

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

325 argument.

326 third: Cube | CubeList

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

328 names.

329 ...

330

331 Returns

332 -------

333 combined_cubelist: CubeList

334 Combined CubeList containing all cubes/CubeLists.

335

336 Raises

337 ------

338 TypeError:

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

340 """

341 # Create empty CubeList to store cubes/CubeList.

342 all_cubes = CubeList()

343 # Combine all CubeLists into a single flat iterable.

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

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

346 if isinstance(item, Cube):

347 # Add cube to CubeList.

348 all_cubes.append(item)

349 else:

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

351 return all_cubes

352

353

354def difference(cubes: CubeList):

355 """Difference of two fields.

356

357 Parameters

358 ----------

359 cubes: CubeList

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

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

362

363 Returns

364 -------

365 Cube

366

367 Raises

368 ------

369 ValueError

370 When the cubes are not compatible.

371

372 Notes

373 -----

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

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

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

377 models or model configurations.

378

379 Examples

380 --------

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

382

383 """

384 if len(cubes) != 2:

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

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

387 other: Cube = cubes.extract_cube(

388 iris.Constraint(

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

390 )

391 )

392

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

394 for cube in cubes:

395 try:

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

397 cube.coord("pressure").points

398 ):

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

400

401 except iris.exceptions.CoordinateNotFoundError:

402 pass

403

404 # Get spatial coord names.

405 base_lat_name, base_lon_name = get_cube_yxcoordname(base)

406 other_lat_name, other_lon_name = get_cube_yxcoordname(other)

407

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

409 # This is triggered if either

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

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

412 # errors.

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

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

415 # for UM and LFRic comparisons.

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

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

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

419 # given this dependency on regridding.

420 if (

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

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

423 ) or (

424 base.long_name

425 in [

426 "eastward_wind_at_10m",

427 "northward_wind_at_10m",

428 "northward_wind_at_cell_centres",

429 "eastward_wind_at_cell_centres",

430 "zonal_wind_at_pressure_levels",

431 "meridional_wind_at_pressure_levels",

432 "potential_vorticity_at_pressure_levels",

433 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",

434 ]

435 ):

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

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

438

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

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

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

442 if base_lat_direction != other_lat_direction:

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

444

445 # Extract just common time points.

446 base, other = _extract_common_time_points(base, other)

447

448 # Equalise attributes so we can merge.

449 fully_equalise_attributes([base, other])

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

451

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

453 difference = base.copy()

454

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

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

457 # its presents.

458 difference.standard_name = None

459 difference.long_name = (

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

461 ) + "_difference"

462 if base.var_name:

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

464 elif base.standard_name:

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

466

467 difference.data = other.data - base.data

468 return difference

469

470

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

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

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

474 time_coord = next(

475 (

476 coord.name()

477 for coord in filter(

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

479 base.coords(),

480 )

481 ),

482 None,

483 )

484 if not time_coord:

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

486 return (base, other)

487 base_time_coord = base.coord(time_coord)

488 other_time_coord = other.coord(time_coord)

489 logger.debug("Base: %s\nOther: %s", base_time_coord, other_time_coord)

490 if time_coord == "hour":

491 # We directly compare points when comparing coordinates with

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

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

494 # comparison for certain coordinate names.

495 base_times = base_time_coord.points

496 other_times = other_time_coord.points

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

498 else:

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

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

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

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

503 logger.debug("Shared times: %s", shared_times)

504 time_constraint = iris.Constraint(

505 coord_values={

506 time_coord: lambda cell, shared_times=shared_times: (

507 cell.point in shared_times

508 )

509 }

510 )

511 # Extract points matching the shared times.

512 base = base.extract(time_constraint)

513 other = other.extract(time_constraint)

514 if base is None or other is None:

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

516 return (base, other)

517

518

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

520 """Convert the units of a cube.

521

522 Arguments

523 ---------

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

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

526

527 units: str

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

529 CF compliant units.

530

531 Returns

532 -------

533 iris.cube.Cube | iris.cube.CubeList

534 The field converted into the specified units.

535

536 Examples

537 --------

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

539

540 """

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

542 for cube in iter_maybe(cubes):

543 # Copy cube to keep original data.

544 cube_a = cube.copy()

545 # Convert cube units.

546 cube_a.convert_units(units)

547 new_cubelist.append(cube_a)

548 if len(new_cubelist) == 1:

549 return new_cubelist[0]

550 else:

551 return new_cubelist

552

553

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

555 """Rename a cube.

556

557 Arguments

558 ---------

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

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

561

562 name: str

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

564

565 Returns

566 -------

567 iris.cube.Cube | iris.cube.CubeList

568 The renamed field.

569

570 Notes

571 -----

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

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

574 long_name. For example, if combining masks

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

576 rather than "mask_for_microphysical_precip_gt_0.0_x_mask_for_microphysical_precip_lt_2.0".

577

578 Examples

579 --------

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

581 """

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

583 for cube in iter_maybe(cubes):

584 cube.rename(name)

585 new_cubelist.append(cube)

586 if len(new_cubelist) == 1:

587 return new_cubelist[0]

588 else:

589 return new_cubelist

590

591

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

593 """

594 Extract levels from a cube for a given coordinate.

595

596 Arguments

597 ---------

598 cube: iris.cube.Cube

599 A Cube to be sliced.

600

601 coord_name: str

602 The coordinate name to be sliced

603

604 levels: list

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

606

607 Returns

608 -------

609 iris.cube.Cube

610 The sliced cube.

611 """

612 coord = cube.coord(coord_name)

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

614

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

616

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

618 slicer[dim_index] = mask

619

620 return cube[tuple(slicer)]

621

622

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

624 """

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

626

627 Parameters

628 ----------

629 cubes: iris.cube.CubeList

630 CubeList containing cubes for which to extract common points.

631

632 coordinate: str

633 The coordinate name to be checked for common points.

634

635 Returns

636 -------

637 iris.cube.CubeList

638 CubeList containing the two cubes sliced to common points

639 for the given coordinate.

640 """

641 # Check type of input

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

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

644

645 # Extract coordinate

646 try:

647 points_list = []

648 for cube in cubes:

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

650 except iris.exceptions.CoordinateNotFoundError as err:

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

652

653 # Find common points

654 common_points = reduce(np.intersect1d, points_list)

655

656 # Check that common points is more than zero.

657 if common_points.size == 0:

658 raise ValueError("No common levels found")

659

660 # Extract common points

661 common_cubes = iris.cube.CubeList()

662 for cube in cubes:

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

664

665 return common_cubes

666

667

668def differentiate(

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

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

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

672

673 Arguments

674 ---------

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

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

677

678 coordinate: str

679 The coordinate that is to be differentiated over.

680

681 Returns

682 -------

683 iris.cube.Cube | iris.cube.CubeList

684 The differential of the cube along the specified coordinate.

685

686 Notes

687 -----

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

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

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

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

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

693

694 Examples

695 --------

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

697 """

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

699 for cube in iter_maybe(cubes):

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

701 new_cubelist.append(dcube)

702 if len(new_cubelist) == 1:

703 return new_cubelist[0]

704 else:

705 return new_cubelist