Coverage for src/CSET/operators/misc.py: 86%
158 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-10 13:27 +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
21import iris
22import iris.analysis.calculus
23import numpy as np
24from iris.cube import Cube, CubeList
26from CSET._common import is_increasing, iter_maybe
27from CSET.operators._utils import fully_equalise_attributes, get_cube_yxcoordname
28from CSET.operators.regrid import regrid_onto_cube
31def noop(x, **kwargs):
32 """Return its input without doing anything to it.
34 Useful for constructing diagnostic chains.
36 Arguments
37 ---------
38 x: Any
39 Input to return.
41 Returns
42 -------
43 x: Any
44 The input that was given.
45 """
46 return x
49def remove_attribute(
50 cubes: Cube | CubeList, attribute: str | Iterable, **kwargs
51) -> CubeList:
52 """Remove a cube attribute.
54 If the attribute is not on the cube, the cube is passed through unchanged.
56 Arguments
57 ---------
58 cubes: Cube | CubeList
59 One or more cubes to remove the attribute from.
60 attribute: str | Iterable
61 Name of attribute (or Iterable of names) to remove.
63 Returns
64 -------
65 cubes: CubeList
66 CubeList of cube(s) with the attribute removed.
67 """
68 # Ensure cubes is a CubeList.
69 if not isinstance(cubes, CubeList):
70 cubes = CubeList(iter_maybe(cubes))
72 for cube in cubes:
73 for attr in iter_maybe(attribute):
74 cube.attributes.pop(attr, None)
76 # Combine things that can be merged due to remove removing the
77 # attributes.
78 cubes = cubes.merge()
79 # combine items that can be merged after removing unwanted attributes
80 cubes = cubes.concatenate()
81 return cubes
84def remove_scalar_coords(
85 cubes: Cube | CubeList, coords: str | Iterable[str]
86) -> CubeList:
87 """Remove scalar coordinates from one or more cubes.
89 Coordinates are only removed if they exist on the cube and are
90 scalar coordinates (i.e. have no associated dimensions). Examples
91 include ``realization`` and ``forecast_reference_time`` on model
92 data cubes. Dimensional and non-scalar auxiliary coordinates are
93 left unchanged.
95 Arguments
96 ---------
97 cubes: Cube | CubeList
98 One or more cubes from which scalar coordinates will be removed.
99 coords: str | Iterable
100 Name of a coordinate (or Iterable of coordinate names) to remove.
102 Returns
103 -------
104 cubes: CubeList
105 CubeList of cube(s) with the requested scalar coordinates
106 removed where present.
107 """
108 if not isinstance(cubes, CubeList): 108 ↛ 111line 108 didn't jump to line 111 because the condition on line 108 was always true
109 cubes = CubeList(iter_maybe(cubes))
111 if isinstance(coords, str): 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true
112 coords = [coords]
114 for cube in cubes:
115 for coord_name in iter_maybe(coords):
116 if cube.coords(coord_name): 116 ↛ 115line 116 didn't jump to line 115 because the condition on line 116 was always true
117 coord = cube.coord(coord_name)
118 # only remove if scalar
119 if cube.coord_dims(coord) == ():
120 cube.remove_coord(coord)
122 return cubes
125def addition(addend_1, addend_2):
126 """Addition of two fields.
128 Parameters
129 ----------
130 addend_1: Cube
131 Any field to have another field added to it.
132 addend_2: Cube
133 Any field to be added to another field.
135 Returns
136 -------
137 Cube
139 Raises
140 ------
141 ValueError, iris.exceptions.NotYetImplementedError
142 When the cubes are not compatible.
144 Notes
145 -----
146 This is a simple operator designed for combination of diagnostics or
147 creating new diagnostics by using recipes.
149 Examples
150 --------
151 >>> field_addition = misc.addition(kinetic_energy_u, kinetic_energy_v)
153 """
154 return addend_1 + addend_2
157def subtraction(minuend, subtrahend):
158 """Subtraction of two fields.
160 Parameters
161 ----------
162 minuend: Cube
163 Any field to have another field subtracted from it.
164 subtrahend: Cube
165 Any field to be subtracted from to another field.
167 Returns
168 -------
169 Cube
171 Raises
172 ------
173 ValueError, iris.exceptions.NotYetImplementedError
174 When the cubes are not compatible.
176 Notes
177 -----
178 This is a simple operator designed for combination of diagnostics or
179 creating new diagnostics by using recipes. It can be used for model
180 differences to allow for comparisons between the same field in different
181 models or model configurations.
183 Examples
184 --------
185 >>> model_diff = misc.subtraction(temperature_model_A, temperature_model_B)
187 """
188 return minuend - subtrahend
191def division(numerator, denominator):
192 """Division of two fields.
194 Parameters
195 ----------
196 numerator: Cube
197 Any field to have the ratio taken with respect to another field.
198 denominator: Cube
199 Any field used to divide another field or provide the reference
200 value in a ratio.
202 Returns
203 -------
204 Cube
206 Raises
207 ------
208 ValueError
209 When the cubes are not compatible.
211 Notes
212 -----
213 This is a simple operator designed for combination of diagnostics or
214 creating new diagnostics by using recipes.
216 Examples
217 --------
218 >>> bowen_ratio = misc.division(sensible_heat_flux, latent_heat_flux)
220 """
221 return numerator / denominator
224def multiplication(
225 multiplicand: Cube | CubeList, multiplier: Cube | CubeList
226) -> Cube | CubeList:
227 """Multiplication of two fields.
229 Parameters
230 ----------
231 multiplicand: Cube | CubeList
232 Any field to be multiplied by another field.
233 multiplier: Cube | CubeList
234 Any field to be multiplied to another field.
236 Returns
237 -------
238 Cube | CubeList
239 The result of multiplicand x multiplier.
241 Raises
242 ------
243 ValueError
244 When the cubes are not compatible.
246 Notes
247 -----
248 This is a simple operator designed for combination of diagnostics or
249 creating new diagnostics by using recipes. CubeLists are multiplied
250 on a strict ordering (e.g. first cube with first cube).
252 Examples
253 --------
254 >>> filtered_CAPE_ratio = misc.multiplication(CAPE_ratio, inflow_layer_properties)
256 """
257 new_cubelist = iris.cube.CubeList([])
258 for cube_a, cube_b in zip(
259 iter_maybe(multiplicand), iter_maybe(multiplier), strict=True
260 ):
261 multiplied_cube = cube_a * cube_b
262 multiplied_cube.rename(f"{cube_a.name()}_x_{cube_b.name()}")
263 new_cubelist.append(multiplied_cube)
264 if len(new_cubelist) == 1:
265 return new_cubelist[0]
266 else:
267 return new_cubelist
270def combine_cubes_into_cubelist(first: Cube | CubeList, **kwargs) -> CubeList:
271 """Operator that combines multiple cubes or CubeLists into one.
273 Arguments
274 ---------
275 first: Cube | CubeList
276 First cube or CubeList to merge into CubeList.
277 second: Cube | CubeList
278 Second cube or CubeList to merge into CubeList. This must be a named
279 argument.
280 third: Cube | CubeList
281 There can be any number of additional arguments, they just need unique
282 names.
283 ...
285 Returns
286 -------
287 combined_cubelist: CubeList
288 Combined CubeList containing all cubes/CubeLists.
290 Raises
291 ------
292 TypeError:
293 If the provided arguments are not either a Cube or CubeList.
294 """
295 # Create empty CubeList to store cubes/CubeList.
296 all_cubes = CubeList()
297 # Combine all CubeLists into a single flat iterable.
298 for item in itertools.chain(iter_maybe(first), *map(iter_maybe, kwargs.values())):
299 # Check each item is a Cube, erroring if not.
300 if isinstance(item, Cube):
301 # Add cube to CubeList.
302 all_cubes.append(item)
303 else:
304 raise TypeError("Not a Cube or CubeList!", item)
305 return all_cubes
308def difference(cubes: CubeList):
309 """Difference of two fields.
311 Parameters
312 ----------
313 cubes: CubeList
314 A list of exactly two cubes. One must have the cset_comparison_base
315 attribute set to 1, and will be used as the base of the comparison.
317 Returns
318 -------
319 Cube
321 Raises
322 ------
323 ValueError
324 When the cubes are not compatible.
326 Notes
327 -----
328 This is a simple operator designed for combination of diagnostics or
329 creating new diagnostics by using recipes. It can be used for model
330 differences to allow for comparisons between the same field in different
331 models or model configurations.
333 Examples
334 --------
335 >>> model_diff = misc.difference(temperature_model_A, temperature_model_B)
337 """
338 if len(cubes) != 2:
339 raise ValueError("cubes should contain exactly 2 cubes.")
340 base: Cube = cubes.extract_cube(iris.AttributeConstraint(cset_comparison_base=1))
341 other: Cube = cubes.extract_cube(
342 iris.Constraint(
343 cube_func=lambda cube: "cset_comparison_base" not in cube.attributes
344 )
345 )
347 # If cubes contain a pressure coordinate, ensure it is increasing.
348 for cube in cubes:
349 try:
350 if len(cube.coord("pressure").points) > 2: 350 ↛ 348line 350 didn't jump to line 348 because the condition on line 350 was always true
351 if not is_increasing(cube.coord("pressure").points):
352 cube.data = np.flip(cube.data, axis=cube.coord_dims("pressure")[0])
354 except iris.exceptions.CoordinateNotFoundError:
355 pass
357 # Get spatial coord names.
358 base_lat_name, base_lon_name = get_cube_yxcoordname(base)
359 other_lat_name, other_lon_name = get_cube_yxcoordname(other)
361 # Ensure cubes to compare are on common differencing grid.
362 # This is triggered if either
363 # i) latitude and longitude shapes are not the same. Note grid points
364 # are not compared directly as these can differ through rounding
365 # errors.
366 # ii) or variables are known to often sit on different grid staggering
367 # in different models (e.g. cell center vs cell edge), as is the case
368 # for UM and LFRic comparisons.
369 # In future greater choice of regridding method might be applied depending
370 # on variable type. Linear regridding can in general be appropriate for smooth
371 # variables. Care should be taken with interpretation of differences
372 # given this dependency on regridding.
373 if (
374 base.coord(base_lat_name).shape != other.coord(other_lat_name).shape
375 or base.coord(base_lon_name).shape != other.coord(other_lon_name).shape
376 ) or (
377 base.long_name
378 in [
379 "eastward_wind_at_10m",
380 "northward_wind_at_10m",
381 "northward_wind_at_cell_centres",
382 "eastward_wind_at_cell_centres",
383 "zonal_wind_at_pressure_levels",
384 "meridional_wind_at_pressure_levels",
385 "potential_vorticity_at_pressure_levels",
386 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",
387 ]
388 ):
389 logging.debug(
390 "Linear regridding base cube to other grid to compute differences"
391 )
392 base = regrid_onto_cube(base, other, method="Linear")
394 # Figure out if we are comparing between UM and LFRic; flip array if so.
395 base_lat_direction = is_increasing(base.coord(base_lat_name).points)
396 other_lat_direction = is_increasing(other.coord(other_lat_name).points)
397 if base_lat_direction != other_lat_direction:
398 other.data = np.flip(other.data, other.coord(other_lat_name).cube_dims(other))
400 # Extract just common time points.
401 base, other = _extract_common_time_points(base, other)
403 # Equalise attributes so we can merge.
404 fully_equalise_attributes([base, other])
405 logging.debug("Base: %s\nOther: %s", base, other)
407 # This currently relies on the cubes having the same underlying data layout.
408 difference = base.copy()
410 # Differences don't have a standard name; long name gets a suffix. We are
411 # assuming we can rely on cubes having a long name, so we don't check for
412 # its presents.
413 difference.standard_name = None
414 difference.long_name = (
415 base.long_name if base.long_name else base.name()
416 ) + "_difference"
417 if base.var_name:
418 difference.var_name = base.var_name + "_difference"
419 elif base.standard_name:
420 difference.var_name = base.standard_name + "_difference"
422 difference.data = other.data - base.data
423 return difference
426def _extract_common_time_points(base: Cube, other: Cube) -> tuple[Cube, Cube]:
427 """Extract common time points from cubes to allow comparison."""
428 # Get the name of the first non-scalar time coordinate.
429 time_coord = next(
430 map(
431 lambda coord: coord.name(),
432 filter(
433 lambda coord: coord.shape > (1,) and coord.name() in ["time", "hour"],
434 base.coords(),
435 ),
436 ),
437 None,
438 )
439 if not time_coord:
440 logging.debug("No time coord, skipping equalisation.")
441 return (base, other)
442 base_time_coord = base.coord(time_coord)
443 other_time_coord = other.coord(time_coord)
444 logging.debug("Base: %s\nOther: %s", base_time_coord, other_time_coord)
445 if time_coord == "hour":
446 # We directly compare points when comparing coordinates with
447 # non-absolute units, such as hour. We can't just check the units are
448 # equal as iris automatically converts to datetime objects in the
449 # comparison for certain coordinate names.
450 base_times = base_time_coord.points
451 other_times = other_time_coord.points
452 shared_times = set.intersection(set(base_times), set(other_times))
453 else:
454 # Units don't match, so converting to datetimes for comparison.
455 base_times = base_time_coord.units.num2date(base_time_coord.points)
456 other_times = other_time_coord.units.num2date(other_time_coord.points)
457 shared_times = set.intersection(set(base_times), set(other_times))
458 logging.debug("Shared times: %s", shared_times)
459 time_constraint = iris.Constraint(
460 coord_values={time_coord: lambda cell: cell.point in shared_times}
461 )
462 # Extract points matching the shared times.
463 base = base.extract(time_constraint)
464 other = other.extract(time_constraint)
465 if base is None or other is None:
466 raise ValueError("No common time points found!")
467 return (base, other)
470def convert_units(cubes: iris.cube.Cube | iris.cube.CubeList, units: str):
471 """Convert the units of a cube.
473 Arguments
474 ---------
475 cubes: iris.cube.Cube | iris.cube.CubeList
476 A Cube or CubeList of a field for its units to be converted.
478 units: str
479 The unit that the original field is to be converted to. It takes
480 CF compliant units.
482 Returns
483 -------
484 iris.cube.Cube | iris.cube.CubeList
485 The field converted into the specified units.
487 Examples
488 --------
489 >>> T_in_F = misc.convert_units(temperature_in_K, "Fahrenheit")
491 """
492 new_cubelist = iris.cube.CubeList([])
493 for cube in iter_maybe(cubes):
494 # Copy cube to keep original data.
495 cube_a = cube.copy()
496 # Convert cube units.
497 cube_a.convert_units(units)
498 new_cubelist.append(cube_a)
499 if len(new_cubelist) == 1:
500 return new_cubelist[0]
501 else:
502 return new_cubelist
505def rename_cube(cubes: iris.cube.Cube | iris.cube.CubeList, name: str):
506 """Rename a cube.
508 Arguments
509 ---------
510 cubes: iris.cube.Cube | iris.cube.CubeList
511 A Cube or CubeList of a field to be renamed.
513 name: str
514 The new name of the cube. It should be CF compliant.
516 Returns
517 -------
518 iris.cube.Cube | iris.cube.CubeList
519 The renamed field.
521 Notes
522 -----
523 This operator is designed to be used when the output field name does not
524 match expectations or needs to be different to defaults in standard_name, var_name or
525 long_name. For example, if combining masks
526 to create light rain you would like the field to be named "mask_for_light_rain"
527 rather than "mask_for_microphysical_precip_gt_0.0_x_mask_for_microphysical_precip_lt_2.0".
529 Examples
530 --------
531 >>> light_rain_mask = misc.rename_cube(light_rain_mask,"mask_for_light_rainfall"
532 """
533 new_cubelist = iris.cube.CubeList([])
534 for cube in iter_maybe(cubes):
535 cube.rename(name)
536 new_cubelist.append(cube)
537 if len(new_cubelist) == 1:
538 return new_cubelist[0]
539 else:
540 return new_cubelist
543def _slice_cube_on_levels(cube: iris.cube.Cube, coord_name: str, levels: list):
544 """
545 Extract levels from a cube for a given coordinate.
547 Arguments
548 ---------
549 cube: iris.cube.Cube
550 A Cube to be sliced.
552 coord_name: str
553 The coordinate name to be sliced
555 levels: list
556 A list containing points to be extracted from the cube.
558 Returns
559 -------
560 iris.cube.Cube
561 The sliced cube.
562 """
563 coord = cube.coord(coord_name)
564 (dim_index,) = cube.coord_dims(coord)
566 mask = np.isin(coord.points, levels)
568 slicer = [slice(None)] * cube.ndim
569 slicer[dim_index] = mask
571 return cube[tuple(slicer)]
574def extract_common_points(cubes: iris.cube.CubeList, coordinate: str):
575 """
576 Extract common points for a given coordinate between two cubes.
578 Parameters
579 ----------
580 cubes: iris.cube.CubeList
581 CubeList containing exactly two cubes.
583 coordinate: str
584 The coordinate name to be checked for common levels.
586 Returns
587 -------
588 iris.cube.CubeList
589 CubeList containing the two cubes sliced to common levels
590 for the given coordinate.
591 """
592 # Check type of input
593 if type(cubes) is not iris.cube.CubeList:
594 raise TypeError(f"Not a CubeList, got type {type(cubes)}")
596 # Check that only two cubes are passed into function.
597 if len(cubes) != 2:
598 raise ValueError(f"Maximum of two cubes allowed, received {len(cubes)}")
600 # Extract coordinate
601 try:
602 p1 = cubes[0].coord(coordinate)
603 p2 = cubes[1].coord(coordinate)
604 except iris.exceptions.CoordinateNotFoundError as err:
605 raise ValueError(f"Both cubes must have an {coordinate} coordinate") from err
607 # Find common points
608 common_points = np.intersect1d(p1.points, p2.points)
610 # Check that common points is more than zero.
611 if common_points.size == 0:
612 raise ValueError("No common levels found")
614 # Extract common points
615 cube0_common = _slice_cube_on_levels(cubes[0], coordinate, common_points)
616 cube1_common = _slice_cube_on_levels(cubes[1], coordinate, common_points)
618 return iris.cube.CubeList([cube0_common, cube1_common])
621def differentiate(
622 cubes: iris.cube.Cube | iris.cube.CubeList, coordinate: str, **kwargs
623) -> iris.cube.Cube | iris.cube.CubeList:
624 """Differentiate a cube on a specified coordinate.
626 Arguments
627 ---------
628 cubes: iris.cube.Cube | iris.cube.CubeList
629 A Cube or CubeList of a field that is to be differentiated.
631 coordinate: str
632 The coordinate that is to be differentiated over.
634 Returns
635 -------
636 iris.cube.Cube | iris.cube.CubeList
637 The differential of the cube along the specified coordinate.
639 Notes
640 -----
641 The differential is calculated based on a carteisan grid. This calculation
642 is then suitable for vertical and temporal derivatives. It is not sensible
643 for horizontal derivatives if they are based on spherical coordinates (e.g.
644 latitude and longitude). In essence this operator is a CSET wrapper around
645 `iris.analysis.calculus.differentiate <https://scitools-iris.readthedocs.io/en/stable/generated/api/iris.analysis.calculus.html#iris.analysis.calculus.differentiate>`_.
647 Examples
648 --------
649 >>> dT_dz = misc.differentiate(temperature, "altitude")
650 """
651 new_cubelist = iris.cube.CubeList([])
652 for cube in iter_maybe(cubes):
653 dcube = iris.analysis.calculus.differentiate(cube, coordinate)
654 new_cubelist.append(dcube)
655 if len(new_cubelist) == 1:
656 return new_cubelist[0]
657 else:
658 return new_cubelist