Coverage for src/CSET/operators/_utils.py: 95%
190 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 08:32 +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
33from iris.time import PartialDateTime
35from CSET._common import iter_maybe
37logger = logging.getLogger(__name__)
40def pdt_fromisoformat(
41 datestring,
42) -> tuple[iris.time.PartialDateTime, timedelta | None]:
43 """Generate PartialDateTime object.
45 Function that takes an ISO 8601 date string and returns a PartialDateTime object.
47 Arguments
48 ---------
49 datestring: str
50 ISO 8601 date.
52 Returns
53 -------
54 time_object: iris.time.PartialDateTime
55 """
57 def make_offset(sign, value) -> timedelta:
58 if len(value) not in [2, 4, 5]: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true
59 raise ValueError(f'expected "hh", "hhmm", or "hh:mm", got {value}')
61 hours = int(value[:2])
62 minutes = 0
63 if len(value) in [4, 5]:
64 minutes = int(value[-2:])
65 return timedelta(hours=sign * hours, minutes=sign * minutes)
67 # Remove the microseconds coord due to no support in PartialDateTime
68 datestring = re.sub(r"\.\d+", "", datestring)
70 datetime_split = datestring.split("T")
71 date = datetime_split[0]
72 if len(datetime_split) == 1:
73 time = ""
74 elif len(datetime_split) == 2: 74 ↛ 77line 74 didn't jump to line 77 because the condition on line 74 was always true
75 time = datetime_split[1]
76 else:
77 raise ValueError("datesting in an unexpected format")
79 offset = None
80 time_split = time.split("+")
81 if len(time_split) == 2:
82 time = time_split[0]
83 offset = make_offset(1, time_split[1])
84 else:
85 time_split = time.split("-")
86 if len(time_split) == 2: 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true
87 time = time_split[0]
88 offset = make_offset(-1, time_split[1])
89 else:
90 offset = None
92 if re.fullmatch(r"\d{8}", date):
93 date = f"{date[0:4]}-{date[4:6]}-{date[6:8]}"
94 elif re.fullmatch(r"\d{6}", date):
95 date = f"{date[0:4]}-{date[4:6]}"
97 if len(date) < 7: 97 ↛ 98line 97 didn't jump to line 98 because the condition on line 97 was never true
98 raise ValueError(f"Invalid datestring: {datestring}, must be at least YYYY-MM")
100 # Returning a PartialDateTime for the special case of string form "YYYY-MM"
101 if re.fullmatch(r"\d{4}-\d{2}", date):
102 pdt = PartialDateTime(
103 year=int(date[0:4]),
104 month=int(date[5:7]),
105 day=None,
106 hour=None,
107 minute=None,
108 second=None,
109 )
110 return pdt, offset
112 year = int(date[0:4])
113 month = int(date[5:7])
114 day = int(date[8:10])
116 kwargs = {
117 "year": year,
118 "month": month,
119 "day": day,
120 "hour": 0,
121 "minute": 0,
122 "second": 0,
123 }
125 # Normalise the time parts into standard format
126 if re.fullmatch(r"\d{4}", time): 126 ↛ 127line 126 didn't jump to line 127 because the condition on line 126 was never true
127 time = f"{time[0:2]}:{time[2:4]}"
128 if re.fullmatch(r"\d{6}", time):
129 time = f"{time[0:2]}:{time[2:4]}:{time[4:6]}"
131 if len(time) >= 2:
132 kwargs["hour"] = int(time[0:2])
133 if len(time) >= 5:
134 kwargs["minute"] = int(time[3:5])
135 if len(time) >= 8:
136 kwargs["second"] = int(time[6:8])
138 pdt = PartialDateTime(**kwargs)
140 return pdt, offset
143def get_cube_yxcoordname(cube: iris.cube.Cube) -> tuple[str, str]:
144 """
145 Return horizontal dimension coordinate name(s) from a given cube.
147 Arguments
148 ---------
150 cube: iris.cube.Cube
151 An iris cube which will be checked to see if it contains coordinate
152 names that match a pre-defined list of acceptable horizontal
153 dimension coordinate names.
155 Returns
156 -------
157 (y_coord, x_coord)
158 A tuple containing the horizontal coordinate name for latitude and longitude respectively
159 found within the cube.
161 Raises
162 ------
163 ValueError
164 If a unique y/x horizontal coordinate cannot be found.
165 """
166 # Acceptable horizontal coordinate names.
167 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
168 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
170 # Get a list of dimension coordinate names for the cube
171 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
172 coord_names = [coord.name() for coord in cube.coords()]
174 # Check which x-coordinate we have, if any
175 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
176 if len(x_coords) != 1:
177 x_coords = [coord for coord in dim_coord_names if coord in X_COORD_NAMES]
178 if len(x_coords) != 1:
179 raise ValueError("Could not identify a unique x-coordinate in cube")
181 # Check which y-coordinate we have, if any
182 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
183 if len(y_coords) != 1:
184 y_coords = [coord for coord in dim_coord_names if coord in Y_COORD_NAMES]
185 if len(y_coords) != 1:
186 raise ValueError("Could not identify a unique y-coordinate in cube")
188 return (y_coords[0], x_coords[0])
191def get_cube_coordindex(cube: iris.cube.Cube, coord_name) -> int:
192 """
193 Return coordinate dimension for a named coordinate from a given cube.
195 Arguments
196 ---------
198 cube: iris.cube.Cube
199 An iris cube which will be checked to see if it contains coordinate
200 names that match a pre-defined list of acceptable horizontal
201 coordinate names.
203 coord_name: str
204 A cube dimension name
206 Returns
207 -------
208 coord_index
209 An integer specifying where in the cube dimension list a specified coordinate name is found.
211 Raises
212 ------
213 ValueError
214 If a specified dimension coordinate cannot be found.
215 """
216 # Get a list of dimension coordinate names for the cube
217 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
219 # Check if requested dimension is found in cube and get index
220 if coord_name in coord_names:
221 coord_index = cube.coord_dims(coord_name)[0]
222 else:
223 raise ValueError("Could not find requested dimension %s", coord_name)
225 return coord_index
228def is_spatialdim(cube: iris.cube.Cube) -> bool:
229 """Determine whether a cube is has two spatial dimension coordinates.
231 If cube has both spatial dims, it will contain two unique coordinates
232 that explain space (latitude and longitude). The coordinates have to
233 be iterable/contain usable dimension data, as cubes may contain these
234 coordinates as scalar dimensions after being collapsed.
236 Arguments
237 ---------
238 cube: iris.cube.Cube
239 An iris cube which will be checked to see if it contains coordinate
240 names that match a pre-defined list of acceptable coordinate names.
242 Returns
243 -------
244 bool
245 If true, then the cube has a spatial projection and thus can be plotted
246 as a map.
247 """
248 # Acceptable horizontal coordinate names.
249 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
250 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
252 # Get a list of coordinate names for the cube
253 coord_names = [coord.name() for coord in cube.dim_coords]
254 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
255 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
257 # If there is one coordinate for both x and y direction return True.
258 return len(x_coords) == 1 and len(y_coords) == 1
261def is_coorddim(cube: iris.cube.Cube, coord_name) -> bool:
262 """Determine whether a cube has specified dimension coordinates.
264 Arguments
265 ---------
266 cube: iris.cube.Cube
267 An iris cube which will be checked to see if it contains coordinate
268 names that match a pre-defined list of acceptable coordinate names.
270 coord_name: str
271 A cube dimension name
273 Returns
274 -------
275 bool
276 If true, then the cube has a spatial projection and thus can be plotted
277 as a map.
278 """
279 # Get a list of dimension coordinate names for the cube
280 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
282 # Check if requested dimension is found in cube and get index
283 return coord_name in coord_names
286def is_transect(cube: iris.cube.Cube) -> bool:
287 """Determine whether a cube is a transect.
289 If cube is a transect, it will contain only one spatial (map) coordinate,
290 and one vertical coordinate (either pressure or model level).
292 Arguments
293 ---------
294 cube: iris.cube.Cube
295 An iris cube which will be checked to see if it contains coordinate
296 names that match a pre-defined list of acceptable coordinate names.
298 Returns
299 -------
300 bool
301 If true, then the cube is a transect that contains one spatial (map)
302 coordinate and one vertical coordinate.
303 """
304 # Acceptable spatial (map) coordinate names.
305 SPATIAL_MAP_COORD_NAMES = [
306 "longitude",
307 "grid_longitude",
308 "projection_x_coordinate",
309 "x",
310 "latitude",
311 "grid_latitude",
312 "projection_y_coordinate",
313 "y",
314 "distance",
315 ]
317 # Acceptable vertical coordinate names
318 VERTICAL_COORD_NAMES = ["pressure", "model_level_number", "level_height"]
320 # Get a list of coordinate names for the cube
321 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
323 # Check which spatial coordinates we have.
324 spatial_coords = [
325 coord for coord in coord_names if coord in SPATIAL_MAP_COORD_NAMES
326 ]
327 if len(spatial_coords) != 1:
328 return False
330 # Check which vertical coordinates we have.
331 vertical_coords = [coord for coord in coord_names if coord in VERTICAL_COORD_NAMES]
332 if len(vertical_coords) != 1: # noqa: SIM103 Clearer to keep separate.
333 return False
335 # Passed criteria so return True
336 return True
339def check_stamp_coordinate(cube: iris.cube.Cube) -> str:
340 """
341 Return stamp dimension coordinate name from a given cube, if exists.
343 If cube contains a valid stamp coordinate as a dimension coordinate,
344 function will return name of the stamp coordinate.
346 Arguments
347 ---------
348 cube: iris.cube.Cube
349 An iris cube which will be checked to see if it contains coordinate
350 names that match a pre-defined list of acceptable coordinate names.
352 Returns
353 -------
354 str
355 If available, then return name of stamp coordinate.
356 Defaults to "realization" if alternative stamp coordinate not found.
357 """
358 # Acceptable stamp coordinate names
359 STAMP_COORD_NAMES = ["realization", "member", "sample", "pseudo_level"]
361 # Check which dimension coordinates we have.
362 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
364 # Check if any acceptable stamp coordinates are cube dimensions.
365 stamp_coords = [coord for coord in dim_coord_names if coord in STAMP_COORD_NAMES]
366 if len(stamp_coords) == 1:
367 stamp_coordinate = stamp_coords[0]
368 else:
369 stamp_coordinate = "realization"
371 return stamp_coordinate
374def fully_equalise_attributes(cubes: iris.cube.CubeList):
375 """Remove any unique attributes between cubes or coordinates in place."""
376 # Equalise cube attributes.
377 removed = iris.util.equalise_attributes(cubes)
378 logger.debug("Removed attributes from cube: %s", removed)
380 # Equalise coordinate attributes.
381 coord_sets = [{coord.name() for coord in cube.coords()} for cube in cubes]
383 all_coords = set.union(*coord_sets)
384 coords_to_equalise = set.intersection(*coord_sets)
385 coords_to_remove = set.difference(all_coords, coords_to_equalise)
387 logger.debug("All coordinates: %s", all_coords)
388 logger.debug("Coordinates to remove: %s", coords_to_remove)
389 logger.debug("Coordinates to equalise: %s", coords_to_equalise)
391 for coord in coords_to_remove:
392 for cube in cubes:
393 try:
394 cube.remove_coord(coord)
395 logger.debug("Removed coordinate %s from %s cube.", coord, cube.name())
396 except iris.exceptions.CoordinateNotFoundError:
397 pass
399 for coord in coords_to_equalise:
400 removed = iris.util.equalise_attributes([cube.coord(coord) for cube in cubes])
401 logger.debug("Removed attributes from coordinate %s: %s", coord, removed)
403 return cubes
406def slice_over_maybe(cube: iris.cube.Cube, coord_name, index):
407 """Test slicing over cube if exists.
409 Return None if not existing.
411 Arguments
412 ---------
413 cube: iris.cube.Cube
414 An iris cube which will be checked to see if it can be sliced over
415 given coordinate.
416 coord_name: coord
417 An iris coordinate over which to slice cube.
418 index:
419 Coordinate index value to extract
421 Returns
422 -------
423 cube_slice: iris.cube.Cube
424 A slice of iris cube, if available to slice.
425 """
426 if cube is None:
427 return None
429 # Check if coord exists as dimension coordinate
430 if not is_coorddim(cube, coord_name):
431 return cube
433 # Use iris to find which axis the dimension coordinate corresponds to
434 dim = cube.coord_dims(coord_name)[0]
436 # Create list of slices for each dimension
437 slices = [slice(None)] * cube.ndim
439 # Only replace the slice for the dim to be extracted
440 slices[dim] = index
442 return cube[tuple(slices)]
445def is_time_aggregatable(cube: iris.cube.Cube) -> bool:
446 """Determine whether a cube can be aggregated in time.
448 If a cube is aggregatable it will contain both a 'forecast_reference_time'
449 and 'forecast_period' coordinate as dimension or scalar coordinates.
451 Arguments
452 ---------
453 cube: iris.cube.Cube
454 An iris cube which will be checked to see if it is aggregatable based
455 on a set of pre-defined dimensional time coordinates:
456 'forecast_period' and 'forecast_reference_time'.
458 Returns
459 -------
460 bool
461 If true, then the cube is aggregatable and contains dimensional
462 coordinates including both 'forecast_reference_time' and
463 'forecast_period'.
464 """
465 # Acceptable time coordinate names for aggregatable cube.
466 TEMPORAL_COORD_NAMES = ["forecast_period", "forecast_reference_time"]
468 def strictly_monotonic(coord: iris.coords.Coord) -> bool:
469 """Return whether a coord is strictly monotonic, catching errors."""
470 try:
471 return coord.is_monotonic()
472 except iris.exceptions.CoordinateMultiDimError:
473 return False
475 # Strictly monotonic coordinate names for the cube.
476 coord_names = [coord.name() for coord in cube.coords() if strictly_monotonic(coord)]
478 # Check which temporal coordinates we have.
479 temporal_coords = [coord for coord in coord_names if coord in TEMPORAL_COORD_NAMES]
480 # Return whether both coordinates are in the temporal coordinates.
481 return len(temporal_coords) == 2
484def check_single_cube(cube: iris.cube.Cube | iris.cube.CubeList) -> iris.cube.Cube:
485 """Ensure a single cube is given.
487 If a CubeList of length one is given that the contained cube is returned,
488 otherwise an error is raised.
490 Parameters
491 ----------
492 cube: Cube | CubeList
493 The cube to check.
495 Returns
496 -------
497 cube: Cube
498 The checked cube.
500 Raises
501 ------
502 TypeError
503 If the input cube is not a Cube or CubeList of a single Cube.
504 """
505 if isinstance(cube, iris.cube.Cube):
506 return cube
507 if isinstance(cube, iris.cube.CubeList):
508 if len(cube) == 1:
509 return cube[0]
510 else:
511 raise ValueError("CubeList did not contain a single cube.", cube)
512 raise TypeError(
513 "check_single_cube requires a Cube or CubeList of a single cube.", cube
514 )
517def check_sequence_coordinate(cubes, sequence_coordinate):
518 # If several histograms are plotted with time as sequence_coordinate for the
519 # time slider option.
520 for cube in iter_maybe(cubes):
521 try:
522 cube.coord(sequence_coordinate)
523 except iris.exceptions.CoordinateNotFoundError as err:
524 raise ValueError(
525 f"Cube must have a {sequence_coordinate} coordinate."
526 ) from err
529def get_num_models(cube: iris.cube.Cube | iris.cube.CubeList) -> int:
530 """Return number of models based on cube attributes."""
531 model_names = {cb.attributes.get("model_name") for cb in iter_maybe(cube)}
533 if not model_names: 533 ↛ 534line 533 didn't jump to line 534 because the condition on line 533 was never true
534 logger.debug("Missing model names. Will assume single model.")
535 return 1
536 else:
537 return len(model_names)
540def validate_cube_shape(
541 cube: iris.cube.Cube | iris.cube.CubeList, num_models: int
542) -> None:
543 """Check all cubes have a model name."""
544 if isinstance(cube, iris.cube.CubeList) and len(cube) != num_models:
545 raise ValueError(
546 f"The number of model names ({num_models}) should equal the number "
547 f"of cubes ({len(cube)})."
548 )
551def validate_cubes_coords(
552 cubes: iris.cube.CubeList, coords: list[iris.coords.Coord]
553) -> None:
554 """Check same number of cubes as sequence coordinate for zip functions."""
555 if len(cubes) != len(coords):
556 raise ValueError(
557 f"The number of CubeList entries ({len(cubes)}) should equal the number "
558 f"of sequence coordinates ({len(coords)})."
559 f"Check that number of time entries in input data are consistent if "
560 f"performing time-averaging steps prior to plotting outputs."
561 )
564def check_if_cylc_workflow() -> Path | None:
565 """Determine if we are running in a Cylc workflow.
567 If running in a Cylc workflow, the ROSE_DATAC environment variable
568 will be set.
570 Returns
571 -------
572 Path | None:
573 If ROSE_DATAC is set, and the path exists, return a Path object
574 containing the path. Otherwise, return None.
575 """
576 # Standard location of ROSE_DATAC data dir in CSET.
577 try:
578 dataloc = Path(os.environ["ROSE_DATAC"])
579 if dataloc.exists():
580 return dataloc
581 except KeyError:
582 pass
584 # If ROSE_DATAC unset or its path does not exist, return None
585 return None