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