Coverage for src/CSET/operators/_utils.py: 95%
232 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"""
16Common operator functionality used across CSET.
18Functions below should only be added if it is not suitable as a standalone
19operator, and will be used across multiple operators.
20"""
22import logging
23import os
24import re
25from datetime import timedelta
26from pathlib import Path
28import iris
29import iris.coords
30import iris.cube
31import iris.exceptions
32import iris.util
33import numpy as np
34from iris.cube import Cube, CubeList
35from iris.time import PartialDateTime
37from CSET._common import iter_maybe
39logger = logging.getLogger(__name__)
42def pdt_fromisoformat(
43 datestring,
44) -> tuple[iris.time.PartialDateTime, timedelta | None]:
45 """Generate PartialDateTime object.
47 Function that takes an ISO 8601 date string and returns a PartialDateTime object.
49 Arguments
50 ---------
51 datestring: str
52 ISO 8601 date.
54 Returns
55 -------
56 time_object: iris.time.PartialDateTime
57 """
59 def make_offset(sign, value) -> timedelta:
60 if len(value) not in [2, 4, 5]: 60 ↛ 61line 60 didn't jump to line 61 because the condition on line 60 was never true
61 raise ValueError(f'expected "hh", "hhmm", or "hh:mm", got {value}')
63 hours = int(value[:2])
64 minutes = 0
65 if len(value) in [4, 5]:
66 minutes = int(value[-2:])
67 return timedelta(hours=sign * hours, minutes=sign * minutes)
69 # Remove the microseconds coord due to no support in PartialDateTime
70 datestring = re.sub(r"\.\d+", "", datestring)
72 datetime_split = datestring.split("T")
73 date = datetime_split[0]
74 if len(datetime_split) == 1:
75 time = ""
76 elif len(datetime_split) == 2: 76 ↛ 79line 76 didn't jump to line 79 because the condition on line 76 was always true
77 time = datetime_split[1]
78 else:
79 raise ValueError("datesting in an unexpected format")
81 offset = None
82 time_split = time.split("+")
83 if len(time_split) == 2:
84 time = time_split[0]
85 offset = make_offset(1, time_split[1])
86 else:
87 time_split = time.split("-")
88 if len(time_split) == 2: 88 ↛ 89line 88 didn't jump to line 89 because the condition on line 88 was never true
89 time = time_split[0]
90 offset = make_offset(-1, time_split[1])
91 else:
92 offset = None
94 if re.fullmatch(r"\d{8}", date):
95 date = f"{date[0:4]}-{date[4:6]}-{date[6:8]}"
96 elif re.fullmatch(r"\d{6}", date):
97 date = f"{date[0:4]}-{date[4:6]}"
99 if len(date) < 7: 99 ↛ 100line 99 didn't jump to line 100 because the condition on line 99 was never true
100 raise ValueError(f"Invalid datestring: {datestring}, must be at least YYYY-MM")
102 # Returning a PartialDateTime for the special case of string form "YYYY-MM"
103 if re.fullmatch(r"\d{4}-\d{2}", date):
104 pdt = PartialDateTime(
105 year=int(date[0:4]),
106 month=int(date[5:7]),
107 day=None,
108 hour=None,
109 minute=None,
110 second=None,
111 )
112 return pdt, offset
114 year = int(date[0:4])
115 month = int(date[5:7])
116 day = int(date[8:10])
118 kwargs = {
119 "year": year,
120 "month": month,
121 "day": day,
122 "hour": 0,
123 "minute": 0,
124 "second": 0,
125 }
127 # Normalise the time parts into standard format
128 if re.fullmatch(r"\d{4}", time): 128 ↛ 129line 128 didn't jump to line 129 because the condition on line 128 was never true
129 time = f"{time[0:2]}:{time[2:4]}"
130 if re.fullmatch(r"\d{6}", time):
131 time = f"{time[0:2]}:{time[2:4]}:{time[4:6]}"
133 if len(time) >= 2:
134 kwargs["hour"] = int(time[0:2])
135 if len(time) >= 5:
136 kwargs["minute"] = int(time[3:5])
137 if len(time) >= 8:
138 kwargs["second"] = int(time[6:8])
140 pdt = PartialDateTime(**kwargs)
142 return pdt, offset
145def get_cube_yxcoordname(cube: iris.cube.Cube) -> tuple[str, str]:
146 """
147 Return horizontal dimension coordinate name(s) from a given cube.
149 Arguments
150 ---------
152 cube: iris.cube.Cube
153 An iris cube which will be checked to see if it contains coordinate
154 names that match a pre-defined list of acceptable horizontal
155 dimension coordinate names.
157 Returns
158 -------
159 (y_coord, x_coord)
160 A tuple containing the horizontal coordinate name for latitude and longitude respectively
161 found within the cube.
163 Raises
164 ------
165 ValueError
166 If a unique y/x horizontal coordinate cannot be found.
167 """
168 # Acceptable horizontal coordinate names.
169 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
170 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
172 # Get a list of dimension coordinate names for the cube
173 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
174 coord_names = [coord.name() for coord in cube.coords()]
176 # Check which x-coordinate we have, if any
177 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
178 if len(x_coords) != 1:
179 x_coords = [coord for coord in dim_coord_names if coord in X_COORD_NAMES]
180 if len(x_coords) != 1:
181 raise ValueError("Could not identify a unique x-coordinate in cube")
183 # Check which y-coordinate we have, if any
184 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
185 if len(y_coords) != 1:
186 y_coords = [coord for coord in dim_coord_names if coord in Y_COORD_NAMES]
187 if len(y_coords) != 1:
188 raise ValueError("Could not identify a unique y-coordinate in cube")
190 return (y_coords[0], x_coords[0])
193def get_cube_coordindex(cube: iris.cube.Cube, coord_name) -> int:
194 """
195 Return coordinate dimension for a named coordinate from a given cube.
197 Arguments
198 ---------
200 cube: iris.cube.Cube
201 An iris cube which will be checked to see if it contains coordinate
202 names that match a pre-defined list of acceptable horizontal
203 coordinate names.
205 coord_name: str
206 A cube dimension name
208 Returns
209 -------
210 coord_index
211 An integer specifying where in the cube dimension list a specified coordinate name is found.
213 Raises
214 ------
215 ValueError
216 If a specified dimension coordinate cannot be found.
217 """
218 # Get a list of dimension coordinate names for the cube
219 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
221 # Check if requested dimension is found in cube and get index
222 if coord_name in coord_names:
223 coord_index = cube.coord_dims(coord_name)[0]
224 else:
225 raise ValueError("Could not find requested dimension %s", coord_name)
227 return coord_index
230def is_spatialdim(cube: iris.cube.Cube) -> bool:
231 """Determine whether a cube is has two spatial dimension coordinates.
233 If cube has both spatial dims, it will contain two unique coordinates
234 that explain space (latitude and longitude). The coordinates have to
235 be iterable/contain usable dimension data, as cubes may contain these
236 coordinates as scalar dimensions after being collapsed.
238 Arguments
239 ---------
240 cube: iris.cube.Cube
241 An iris cube which will be checked to see if it contains coordinate
242 names that match a pre-defined list of acceptable coordinate names.
244 Returns
245 -------
246 bool
247 If true, then the cube has a spatial projection and thus can be plotted
248 as a map.
249 """
250 # Acceptable horizontal coordinate names.
251 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
252 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
254 # Get a list of coordinate names for the cube
255 coord_names = [coord.name() for coord in cube.dim_coords]
256 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
257 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
259 # If there is one coordinate for both x and y direction return True.
260 return len(x_coords) == 1 and len(y_coords) == 1
263def is_coorddim(cube: iris.cube.Cube, coord_name) -> bool:
264 """Determine whether a cube has specified dimension coordinates.
266 Arguments
267 ---------
268 cube: iris.cube.Cube
269 An iris cube which will be checked to see if it contains coordinate
270 names that match a pre-defined list of acceptable coordinate names.
272 coord_name: str
273 A cube dimension name
275 Returns
276 -------
277 bool
278 If true, then the cube has a spatial projection and thus can be plotted
279 as a map.
280 """
281 # Get a list of dimension coordinate names for the cube
282 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
284 # Check if requested dimension is found in cube and get index
285 return coord_name in coord_names
288def is_transect(cube: iris.cube.Cube) -> bool:
289 """Determine whether a cube is a transect.
291 If cube is a transect, it will contain only one spatial (map) coordinate,
292 and one vertical coordinate (either pressure or model level).
294 Arguments
295 ---------
296 cube: iris.cube.Cube
297 An iris cube which will be checked to see if it contains coordinate
298 names that match a pre-defined list of acceptable coordinate names.
300 Returns
301 -------
302 bool
303 If true, then the cube is a transect that contains one spatial (map)
304 coordinate and one vertical coordinate.
305 """
306 # Acceptable spatial (map) coordinate names.
307 SPATIAL_MAP_COORD_NAMES = [
308 "longitude",
309 "grid_longitude",
310 "projection_x_coordinate",
311 "x",
312 "latitude",
313 "grid_latitude",
314 "projection_y_coordinate",
315 "y",
316 "distance",
317 ]
319 # Acceptable vertical coordinate names
320 VERTICAL_COORD_NAMES = ["pressure", "model_level_number", "level_height"]
322 # Get a list of coordinate names for the cube
323 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
325 # Check which spatial coordinates we have.
326 spatial_coords = [
327 coord for coord in coord_names if coord in SPATIAL_MAP_COORD_NAMES
328 ]
329 if len(spatial_coords) != 1:
330 return False
332 # Check which vertical coordinates we have.
333 vertical_coords = [coord for coord in coord_names if coord in VERTICAL_COORD_NAMES]
334 if len(vertical_coords) != 1: # noqa: SIM103 Clearer to keep separate.
335 return False
337 # Passed criteria so return True
338 return True
341def check_stamp_coordinate(cube: iris.cube.Cube) -> str:
342 """
343 Return stamp dimension coordinate name from a given cube, if exists.
345 If cube contains a valid stamp coordinate as a dimension coordinate,
346 function will return name of the stamp coordinate.
348 Arguments
349 ---------
350 cube: iris.cube.Cube
351 An iris cube which will be checked to see if it contains coordinate
352 names that match a pre-defined list of acceptable coordinate names.
354 Returns
355 -------
356 str
357 If available, then return name of stamp coordinate.
358 Defaults to "realization" if alternative stamp coordinate not found.
359 """
360 # Acceptable stamp coordinate names
361 STAMP_COORD_NAMES = ["realization", "member", "sample", "pseudo_level"]
363 # Check which dimension coordinates we have.
364 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
366 # Check if any acceptable stamp coordinates are cube dimensions.
367 stamp_coords = [coord for coord in dim_coord_names if coord in STAMP_COORD_NAMES]
368 if len(stamp_coords) == 1:
369 stamp_coordinate = stamp_coords[0]
370 else:
371 stamp_coordinate = "realization"
373 return stamp_coordinate
376def fully_equalise_attributes(cubes: iris.cube.CubeList):
377 """Remove any unique attributes between cubes or coordinates in place."""
378 # Equalise cube attributes.
379 removed = iris.util.equalise_attributes(cubes)
380 logger.debug("Removed attributes from cube: %s", removed)
382 # Equalise coordinate attributes.
383 coord_sets = [{coord.name() for coord in cube.coords()} for cube in cubes]
385 all_coords = set.union(*coord_sets)
386 coords_to_equalise = set.intersection(*coord_sets)
387 coords_to_remove = set.difference(all_coords, coords_to_equalise)
389 logger.debug("All coordinates: %s", all_coords)
390 logger.debug("Coordinates to remove: %s", coords_to_remove)
391 logger.debug("Coordinates to equalise: %s", coords_to_equalise)
393 for coord in coords_to_remove:
394 for cube in cubes:
395 try:
396 cube.remove_coord(coord)
397 logger.debug("Removed coordinate %s from %s cube.", coord, cube.name())
398 except iris.exceptions.CoordinateNotFoundError:
399 pass
401 for coord in coords_to_equalise:
402 removed = iris.util.equalise_attributes([cube.coord(coord) for cube in cubes])
403 logger.debug("Removed attributes from coordinate %s: %s", coord, removed)
405 return cubes
408def slice_over_maybe(cube: iris.cube.Cube, coord_name, index):
409 """Test slicing over cube if exists.
411 Return None if not existing.
413 Arguments
414 ---------
415 cube: iris.cube.Cube
416 An iris cube which will be checked to see if it can be sliced over
417 given coordinate.
418 coord_name: coord
419 An iris coordinate over which to slice cube.
420 index:
421 Coordinate index value to extract
423 Returns
424 -------
425 cube_slice: iris.cube.Cube
426 A slice of iris cube, if available to slice.
427 """
428 if cube is None:
429 return None
431 # Check if coord exists as dimension coordinate
432 if not is_coorddim(cube, coord_name):
433 return cube
435 # Use iris to find which axis the dimension coordinate corresponds to
436 dim = cube.coord_dims(coord_name)[0]
438 # Create list of slices for each dimension
439 slices = [slice(None)] * cube.ndim
441 # Only replace the slice for the dim to be extracted
442 slices[dim] = index
444 return cube[tuple(slices)]
447def is_time_aggregatable(cube: iris.cube.Cube) -> bool:
448 """Determine whether a cube can be aggregated in time.
450 If a cube is aggregatable it will contain both a 'forecast_reference_time'
451 and 'forecast_period' coordinate as dimension or scalar coordinates.
453 Arguments
454 ---------
455 cube: iris.cube.Cube
456 An iris cube which will be checked to see if it is aggregatable based
457 on a set of pre-defined dimensional time coordinates:
458 'forecast_period' and 'forecast_reference_time'.
460 Returns
461 -------
462 bool
463 If true, then the cube is aggregatable and contains dimensional
464 coordinates including both 'forecast_reference_time' and
465 'forecast_period'.
466 """
467 # Acceptable time coordinate names for aggregatable cube.
468 TEMPORAL_COORD_NAMES = ["forecast_period", "forecast_reference_time"]
470 def strictly_monotonic(coord: iris.coords.Coord) -> bool:
471 """Return whether a coord is strictly monotonic, catching errors."""
472 try:
473 return coord.is_monotonic()
474 except iris.exceptions.CoordinateMultiDimError:
475 return False
477 # Strictly monotonic coordinate names for the cube.
478 coord_names = [coord.name() for coord in cube.coords() if strictly_monotonic(coord)]
480 # Check which temporal coordinates we have.
481 temporal_coords = [coord for coord in coord_names if coord in TEMPORAL_COORD_NAMES]
482 # Return whether both coordinates are in the temporal coordinates.
483 return len(temporal_coords) == 2
486def check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube:
487 """Ensure a single cube is given.
489 If a CubeList of length one is given that the contained cube is returned,
490 otherwise an error is raised.
492 Parameters
493 ----------
494 cube: Cube | CubeList
495 The cube to check.
497 Returns
498 -------
499 cube: Cube
500 The checked cube.
502 Raises
503 ------
504 TypeError
505 If the input cube is not a Cube or CubeList of a single Cube.
506 """
507 if isinstance(cube, iris.cube.Cube):
508 return cube
509 if isinstance(cube, iris.cube.CubeList):
510 if len(cube) == 1:
511 return cube[0]
512 else:
513 raise ValueError("CubeList did not contain a single cube.", cube)
514 raise TypeError(
515 "check_single_cube requires a Cube or CubeList of a single cube.", cube
516 )
519def check_sequence_coordinate(cubes, sequence_coordinate):
520 # If several histograms are plotted with time as sequence_coordinate for the
521 # time slider option.
522 for cube in iter_maybe(cubes):
523 try:
524 cube.coord(sequence_coordinate)
525 except iris.exceptions.CoordinateNotFoundError as err:
526 raise ValueError(
527 f"Cube must have a {sequence_coordinate} coordinate."
528 ) from err
531def get_num_models(cube: iris.cube.Cube | iris.cube.CubeList) -> int:
532 """Return number of models based on cube attributes."""
533 model_names = {cb.attributes.get("model_name") for cb in iter_maybe(cube)}
535 if not model_names: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true
536 logger.debug("Missing model names. Will assume single model.")
537 return 1
538 else:
539 return len(model_names)
542def validate_cube_shape(
543 cube: iris.cube.Cube | iris.cube.CubeList, num_models: int
544) -> None:
545 """Check all cubes have a model name."""
546 if isinstance(cube, iris.cube.CubeList) and len(cube) != num_models:
547 raise ValueError(
548 f"The number of model names ({num_models}) should equal the number "
549 f"of cubes ({len(cube)})."
550 )
553def validate_cubes_coords(
554 cubes: iris.cube.CubeList, coords: list[iris.coords.Coord]
555) -> None:
556 """Check same number of cubes as sequence coordinate for zip functions."""
557 if len(cubes) != len(coords):
558 raise ValueError(
559 f"The number of CubeList entries ({len(cubes)}) should equal the number "
560 f"of sequence coordinates ({len(coords)})."
561 f"Check that number of time entries in input data are consistent if "
562 f"performing time-averaging steps prior to plotting outputs."
563 )
566def check_if_cylc_workflow() -> Path | None:
567 """Determine if we are running in a Cylc workflow.
569 If running in a Cylc workflow, the ROSE_DATAC environment variable
570 will be set.
572 Returns
573 -------
574 Path | None:
575 If ROSE_DATAC is set, and the path exists, return a Path object
576 containing the path. Otherwise, return None.
577 """
578 # Standard location of ROSE_DATAC data dir in CSET.
579 try:
580 dataloc = Path(os.environ["ROSE_DATAC"])
581 if dataloc.exists():
582 return dataloc
583 except KeyError:
584 pass
586 # If ROSE_DATAC unset or its path does not exist, return None
587 return None
590def calc_array_stats(
591 array: np.ndarray | np.ma.MaskedArray,
592) -> tuple[float, float, float]:
593 """Calculate the min, max, and mean of an array.
595 NaNs/Masked data is ignored.
597 Returns
598 -------
599 stats:
600 A tuple of (min, max, mean).
601 """
602 if np.ma.isMaskedArray(array):
603 array_min = array.min()
604 array_max = array.max()
605 array_mean = array.mean()
606 else:
607 array_min = np.nanmin(array)
608 array_max = np.nanmax(array)
609 array_mean = np.nanmean(array)
610 return array_min, array_max, array_mean
613def get_time_coord(cube: Cube | CubeList) -> str | None:
614 """Return the time coordinate name or None if no time coordinate is found."""
615 cube = cube[0] if isinstance(cube, CubeList) else cube
616 for name in ("time", "hour"):
617 coords = cube.coords(name)
618 if coords and len(coords[0].points) > 1:
619 return name
620 return None
623def ensure_common_time_points(base: Cube, other: Cube):
624 """If a time coordinate is present, raise if base and other don't have the same time points."""
625 time_coord = get_time_coord(base)
626 if not time_coord:
627 return
628 shared_times = get_shared_times(
629 base, other, time_coord
630 ) # raises if no shared times
631 n_base = len(base.coord(time_coord).points)
632 n_other = len(other.coord(time_coord).points)
633 if not len(shared_times) == n_base == n_other: 633 ↛ 634line 633 didn't jump to line 634 because the condition on line 633 was never true
634 raise ValueError("Base and other cubes have different time points.")
637def get_shared_times(base: Cube, other: Cube, time_coord) -> set[int]:
638 """Extract common time points from cubes to allow comparison."""
639 # Get the name of the first non-scalar time coordinate.
640 base_time_coord = base.coord(time_coord)
641 other_time_coord = other.coord(time_coord)
642 logger.debug("Base: %s\nOther: %s", base_time_coord, other_time_coord)
643 if time_coord == "hour":
644 # We directly compare points when comparing coordinates with
645 # non-absolute units, such as hour. We can't just check the units are
646 # equal as iris automatically converts to datetime objects in the
647 # comparison for certain coordinate names.
648 base_times = base_time_coord.points
649 other_times = other_time_coord.points
650 shared_times = set.intersection(set(base_times), set(other_times))
651 else:
652 # Units don't match, so converting to datetimes for comparison.
653 base_times = base_time_coord.units.num2date(base_time_coord.points)
654 other_times = other_time_coord.units.num2date(other_time_coord.points)
655 shared_times = set.intersection(set(base_times), set(other_times))
657 logger.debug("Shared times: %s", shared_times)
659 if not shared_times:
660 raise ValueError("No common time points found!")
661 return shared_times