Coverage for src/CSET/operators/_utils.py: 95%
200 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-28 09:54 +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.time import PartialDateTime
36from CSET._common import iter_maybe
38logger = logging.getLogger(__name__)
41def pdt_fromisoformat(
42 datestring,
43) -> tuple[iris.time.PartialDateTime, timedelta | None]:
44 """Generate PartialDateTime object.
46 Function that takes an ISO 8601 date string and returns a PartialDateTime object.
48 Arguments
49 ---------
50 datestring: str
51 ISO 8601 date.
53 Returns
54 -------
55 time_object: iris.time.PartialDateTime
56 """
58 def make_offset(sign, value) -> timedelta:
59 if len(value) not in [2, 4, 5]: 59 ↛ 60line 59 didn't jump to line 60 because the condition on line 59 was never true
60 raise ValueError(f'expected "hh", "hhmm", or "hh:mm", got {value}')
62 hours = int(value[:2])
63 minutes = 0
64 if len(value) in [4, 5]:
65 minutes = int(value[-2:])
66 return timedelta(hours=sign * hours, minutes=sign * minutes)
68 # Remove the microseconds coord due to no support in PartialDateTime
69 datestring = re.sub(r"\.\d+", "", datestring)
71 datetime_split = datestring.split("T")
72 date = datetime_split[0]
73 if len(datetime_split) == 1:
74 time = ""
75 elif len(datetime_split) == 2: 75 ↛ 78line 75 didn't jump to line 78 because the condition on line 75 was always true
76 time = datetime_split[1]
77 else:
78 raise ValueError("datesting in an unexpected format")
80 offset = None
81 time_split = time.split("+")
82 if len(time_split) == 2:
83 time = time_split[0]
84 offset = make_offset(1, time_split[1])
85 else:
86 time_split = time.split("-")
87 if len(time_split) == 2: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true
88 time = time_split[0]
89 offset = make_offset(-1, time_split[1])
90 else:
91 offset = None
93 if re.fullmatch(r"\d{8}", date):
94 date = f"{date[0:4]}-{date[4:6]}-{date[6:8]}"
95 elif re.fullmatch(r"\d{6}", date):
96 date = f"{date[0:4]}-{date[4:6]}"
98 if len(date) < 7: 98 ↛ 99line 98 didn't jump to line 99 because the condition on line 98 was never true
99 raise ValueError(f"Invalid datestring: {datestring}, must be at least YYYY-MM")
101 # Returning a PartialDateTime for the special case of string form "YYYY-MM"
102 if re.fullmatch(r"\d{4}-\d{2}", date):
103 pdt = PartialDateTime(
104 year=int(date[0:4]),
105 month=int(date[5:7]),
106 day=None,
107 hour=None,
108 minute=None,
109 second=None,
110 )
111 return pdt, offset
113 year = int(date[0:4])
114 month = int(date[5:7])
115 day = int(date[8:10])
117 kwargs = {
118 "year": year,
119 "month": month,
120 "day": day,
121 "hour": 0,
122 "minute": 0,
123 "second": 0,
124 }
126 # Normalise the time parts into standard format
127 if re.fullmatch(r"\d{4}", time): 127 ↛ 128line 127 didn't jump to line 128 because the condition on line 127 was never true
128 time = f"{time[0:2]}:{time[2:4]}"
129 if re.fullmatch(r"\d{6}", time):
130 time = f"{time[0:2]}:{time[2:4]}:{time[4:6]}"
132 if len(time) >= 2:
133 kwargs["hour"] = int(time[0:2])
134 if len(time) >= 5:
135 kwargs["minute"] = int(time[3:5])
136 if len(time) >= 8:
137 kwargs["second"] = int(time[6:8])
139 pdt = PartialDateTime(**kwargs)
141 return pdt, offset
144def get_cube_yxcoordname(cube: iris.cube.Cube) -> tuple[str, str]:
145 """
146 Return horizontal dimension coordinate name(s) from a given cube.
148 Arguments
149 ---------
151 cube: iris.cube.Cube
152 An iris cube which will be checked to see if it contains coordinate
153 names that match a pre-defined list of acceptable horizontal
154 dimension coordinate names.
156 Returns
157 -------
158 (y_coord, x_coord)
159 A tuple containing the horizontal coordinate name for latitude and longitude respectively
160 found within the cube.
162 Raises
163 ------
164 ValueError
165 If a unique y/x horizontal coordinate cannot be found.
166 """
167 # Acceptable horizontal coordinate names.
168 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
169 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
171 # Get a list of dimension coordinate names for the cube
172 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
173 coord_names = [coord.name() for coord in cube.coords()]
175 # Check which x-coordinate we have, if any
176 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
177 if len(x_coords) != 1:
178 x_coords = [coord for coord in dim_coord_names if coord in X_COORD_NAMES]
179 if len(x_coords) != 1:
180 raise ValueError("Could not identify a unique x-coordinate in cube")
182 # Check which y-coordinate we have, if any
183 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
184 if len(y_coords) != 1:
185 y_coords = [coord for coord in dim_coord_names if coord in Y_COORD_NAMES]
186 if len(y_coords) != 1:
187 raise ValueError("Could not identify a unique y-coordinate in cube")
189 return (y_coords[0], x_coords[0])
192def get_cube_coordindex(cube: iris.cube.Cube, coord_name) -> int:
193 """
194 Return coordinate dimension for a named coordinate from a given cube.
196 Arguments
197 ---------
199 cube: iris.cube.Cube
200 An iris cube which will be checked to see if it contains coordinate
201 names that match a pre-defined list of acceptable horizontal
202 coordinate names.
204 coord_name: str
205 A cube dimension name
207 Returns
208 -------
209 coord_index
210 An integer specifying where in the cube dimension list a specified coordinate name is found.
212 Raises
213 ------
214 ValueError
215 If a specified dimension coordinate cannot be found.
216 """
217 # Get a list of dimension coordinate names for the cube
218 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
220 # Check if requested dimension is found in cube and get index
221 if coord_name in coord_names:
222 coord_index = cube.coord_dims(coord_name)[0]
223 else:
224 raise ValueError("Could not find requested dimension %s", coord_name)
226 return coord_index
229def is_spatialdim(cube: iris.cube.Cube) -> bool:
230 """Determine whether a cube is has two spatial dimension coordinates.
232 If cube has both spatial dims, it will contain two unique coordinates
233 that explain space (latitude and longitude). The coordinates have to
234 be iterable/contain usable dimension data, as cubes may contain these
235 coordinates as scalar dimensions after being collapsed.
237 Arguments
238 ---------
239 cube: iris.cube.Cube
240 An iris cube which will be checked to see if it contains coordinate
241 names that match a pre-defined list of acceptable coordinate names.
243 Returns
244 -------
245 bool
246 If true, then the cube has a spatial projection and thus can be plotted
247 as a map.
248 """
249 # Acceptable horizontal coordinate names.
250 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
251 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
253 # Get a list of coordinate names for the cube
254 coord_names = [coord.name() for coord in cube.dim_coords]
255 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
256 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
258 # If there is one coordinate for both x and y direction return True.
259 return len(x_coords) == 1 and len(y_coords) == 1
262def is_coorddim(cube: iris.cube.Cube, coord_name) -> bool:
263 """Determine whether a cube has specified dimension coordinates.
265 Arguments
266 ---------
267 cube: iris.cube.Cube
268 An iris cube which will be checked to see if it contains coordinate
269 names that match a pre-defined list of acceptable coordinate names.
271 coord_name: str
272 A cube dimension name
274 Returns
275 -------
276 bool
277 If true, then the cube has a spatial projection and thus can be plotted
278 as a map.
279 """
280 # Get a list of dimension coordinate names for the cube
281 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
283 # Check if requested dimension is found in cube and get index
284 return coord_name in coord_names
287def is_transect(cube: iris.cube.Cube) -> bool:
288 """Determine whether a cube is a transect.
290 If cube is a transect, it will contain only one spatial (map) coordinate,
291 and one vertical coordinate (either pressure or model level).
293 Arguments
294 ---------
295 cube: iris.cube.Cube
296 An iris cube which will be checked to see if it contains coordinate
297 names that match a pre-defined list of acceptable coordinate names.
299 Returns
300 -------
301 bool
302 If true, then the cube is a transect that contains one spatial (map)
303 coordinate and one vertical coordinate.
304 """
305 # Acceptable spatial (map) coordinate names.
306 SPATIAL_MAP_COORD_NAMES = [
307 "longitude",
308 "grid_longitude",
309 "projection_x_coordinate",
310 "x",
311 "latitude",
312 "grid_latitude",
313 "projection_y_coordinate",
314 "y",
315 "distance",
316 ]
318 # Acceptable vertical coordinate names
319 VERTICAL_COORD_NAMES = ["pressure", "model_level_number", "level_height"]
321 # Get a list of coordinate names for the cube
322 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
324 # Check which spatial coordinates we have.
325 spatial_coords = [
326 coord for coord in coord_names if coord in SPATIAL_MAP_COORD_NAMES
327 ]
328 if len(spatial_coords) != 1:
329 return False
331 # Check which vertical coordinates we have.
332 vertical_coords = [coord for coord in coord_names if coord in VERTICAL_COORD_NAMES]
333 if len(vertical_coords) != 1: # noqa: SIM103 Clearer to keep separate.
334 return False
336 # Passed criteria so return True
337 return True
340def check_stamp_coordinate(cube: iris.cube.Cube) -> str:
341 """
342 Return stamp dimension coordinate name from a given cube, if exists.
344 If cube contains a valid stamp coordinate as a dimension coordinate,
345 function will return name of the stamp coordinate.
347 Arguments
348 ---------
349 cube: iris.cube.Cube
350 An iris cube which will be checked to see if it contains coordinate
351 names that match a pre-defined list of acceptable coordinate names.
353 Returns
354 -------
355 str
356 If available, then return name of stamp coordinate.
357 Defaults to "realization" if alternative stamp coordinate not found.
358 """
359 # Acceptable stamp coordinate names
360 STAMP_COORD_NAMES = ["realization", "member", "sample", "pseudo_level"]
362 # Check which dimension coordinates we have.
363 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
365 # Check if any acceptable stamp coordinates are cube dimensions.
366 stamp_coords = [coord for coord in dim_coord_names if coord in STAMP_COORD_NAMES]
367 if len(stamp_coords) == 1:
368 stamp_coordinate = stamp_coords[0]
369 else:
370 stamp_coordinate = "realization"
372 return stamp_coordinate
375def fully_equalise_attributes(cubes: iris.cube.CubeList):
376 """Remove any unique attributes between cubes or coordinates in place."""
377 # Equalise cube attributes.
378 removed = iris.util.equalise_attributes(cubes)
379 logger.debug("Removed attributes from cube: %s", removed)
381 # Equalise coordinate attributes.
382 coord_sets = [{coord.name() for coord in cube.coords()} for cube in cubes]
384 all_coords = set.union(*coord_sets)
385 coords_to_equalise = set.intersection(*coord_sets)
386 coords_to_remove = set.difference(all_coords, coords_to_equalise)
388 logger.debug("All coordinates: %s", all_coords)
389 logger.debug("Coordinates to remove: %s", coords_to_remove)
390 logger.debug("Coordinates to equalise: %s", coords_to_equalise)
392 for coord in coords_to_remove:
393 for cube in cubes:
394 try:
395 cube.remove_coord(coord)
396 logger.debug("Removed coordinate %s from %s cube.", coord, cube.name())
397 except iris.exceptions.CoordinateNotFoundError:
398 pass
400 for coord in coords_to_equalise:
401 removed = iris.util.equalise_attributes([cube.coord(coord) for cube in cubes])
402 logger.debug("Removed attributes from coordinate %s: %s", coord, removed)
404 return cubes
407def slice_over_maybe(cube: iris.cube.Cube, coord_name, index):
408 """Test slicing over cube if exists.
410 Return None if not existing.
412 Arguments
413 ---------
414 cube: iris.cube.Cube
415 An iris cube which will be checked to see if it can be sliced over
416 given coordinate.
417 coord_name: coord
418 An iris coordinate over which to slice cube.
419 index:
420 Coordinate index value to extract
422 Returns
423 -------
424 cube_slice: iris.cube.Cube
425 A slice of iris cube, if available to slice.
426 """
427 if cube is None:
428 return None
430 # Check if coord exists as dimension coordinate
431 if not is_coorddim(cube, coord_name):
432 return cube
434 # Use iris to find which axis the dimension coordinate corresponds to
435 dim = cube.coord_dims(coord_name)[0]
437 # Create list of slices for each dimension
438 slices = [slice(None)] * cube.ndim
440 # Only replace the slice for the dim to be extracted
441 slices[dim] = index
443 return cube[tuple(slices)]
446def is_time_aggregatable(cube: iris.cube.Cube) -> bool:
447 """Determine whether a cube can be aggregated in time.
449 If a cube is aggregatable it will contain both a 'forecast_reference_time'
450 and 'forecast_period' coordinate as dimension or scalar coordinates.
452 Arguments
453 ---------
454 cube: iris.cube.Cube
455 An iris cube which will be checked to see if it is aggregatable based
456 on a set of pre-defined dimensional time coordinates:
457 'forecast_period' and 'forecast_reference_time'.
459 Returns
460 -------
461 bool
462 If true, then the cube is aggregatable and contains dimensional
463 coordinates including both 'forecast_reference_time' and
464 'forecast_period'.
465 """
466 # Acceptable time coordinate names for aggregatable cube.
467 TEMPORAL_COORD_NAMES = ["forecast_period", "forecast_reference_time"]
469 def strictly_monotonic(coord: iris.coords.Coord) -> bool:
470 """Return whether a coord is strictly monotonic, catching errors."""
471 try:
472 return coord.is_monotonic()
473 except iris.exceptions.CoordinateMultiDimError:
474 return False
476 # Strictly monotonic coordinate names for the cube.
477 coord_names = [coord.name() for coord in cube.coords() if strictly_monotonic(coord)]
479 # Check which temporal coordinates we have.
480 temporal_coords = [coord for coord in coord_names if coord in TEMPORAL_COORD_NAMES]
481 # Return whether both coordinates are in the temporal coordinates.
482 return len(temporal_coords) == 2
485def check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube:
486 """Ensure a single cube is given.
488 If a CubeList of length one is given that the contained cube is returned,
489 otherwise an error is raised.
491 Parameters
492 ----------
493 cube: Cube | CubeList
494 The cube to check.
496 Returns
497 -------
498 cube: Cube
499 The checked cube.
501 Raises
502 ------
503 TypeError
504 If the input cube is not a Cube or CubeList of a single Cube.
505 """
506 if isinstance(cube, iris.cube.Cube):
507 return cube
508 if isinstance(cube, iris.cube.CubeList):
509 if len(cube) == 1:
510 return cube[0]
511 else:
512 raise ValueError("CubeList did not contain a single cube.", cube)
513 raise TypeError(
514 "check_single_cube requires a Cube or CubeList of a single cube.", cube
515 )
518def check_sequence_coordinate(cubes, sequence_coordinate):
519 # If several histograms are plotted with time as sequence_coordinate for the
520 # time slider option.
521 for cube in iter_maybe(cubes):
522 try:
523 cube.coord(sequence_coordinate)
524 except iris.exceptions.CoordinateNotFoundError as err:
525 raise ValueError(
526 f"Cube must have a {sequence_coordinate} coordinate."
527 ) from err
530def get_num_models(cube: iris.cube.Cube | iris.cube.CubeList) -> int:
531 """Return number of models based on cube attributes."""
532 model_names = {cb.attributes.get("model_name") for cb in iter_maybe(cube)}
534 if not model_names: 534 ↛ 535line 534 didn't jump to line 535 because the condition on line 534 was never true
535 logger.debug("Missing model names. Will assume single model.")
536 return 1
537 else:
538 return len(model_names)
541def validate_cube_shape(
542 cube: iris.cube.Cube | iris.cube.CubeList, num_models: int
543) -> None:
544 """Check all cubes have a model name."""
545 if isinstance(cube, iris.cube.CubeList) and len(cube) != num_models:
546 raise ValueError(
547 f"The number of model names ({num_models}) should equal the number "
548 f"of cubes ({len(cube)})."
549 )
552def validate_cubes_coords(
553 cubes: iris.cube.CubeList, coords: list[iris.coords.Coord]
554) -> None:
555 """Check same number of cubes as sequence coordinate for zip functions."""
556 if len(cubes) != len(coords):
557 raise ValueError(
558 f"The number of CubeList entries ({len(cubes)}) should equal the number "
559 f"of sequence coordinates ({len(coords)})."
560 f"Check that number of time entries in input data are consistent if "
561 f"performing time-averaging steps prior to plotting outputs."
562 )
565def check_if_cylc_workflow() -> Path | None:
566 """Determine if we are running in a Cylc workflow.
568 If running in a Cylc workflow, the ROSE_DATAC environment variable
569 will be set.
571 Returns
572 -------
573 Path | None:
574 If ROSE_DATAC is set, and the path exists, return a Path object
575 containing the path. Otherwise, return None.
576 """
577 # Standard location of ROSE_DATAC data dir in CSET.
578 try:
579 dataloc = Path(os.environ["ROSE_DATAC"])
580 if dataloc.exists():
581 return dataloc
582 except KeyError:
583 pass
585 # If ROSE_DATAC unset or its path does not exist, return None
586 return None
589def calc_array_stats(
590 array: np.ndarray | np.ma.MaskedArray,
591) -> tuple[float, float, float]:
592 """Calculate the min, max, and mean of an array.
594 NaNs/Masked data is ignored.
596 Returns
597 -------
598 stats:
599 A tuple of (min, max, mean).
600 """
601 if np.ma.isMaskedArray(array):
602 array_min = array.min()
603 array_max = array.max()
604 array_mean = array.mean()
605 else:
606 array_min = np.nanmin(array)
607 array_max = np.nanmax(array)
608 array_mean = np.nanmean(array)
609 return array_min, array_max, array_mean