Coverage for src/CSET/operators/_utils.py: 100%
73 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-05 21:08 +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
24import iris
25import iris.cube
26import iris.exceptions
27import iris.util
30def get_cube_yxcoordname(cube: iris.cube.Cube) -> tuple[str, str]:
31 """
32 Return horizontal dimension coordinate name(s) from a given cube.
34 Arguments
35 ---------
37 cube: iris.cube.Cube
38 An iris cube which will be checked to see if it contains coordinate
39 names that match a pre-defined list of acceptable horizontal
40 dimension coordinate names.
42 Returns
43 -------
44 (y_coord, x_coord)
45 A tuple containing the horizontal coordinate name for latitude and longitude respectively
46 found within the cube.
48 Raises
49 ------
50 ValueError
51 If a unique y/x horizontal coordinate cannot be found.
52 """
53 # Acceptable horizontal coordinate names.
54 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
55 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
57 # Get a list of dimension coordinate names for the cube
58 dim_coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
59 coord_names = [coord.name() for coord in cube.coords()]
61 # Check which x-coordinate we have, if any
62 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
63 if len(x_coords) != 1:
64 x_coords = [coord for coord in dim_coord_names if coord in X_COORD_NAMES]
65 if len(x_coords) != 1:
66 raise ValueError("Could not identify a unique x-coordinate in cube")
68 # Check which y-coordinate we have, if any
69 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
70 if len(y_coords) != 1:
71 y_coords = [coord for coord in dim_coord_names if coord in Y_COORD_NAMES]
72 if len(y_coords) != 1:
73 raise ValueError("Could not identify a unique y-coordinate in cube")
75 return (y_coords[0], x_coords[0])
78def get_cube_coordindex(cube: iris.cube.Cube, coord_name) -> int:
79 """
80 Return coordinate dimension for a named coordinate from a given cube.
82 Arguments
83 ---------
85 cube: iris.cube.Cube
86 An iris cube which will be checked to see if it contains coordinate
87 names that match a pre-defined list of acceptable horizontal
88 coordinate names.
90 coord_name: str
91 A cube dimension name
93 Returns
94 -------
95 coord_index
96 An integer specifying where in the cube dimension list a specified coordinate name is found.
98 Raises
99 ------
100 ValueError
101 If a specified dimension coordinate cannot be found.
102 """
103 # Get a list of dimension coordinate names for the cube
104 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
106 # Check if requested dimension is found in cube and get index
107 if coord_name in coord_names:
108 coord_index = cube.coord_dims(coord_name)[0]
109 else:
110 raise ValueError("Could not find requested dimension %s", coord_name)
112 return coord_index
115def is_spatialdim(cube: iris.cube.Cube) -> bool:
116 """Determine whether a cube is has two spatial dimension coordinates.
118 If cube has both spatial dims, it will contain two unique coordinates
119 that explain space (latitude and longitude). The coordinates have to
120 be iterable/contain usable dimension data, as cubes may contain these
121 coordinates as scalar dimensions after being collapsed.
123 Arguments
124 ---------
125 cube: iris.cube.Cube
126 An iris cube which will be checked to see if it contains coordinate
127 names that match a pre-defined list of acceptable coordinate names.
129 Returns
130 -------
131 bool
132 If true, then the cube has a spatial projection and thus can be plotted
133 as a map.
134 """
135 # Acceptable horizontal coordinate names.
136 X_COORD_NAMES = ["longitude", "grid_longitude", "projection_x_coordinate", "x"]
137 Y_COORD_NAMES = ["latitude", "grid_latitude", "projection_y_coordinate", "y"]
139 # Get a list of coordinate names for the cube
140 coord_names = [coord.name() for coord in cube.dim_coords]
141 x_coords = [coord for coord in coord_names if coord in X_COORD_NAMES]
142 y_coords = [coord for coord in coord_names if coord in Y_COORD_NAMES]
144 # If there is one coordinate for both x and y direction return True.
145 if len(x_coords) == 1 and len(y_coords) == 1:
146 return True
147 else:
148 return False
151def is_transect(cube: iris.cube.Cube) -> bool:
152 """Determine whether a cube is a transect.
154 If cube is a transect, it will contain only one spatial (map) coordinate,
155 and one vertical coordinate (either pressure or model level).
157 Arguments
158 ---------
159 cube: iris.cube.Cube
160 An iris cube which will be checked to see if it contains coordinate
161 names that match a pre-defined list of acceptable coordinate names.
163 Returns
164 -------
165 bool
166 If true, then the cube is a transect that contains one spatial (map)
167 coordinate and one vertical coordinate.
168 """
169 # Acceptable spatial (map) coordinate names.
170 SPATIAL_MAP_COORD_NAMES = [
171 "longitude",
172 "grid_longitude",
173 "projection_x_coordinate",
174 "x",
175 "latitude",
176 "grid_latitude",
177 "projection_y_coordinate",
178 "y",
179 "distance",
180 ]
182 # Acceptable vertical coordinate names
183 VERTICAL_COORD_NAMES = ["pressure", "model_level_number", "level_height"]
185 # Get a list of coordinate names for the cube
186 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
188 # Check which spatial coordinates we have.
189 spatial_coords = [
190 coord for coord in coord_names if coord in SPATIAL_MAP_COORD_NAMES
191 ]
192 if len(spatial_coords) != 1:
193 return False
195 # Check which vertical coordinates we have.
196 vertical_coords = [coord for coord in coord_names if coord in VERTICAL_COORD_NAMES]
197 if len(vertical_coords) != 1:
198 return False
200 # Passed criteria so return True
201 return True
204def fully_equalise_attributes(cubes: iris.cube.CubeList):
205 """Remove any unique attributes between cubes or coordinates in place."""
206 # Equalise cube attributes.
207 removed = iris.util.equalise_attributes(cubes)
208 logging.debug("Removed attributes from cube: %s", removed)
210 # Equalise coordinate attributes.
211 coord_sets = [{coord.name() for coord in cube.coords()} for cube in cubes]
213 all_coords = set.union(*coord_sets)
214 coords_to_equalise = set.intersection(*coord_sets)
215 coords_to_remove = set.difference(all_coords, coords_to_equalise)
217 logging.debug("All coordinates: %s", all_coords)
218 logging.debug("Coordinates to remove: %s", coords_to_remove)
219 logging.debug("Coordinates to equalise: %s", coords_to_equalise)
221 for coord in coords_to_remove:
222 for cube in cubes:
223 try:
224 cube.remove_coord(coord)
225 logging.debug("Removed coordinate %s from %s cube.", coord, cube.name())
226 except iris.exceptions.CoordinateNotFoundError:
227 pass
229 for coord in coords_to_equalise:
230 removed = iris.util.equalise_attributes([cube.coord(coord) for cube in cubes])
231 logging.debug("Removed attributes from coordinate %s: %s", coord, removed)
233 return cubes
236def is_time_aggregatable(cube: iris.cube.Cube) -> bool:
237 """Determine whether a cube can be aggregated in time.
239 If a cube is aggregatable it will contain both a 'forecast_reference_time'
240 and 'forecast_period' coordinate as dimensional coordinates.
242 Arguments
243 ---------
244 cube: iris.cube.Cube
245 An iris cube which will be checked to see if it is aggregatable based
246 on a set of pre-defined dimensional time coordinates:
247 'forecast_period' and 'forecast_reference_time'.
249 Returns
250 -------
251 bool
252 If true, then the cube is aggregatable and contains dimensional
253 coordinates including both 'forecast_reference_time' and
254 'forecast_period'.
255 """
256 # Acceptable time coordinate names for aggregatable cube.
257 TEMPORAL_COORD_NAMES = ["forecast_period", "forecast_reference_time"]
259 # Coordinate names for the cube.
260 coord_names = [coord.name() for coord in cube.coords(dim_coords=True)]
262 # Check which temporal coordinates we have.
263 temporal_coords = [coord for coord in coord_names if coord in TEMPORAL_COORD_NAMES]
264 # Return whether both coordinates are in the temporal coordinates.
265 return len(temporal_coords) == 2