Coverage for src/CSET/operators/misc.py: 96%
192 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-20 14:01 +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.
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.
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.
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))
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]
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)
131 return cubes
134def addition(addend_1, addend_2):
135 """Addition of two fields.
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.
144 Returns
145 -------
146 Cube
148 Raises
149 ------
150 ValueError, iris.exceptions.NotYetImplementedError
151 When the cubes are not compatible.
153 Notes
154 -----
155 This is a simple operator designed for combination of diagnostics or
156 creating new diagnostics by using recipes.
158 Examples
159 --------
160 >>> field_addition = misc.addition(kinetic_energy_u, kinetic_energy_v)
162 """
163 return addend_1 + addend_2
166def subtraction(
167 minuend: Cube | CubeList, subtrahend: Cube | CubeList
168) -> Cube | CubeList:
169 """Subtraction of two fields.
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.
178 Returns
179 -------
180 Cube | CubeList
182 Raises
183 ------
184 ValueError, iris.exceptions.NotYetImplementedError
185 When the cubes are not compatible.
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.
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.
199 Examples
200 --------
201 >>> model_diff = misc.subtraction(temperature_model_A, temperature_model_B)
203 """
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
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)
214 # Check if minuend is iterable
215 cubes_a = iter_maybe(minuend)
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 )
239 # Return single cube if only one result, else return CubeList
240 return result[0] if len(result) == 1 else result
243def division(numerator, denominator):
244 """Division of two fields.
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.
254 Returns
255 -------
256 Cube
258 Raises
259 ------
260 ValueError
261 When the cubes are not compatible.
263 Notes
264 -----
265 This is a simple operator designed for combination of diagnostics or
266 creating new diagnostics by using recipes.
268 Examples
269 --------
270 >>> bowen_ratio = misc.division(sensible_heat_flux, latent_heat_flux)
272 """
273 return numerator / denominator
276def multiplication(
277 multiplicand: Cube | CubeList, multiplier: Cube | CubeList
278) -> Cube | CubeList:
279 """Multiplication of two fields.
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.
288 Returns
289 -------
290 Cube | CubeList
291 The result of multiplicand x multiplier.
293 Raises
294 ------
295 ValueError
296 When the cubes are not compatible.
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).
304 Examples
305 --------
306 >>> filtered_CAPE_ratio = misc.multiplication(CAPE_ratio, inflow_layer_properties)
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
322def combine_cubes_into_cubelist(first: Cube | CubeList, **kwargs) -> CubeList:
323 """Operator that combines multiple cubes or CubeLists into one.
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 ...
337 Returns
338 -------
339 combined_cubelist: CubeList
340 Combined CubeList containing all cubes/CubeLists.
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
360def difference(cubes: CubeList):
361 """Difference of two fields.
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.
369 Returns
370 -------
371 Cube
373 Raises
374 ------
375 ValueError
376 When the cubes are not compatible.
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.
385 Examples
386 --------
387 >>> model_diff = misc.difference(temperature_model_A, temperature_model_B)
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 )
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])
407 except iris.exceptions.CoordinateNotFoundError:
408 pass
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)
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")
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))
451 # Equalise attributes so we can merge.
452 fully_equalise_attributes([base, other])
453 logger.debug("Base: %s\nOther: %s", base, other)
455 # This currently relies on the cubes having the same underlying data layout.
456 difference = base.copy()
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"
470 # will raise if not common times
471 ensure_common_time_points(base, other)
473 difference.data = other.data - base.data
474 return difference
477def extract_common_time_points(cubes: CubeList) -> CubeList:
478 """Return cubes restricted to their common time points.
480 If the first cube has more than one FRT point, every cube must have identical
481 FRT points. The cubes are then trimmed to their shared
482 forecast_period values, so lead times line up across cubes.
483 Otherwise (no FRT coordinate, or a single FRT), the cubes are
484 trimmed to their shared times.
486 Parameters
487 ----------
488 cubes: CubeList
489 Cubes to align in time. The first cube is used as the reference.
491 Returns
492 -------
493 CubeList
494 Cubes containing only the time points common to all inputs.
496 Raises
497 ------
498 ValueError
499 If the first cube has multiple FRT points and any other cube either
500 lacks a forecast_reference_time coordinate or has different FRT
501 values (compared element-wise, so order matters).
502 """
503 frt_coord = cubes[0].coords("forecast_reference_time")
505 if frt_coord and len(frt_coord[0].points) > 1:
506 for cb in cubes[1:]:
507 reference_frts = cubes[0].coord("forecast_reference_time").points
508 cb_frt_coord = cb.coords("forecast_reference_time")
509 if not cb_frt_coord:
510 raise ValueError(
511 "All cubes must have a forecast_reference_time coordinate."
512 )
513 cube_frts = cb.coord("forecast_reference_time").points
514 if not np.array_equal(reference_frts, cube_frts):
515 raise ValueError(
516 "Cubes do not share the same forecast_reference_time values."
517 )
519 return extract_common_points(cubes, "forecast_period")
520 else:
521 return _extract_common_time_points_one_frt(cubes)
524def _extract_common_time_points_one_frt(cubes: CubeList) -> CubeList:
526 # typehints
527 obs: Cube | None
528 models: CubeList
529 obs, models = split_obs_and_models(cubes)
530 if obs is not None: 530 ↛ 531line 530 didn't jump to line 531 because the condition on line 530 was never true
531 base = obs
532 models[0:]
533 else:
534 base = models[0]
535 others = models[1:]
536 list_shared_times: list[set[int]] = []
538 time_coord = get_time_coord(cubes)
540 if not time_coord:
541 logger.debug("No time coord, skipping equalisation.")
542 return cubes
544 for other in iter_maybe(others):
545 list_shared_times.append(get_shared_times(base, other, time_coord))
547 shared_times_all = reduce(lambda x, y: x.intersection(y), list_shared_times)
549 logger.debug("Shared times: %s", shared_times_all)
551 time_constraint = iris.Constraint(
552 coord_values={
553 time_coord: lambda cell, shared_times=shared_times_all: (
554 cell.point in shared_times
555 )
556 }
557 )
559 # Extract points matching the shared times.
560 cubes_common_time = CubeList()
561 for cube in cubes:
562 cubes_common_time.append(cube.extract(time_constraint))
564 return cubes_common_time
567def convert_units(cubes: iris.cube.Cube | iris.cube.CubeList, units: str):
568 """Convert the units of a cube.
570 Arguments
571 ---------
572 cubes: iris.cube.Cube | iris.cube.CubeList
573 A Cube or CubeList of a field for its units to be converted.
575 units: str
576 The unit that the original field is to be converted to. It takes
577 CF compliant units.
579 Returns
580 -------
581 iris.cube.Cube | iris.cube.CubeList
582 The field converted into the specified units.
584 Examples
585 --------
586 >>> T_in_F = misc.convert_units(temperature_in_K, "Fahrenheit")
588 """
589 new_cubelist = iris.cube.CubeList([])
590 for cube in iter_maybe(cubes):
591 # Copy cube to keep original data.
592 cube_a = cube.copy()
593 # Convert cube units.
594 cube_a.convert_units(units)
595 new_cubelist.append(cube_a)
596 if len(new_cubelist) == 1:
597 return new_cubelist[0]
598 else:
599 return new_cubelist
602def rename_cube(cubes: iris.cube.Cube | iris.cube.CubeList, name: str):
603 """Rename a cube.
605 Arguments
606 ---------
607 cubes: iris.cube.Cube | iris.cube.CubeList
608 A Cube or CubeList of a field to be renamed.
610 name: str
611 The new name of the cube. It should be CF compliant.
613 Returns
614 -------
615 iris.cube.Cube | iris.cube.CubeList
616 The renamed field.
618 Notes
619 -----
620 This operator is designed to be used when the output field name does not
621 match expectations or needs to be different to defaults in standard_name, var_name or
622 long_name. For example, if combining masks
623 to create light rain you would like the field to be named "mask_for_light_rain"
624 rather than "mask_for_microphysical_precip_gt_0.0_x_mask_for_microphysical_precip_lt_2.0".
626 Examples
627 --------
628 >>> light_rain_mask = misc.rename_cube(light_rain_mask,"mask_for_light_rainfall"
629 """
630 new_cubelist = iris.cube.CubeList([])
631 for cube in iter_maybe(cubes):
632 cube.rename(name)
633 new_cubelist.append(cube)
634 if len(new_cubelist) == 1:
635 return new_cubelist[0]
636 else:
637 return new_cubelist
640def _slice_cube_on_levels(cube: iris.cube.Cube, coord_name: str, levels: list):
641 """
642 Extract levels from a cube for a given coordinate.
644 Arguments
645 ---------
646 cube: iris.cube.Cube
647 A Cube to be sliced.
649 coord_name: str
650 The coordinate name to be sliced
652 levels: list
653 A list containing points to be extracted from the cube.
655 Returns
656 -------
657 iris.cube.Cube
658 The sliced cube.
659 """
660 coord = cube.coord(coord_name)
661 (dim_index,) = cube.coord_dims(coord)
663 mask = np.isin(coord.points, levels)
665 slicer = [slice(None)] * cube.ndim
666 slicer[dim_index] = mask
668 return cube[tuple(slicer)]
671def extract_common_points(cubes: iris.cube.CubeList, coordinate: str):
672 """
673 Extract common points for a given coordinate between cubes in a CubeList.
675 Parameters
676 ----------
677 cubes: iris.cube.CubeList
678 CubeList containing cubes for which to extract common points.
680 coordinate: str
681 The coordinate name to be checked for common points.
683 Returns
684 -------
685 iris.cube.CubeList
686 CubeList containing the two cubes sliced to common points
687 for the given coordinate.
688 """
689 # Check type of input
690 if type(cubes) is not iris.cube.CubeList:
691 raise TypeError(f"Not a CubeList, got type {type(cubes)}")
693 # Extract coordinate
694 try:
695 points_list = []
696 for cube in cubes:
697 points_list.append(cube.coord(coordinate).points)
698 except iris.exceptions.CoordinateNotFoundError as err:
699 raise ValueError(f"Both cubes must have an {coordinate} coordinate") from err
701 # Find common points
702 common_points = reduce(np.intersect1d, points_list)
704 # Check that common points is more than zero.
705 if common_points.size == 0:
706 raise ValueError("No common levels found")
708 # Extract common points
709 common_cubes = iris.cube.CubeList()
710 for cube in cubes:
711 common_cubes.append(_slice_cube_on_levels(cube, coordinate, common_points))
713 return common_cubes
716def differentiate(
717 cubes: iris.cube.Cube | iris.cube.CubeList, coordinate: str, **kwargs
718) -> iris.cube.Cube | iris.cube.CubeList:
719 """Differentiate a cube on a specified coordinate.
721 Arguments
722 ---------
723 cubes: iris.cube.Cube | iris.cube.CubeList
724 A Cube or CubeList of a field that is to be differentiated.
726 coordinate: str
727 The coordinate that is to be differentiated over.
729 Returns
730 -------
731 iris.cube.Cube | iris.cube.CubeList
732 The differential of the cube along the specified coordinate.
734 Notes
735 -----
736 The differential is calculated based on a carteisan grid. This calculation
737 is then suitable for vertical and temporal derivatives. It is not sensible
738 for horizontal derivatives if they are based on spherical coordinates (e.g.
739 latitude and longitude). In essence this operator is a CSET wrapper around
740 `iris.analysis.calculus.differentiate <https://scitools-iris.readthedocs.io/en/stable/generated/api/iris.analysis.calculus.html#iris.analysis.calculus.differentiate>`_.
742 Examples
743 --------
744 >>> dT_dz = misc.differentiate(temperature, "altitude")
745 """
746 new_cubelist = iris.cube.CubeList([])
747 for cube in iter_maybe(cubes):
748 dcube = iris.analysis.calculus.differentiate(cube, coordinate)
749 new_cubelist.append(dcube)
750 if len(new_cubelist) == 1:
751 return new_cubelist[0]
752 else:
753 return new_cubelist
756def split_obs_and_models(cubes: CubeList) -> tuple[Cube | None, CubeList]:
757 """Return (obs_cube or None, CubeList of model cubes)."""
758 obs = [cb for cb in cubes if "observed" in (cb.long_name or "")]
760 if len(obs) > 1: 760 ↛ 761line 760 didn't jump to line 761 because the condition on line 760 was never true
761 raise ValueError(f"Expected at most one 'observed' cube, found {len(obs)}")
762 if not obs: 762 ↛ 765line 762 didn't jump to line 765 because the condition on line 762 was always true
763 return None, CubeList(cubes)
765 base = obs[0]
766 return base, CubeList(cb for cb in cubes if cb is not base)