Coverage for src/CSET/operators/misc.py: 94%
171 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-11 15:11 +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 combine_cubes_into_cubelist(first: Cube | CubeList, **kwargs) -> CubeList:
135 """Operator that combines multiple cubes or CubeLists into one.
137 Arguments
138 ---------
139 first: Cube | CubeList
140 First cube or CubeList to merge into CubeList.
141 second: Cube | CubeList
142 Second cube or CubeList to merge into CubeList. This must be a named
143 argument.
144 third: Cube | CubeList
145 There can be any number of additional arguments, they just need unique
146 names.
147 ...
149 Returns
150 -------
151 combined_cubelist: CubeList
152 Combined CubeList containing all cubes/CubeLists.
154 Raises
155 ------
156 TypeError:
157 If the provided arguments are not either a Cube or CubeList.
158 """
159 # Create empty CubeList to store cubes/CubeList.
160 all_cubes = CubeList()
162 # Combine all CubeLists into a single flat iterable.
163 for item in itertools.chain(iter_maybe(first), *map(iter_maybe, kwargs.values())):
164 # Check each item is a Cube, erroring if not.
165 if isinstance(item, Cube):
166 # Add cube to CubeList.
167 all_cubes.append(item)
168 else:
169 raise TypeError("Not a Cube or CubeList!", item)
171 return all_cubes
174def difference(cubes: CubeList):
175 """Difference of two fields.
177 Parameters
178 ----------
179 cubes: CubeList
180 A list of exactly two cubes. One must have the cset_comparison_base
181 attribute set to 1, and will be used as the base of the comparison.
183 Returns
184 -------
185 Cube
187 Raises
188 ------
189 ValueError
190 When the cubes are not compatible.
192 Notes
193 -----
194 This is a simple operator designed for combination of diagnostics or
195 creating new diagnostics by using recipes. It can be used for model
196 differences to allow for comparisons between the same field in different
197 models or model configurations.
199 Examples
200 --------
201 >>> model_diff = misc.difference(temperature_model_A, temperature_model_B)
203 """
204 if len(cubes) != 2:
205 raise ValueError("cubes should contain exactly 2 cubes.")
206 base: Cube = cubes.extract_cube(iris.AttributeConstraint(cset_comparison_base=1))
207 other: Cube = cubes.extract_cube(
208 iris.Constraint(
209 cube_func=lambda cube: "cset_comparison_base" not in cube.attributes
210 )
211 )
213 # If cubes contain a pressure coordinate, ensure it is increasing.
214 for cube in cubes:
215 try:
216 if len(cube.coord("pressure").points) > 2 and not is_increasing(
217 cube.coord("pressure").points
218 ):
219 cube.data = np.flip(cube.data, axis=cube.coord_dims("pressure")[0])
221 except iris.exceptions.CoordinateNotFoundError:
222 pass
224 # Get spatial coord names.
225 base_lat_name, base_lon_name = get_cube_yxcoordname(base)
226 other_lat_name, other_lon_name = get_cube_yxcoordname(other)
228 # Ensure cubes to compare are on common differencing grid.
229 # This is triggered if either
230 # i) latitude and longitude shapes are not the same. Note grid points
231 # are not compared directly as these can differ through rounding
232 # errors.
233 # ii) or variables are known to often sit on different grid staggering
234 # in different models (e.g. cell center vs cell edge), as is the case
235 # for UM and LFRic comparisons.
236 # In future greater choice of regridding method might be applied depending
237 # on variable type. Linear regridding can in general be appropriate for smooth
238 # variables. Care should be taken with interpretation of differences
239 # given this dependency on regridding.
240 if (
241 base.coord(base_lat_name).shape != other.coord(other_lat_name).shape
242 or base.coord(base_lon_name).shape != other.coord(other_lon_name).shape
243 ) or (
244 base.long_name
245 in [
246 "eastward_wind_at_10m",
247 "northward_wind_at_10m",
248 "northward_wind_at_cell_centres",
249 "eastward_wind_at_cell_centres",
250 "zonal_wind_at_pressure_levels",
251 "meridional_wind_at_pressure_levels",
252 "potential_vorticity_at_pressure_levels",
253 "vapour_specific_humidity_at_pressure_levels_for_climate_averaging",
254 ]
255 ):
256 logger.debug("Linear regridding base cube to other grid to compute differences")
257 base = regrid_onto_cube(base, other, method="Linear")
259 # Figure out if we are comparing between UM and LFRic; flip array if so.
260 base_lat_direction = is_increasing(base.coord(base_lat_name).points)
261 other_lat_direction = is_increasing(other.coord(other_lat_name).points)
262 if base_lat_direction != other_lat_direction:
263 other.data = np.flip(other.data, other.coord(other_lat_name).cube_dims(other))
265 # Equalise attributes so we can merge.
266 fully_equalise_attributes([base, other])
267 logger.debug("Base: %s\nOther: %s", base, other)
269 # This currently relies on the cubes having the same underlying data layout.
270 difference = base.copy()
272 # Differences don't have a standard name; long name gets a suffix. We are
273 # assuming we can rely on cubes having a long name, so we don't check for
274 # its presents.
275 difference.standard_name = None
276 difference.long_name = (
277 base.long_name if base.long_name else base.name()
278 ) + "_difference"
279 if base.var_name:
280 difference.var_name = base.var_name + "_difference"
281 elif base.standard_name:
282 difference.var_name = base.standard_name + "_difference"
284 # will raise if not common times
285 ensure_common_time_points(base, other)
287 difference.data = other.data - base.data
288 return difference
291def extract_common_time_points(cubes: CubeList | Cube) -> CubeList | Cube:
292 """Return cubes restricted to their common time points.
294 If the first cube has more than one FRT point, every cube must have identical
295 FRT points. The cubes are then trimmed to their shared
296 forecast_period values, so lead times line up across cubes.
297 Otherwise (no FRT coordinate, or a single FRT), the cubes are
298 trimmed to their shared times.
300 Parameters
301 ----------
302 cubes: CubeList | Cube
303 Cubes to align in time. The first cube is used as the reference.
305 Returns
306 -------
307 CubeList | Cube
308 Either cubes containing only the time points common to all inputs
309 or original cube/cube list (no op if single cube or cubelist of length 1 given)
311 Raises
312 ------
313 ValueError
314 If the first cube has multiple FRT points and any other cube either
315 lacks a forecast_reference_time coordinate or has different FRT
316 values (compared element-wise, so order matters).
317 """
318 if isinstance(cubes, Cube) or len(cubes) == 1:
319 logger.debug("Only one cube provided, skipping equalisation.")
320 return cubes
322 multi_frt = any(
323 len(c.coord("forecast_reference_time").points) > 1
324 for c in cubes
325 if c.coords("forecast_reference_time")
326 )
328 if multi_frt:
329 for cb in cubes[1:]:
330 cb_frt_coord = cb.coords("forecast_reference_time")
331 if not cb_frt_coord:
332 raise ValueError(
333 "All cubes must have a forecast_reference_time coordinate."
334 )
335 cubes_common = extract_common_points(cubes, "forecast_period")
337 return cubes_common
338 else:
339 return _extract_common_time_points_one_frt(cubes)
342def _ensure_same_frt(cubes: CubeList) -> None:
343 """Ensure that all cubes have the same forecast_reference_time(s)."""
345 def _frt_points(cb: Cube):
346 coords = cb.coords("forecast_reference_time")
347 return coords[0].points if coords else None
349 reference = _frt_points(cubes[0])
350 if reference is None: 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
351 return
353 for cube in cubes[1:]:
354 points = _frt_points(cube)
355 if points is None:
356 raise ValueError("Cube does not contain any forecast_reference_times.")
357 elif len(points) != len(reference):
358 raise ValueError(
359 "Cubes do not have same number of forecast_reference_times"
360 )
361 elif not np.array_equal(reference, points):
362 raise ValueError(
363 "Cubes do not share the same forecast_reference_time values."
364 )
367def _extract_common_time_points_one_frt(cubes: CubeList) -> CubeList:
368 """Extract common time points from cubes where there is only one frt."""
369 obs: Cube | None
370 models: CubeList
371 obs, models = split_obs_and_models(cubes)
372 if obs is not None: 372 ↛ 373line 372 didn't jump to line 373 because the condition on line 372 was never true
373 base = obs
374 others = models[0:]
375 else:
376 base = models[0]
377 others = models[1:]
378 list_shared_times: list[set[int]] = []
379 time_coord = get_time_coord(cubes)
381 if not time_coord:
382 logger.debug("No time coord, skipping equalisation.")
383 return cubes
385 for other in iter_maybe(others):
386 list_shared_times.append(get_shared_times(base, other, time_coord))
387 shared_times_all = reduce(lambda x, y: x.intersection(y), list_shared_times)
389 logger.debug("Shared times: %s", shared_times_all)
391 time_constraint = iris.Constraint(
392 coord_values={
393 time_coord: lambda cell, shared_times=shared_times_all: (
394 cell.point in shared_times
395 )
396 }
397 )
399 # Extract points matching the shared times.
400 cubes_common_time = CubeList()
401 for cube in cubes:
402 cubes_common_time.append(cube.extract(time_constraint))
404 return cubes_common_time
407def convert_units(cubes: iris.cube.Cube | iris.cube.CubeList, units: str):
408 """Convert the units of a cube.
410 Arguments
411 ---------
412 cubes: iris.cube.Cube | iris.cube.CubeList
413 A Cube or CubeList of a field for its units to be converted.
415 units: str
416 The unit that the original field is to be converted to. It takes
417 CF compliant units.
419 Returns
420 -------
421 iris.cube.Cube | iris.cube.CubeList
422 The field converted into the specified units.
424 Examples
425 --------
426 >>> T_in_F = misc.convert_units(temperature_in_K, "Fahrenheit")
428 """
429 new_cubelist = iris.cube.CubeList([])
430 for cube in iter_maybe(cubes):
431 # Copy cube to keep original data.
432 cube_a = cube.copy()
433 # Convert cube units.
434 cube_a.convert_units(units)
435 new_cubelist.append(cube_a)
436 if len(new_cubelist) == 1:
437 return new_cubelist[0]
438 else:
439 return new_cubelist
442def rename_cube(cubes: iris.cube.Cube | iris.cube.CubeList, name: str):
443 """Rename a cube.
445 Arguments
446 ---------
447 cubes: iris.cube.Cube | iris.cube.CubeList
448 A Cube or CubeList of a field to be renamed.
450 name: str
451 The new name of the cube. It should be CF compliant.
453 Returns
454 -------
455 iris.cube.Cube | iris.cube.CubeList
456 The renamed field.
458 Notes
459 -----
460 This operator is designed to be used when the output field name does not
461 match expectations or needs to be different to defaults in standard_name, var_name or
462 long_name. For example, if combining masks
463 to create light rain you would like the field to be named "mask_for_light_rain"
464 rather than "mask_for_microphysical_precip_gt_0.0_x_mask_for_microphysical_precip_lt_2.0".
466 Examples
467 --------
468 >>> light_rain_mask = misc.rename_cube(light_rain_mask,"mask_for_light_rainfall"
469 """
470 new_cubelist = iris.cube.CubeList([])
471 for cube in iter_maybe(cubes):
472 cube.rename(name)
473 new_cubelist.append(cube)
474 if len(new_cubelist) == 1:
475 return new_cubelist[0]
476 else:
477 return new_cubelist
480def _slice_cube_on_levels(cube: iris.cube.Cube, coord_name: str, levels: list):
481 """
482 Extract levels from a cube for a given coordinate.
484 Arguments
485 ---------
486 cube: iris.cube.Cube
487 A Cube to be sliced.
489 coord_name: str
490 The coordinate name to be sliced
492 levels: list
493 A list containing points to be extracted from the cube.
495 Returns
496 -------
497 iris.cube.Cube
498 The sliced cube.
499 """
500 coord = cube.coord(coord_name)
501 (dim_index,) = cube.coord_dims(coord)
503 mask = np.isin(coord.points, levels)
505 slicer = [slice(None)] * cube.ndim
506 slicer[dim_index] = mask
508 return cube[tuple(slicer)]
511def extract_common_points(cubes: iris.cube.CubeList, coordinate: str):
512 """
513 Extract common points for a given coordinate between cubes in a CubeList.
515 Parameters
516 ----------
517 cubes: iris.cube.CubeList
518 CubeList containing cubes for which to extract common points.
520 coordinate: str
521 The coordinate name to be checked for common points.
523 Returns
524 -------
525 iris.cube.CubeList
526 CubeList containing the two cubes sliced to common points
527 for the given coordinate.
528 """
529 # Check type of input
530 if type(cubes) is not iris.cube.CubeList:
531 raise TypeError(f"Not a CubeList, got type {type(cubes)}")
533 # Extract coordinate
534 try:
535 points_list = []
536 for cube in cubes:
537 points_list.append(cube.coord(coordinate).points)
538 except iris.exceptions.CoordinateNotFoundError as err:
539 raise ValueError(f"Both cubes must have an {coordinate} coordinate") from err
541 # Find common points
542 common_points = reduce(np.intersect1d, points_list)
544 # Check that common points is more than zero.
545 if common_points.size == 0:
546 raise ValueError("No common levels found")
548 # Extract common points
549 common_cubes = iris.cube.CubeList()
550 for cube in cubes:
551 common_cubes.append(_slice_cube_on_levels(cube, coordinate, common_points))
552 return common_cubes
555def split_obs_and_models(cubes: CubeList) -> tuple[Cube | None, CubeList]:
556 """Return (obs_cube or None, CubeList of model cubes)."""
557 obs = [cb for cb in cubes if "observed" in (cb.long_name or "")]
559 if len(obs) > 1: 559 ↛ 560line 559 didn't jump to line 560 because the condition on line 559 was never true
560 raise ValueError(f"Expected at most one 'observed' cube, found {len(obs)}")
561 if not obs: 561 ↛ 564line 561 didn't jump to line 564 because the condition on line 561 was always true
562 return None, CubeList(cubes)
564 base = obs[0]
565 return base, CubeList(cb for cb in cubes if cb is not base)