Coverage for src/CSET/operators/collapse.py: 93%
155 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"""Operators to perform various kind of collapse on either 1 or 2 dimensions."""
17import datetime
18import logging
19import warnings
21import iris
22import iris.analysis
23import iris.coord_categorisation
24import iris.coords
25import iris.cube
26import iris.exceptions
27import iris.util
28import numpy as np
30from CSET._common import iter_maybe
31from CSET.operators.aggregate import add_hour_coordinate
33logger = logging.getLogger(__name__)
36def collapse(
37 cubes: iris.cube.Cube | iris.cube.CubeList,
38 coordinate: str | list[str],
39 method: str,
40 additional_percent: float | None = None,
41 **kwargs,
42) -> iris.cube.Cube | iris.cube.CubeList:
43 """Collapse coordinate(s) of a single cube or of every cube in a cube list.
45 Collapses similar fields in each cube into a cube collapsing around the
46 specified coordinate(s) and method. This could be a (weighted) mean or
47 percentile.
49 Arguments
50 ---------
51 cubes: iris.cube.Cube | iris.cube.CubeList
52 Cube or CubeList to collapse and iterate over one dimension
53 coordinate: str | list[str]
54 Coordinate(s) to collapse over e.g. 'time', 'longitude', 'latitude',
55 'model_level_number', 'realization'. A list of multiple coordinates can
56 be given.
57 method: str
58 Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN',
59 'PERCENTILE' getattr creates iris.analysis.MEAN, etc. For PERCENTILE YAML
60 file requires i.e. method: 'PERCENTILE' additional_percent: 90.
61 additional_percent: float, optional
62 Required for the PERCENTILE method. This is a number between 0 and 100.
64 Returns
65 -------
66 collapsed_cubes: iris.cube.Cube | iris.cube.CubeList
67 Single variable but several methods of aggregation
69 Raises
70 ------
71 ValueError
72 If additional_percent wasn't supplied while using PERCENTILE method.
73 """
74 if method == "SEQ" or method == "" or method is None:
75 return cubes
76 if method == "PERCENTILE" and additional_percent is None:
77 raise ValueError("Must specify additional_percent")
79 # Retain only common time points between different models if multiple model inputs.
80 if isinstance(cubes, iris.cube.CubeList) and len(cubes) > 1:
81 logger.debug("Extracting common time points as multiple model inputs detected.")
82 for cube in cubes:
83 cube.coord("forecast_reference_time").bounds = None
84 cube.coord("forecast_period").bounds = None
85 cubes = cubes.extract_overlapping(
86 ["forecast_reference_time", "forecast_period"]
87 )
88 if len(cubes) == 0:
89 raise ValueError("No overlapping times detected in input cubes.")
91 collapsed_cubes = iris.cube.CubeList([])
92 with warnings.catch_warnings():
93 warnings.filterwarnings(
94 "ignore", "Cannot check if coordinate is contiguous", UserWarning
95 )
96 warnings.filterwarnings(
97 "ignore", "Collapsing spatial coordinate.+without weighting", UserWarning
98 )
99 for cube in iter_maybe(cubes):
100 # Apply a mask to check for invalid data, this will allow NaNs to
101 # be ignored.
102 cube.data = np.ma.masked_invalid(cube.data)
103 if method == "PERCENTILE":
104 collapsed_cubes.append(
105 cube.collapsed(
106 coordinate,
107 getattr(iris.analysis, method),
108 percent=additional_percent,
109 )
110 )
111 elif method == "RANGE":
112 cube_max = cube.collapsed(coordinate, iris.analysis.MAX)
113 cube_min = cube.collapsed(coordinate, iris.analysis.MIN)
114 collapsed_cubes.append(cube_max - cube_min)
115 else:
116 collapsed_cubes.append(
117 cube.collapsed(coordinate, getattr(iris.analysis, method))
118 )
119 if len(collapsed_cubes) == 1:
120 return collapsed_cubes[0]
121 else:
122 return collapsed_cubes
125def collapse_by_hour_of_day(
126 cubes: iris.cube.Cube | iris.cube.CubeList,
127 method: str,
128 additional_percent: float | None = None,
129 **kwargs,
130) -> iris.cube.Cube:
131 """Collapse a cube by hour of the day.
133 Collapses a cube by hour of the day in the time coordinates provided by the
134 model. It is useful for creating diurnal cycle plots. It aggregates all 00
135 UTC together regardless of lead time.
137 Arguments
138 ---------
139 cubes: iris.cube.Cube | iris.cube.CubeList
140 Cube to collapse and iterate over one dimension or CubeList to convert
141 to a cube and then collapse prior to aggregating by hour. If a CubeList
142 is provided each cube is handled separately.
143 method: str
144 Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN',
145 'PERCENTILE'. For 'PERCENTILE' the additional_percent must be specified.
147 Returns
148 -------
149 cube: iris.cube.Cube
150 Single variable but several methods of aggregation.
152 Raises
153 ------
154 ValueError
155 If additional_percent wasn't supplied while using PERCENTILE method.
157 Notes
158 -----
159 Collapsing of the cube is around the 'time' coordinate. The coordinates are
160 first grouped by the hour of day, and then aggregated by the hour of day to
161 create a diurnal cycle. This operator is applicable for both single
162 forecasts and for multiple forecasts. The hour used is based on the units of
163 the time coordinate. If the time coordinate is in UTC, hour will be in UTC.
165 To apply this operator successfully there must only be one time dimension.
166 Should a MultiDim exception be raised the user first needs to apply the
167 collapse operator to reduce the time dimensions before applying this
168 operator. A cube containing the two time dimensions
169 'forecast_reference_time' and 'forecast_period' will be automatically
170 collapsed by lead time before being being collapsed by hour of day.
171 """
172 if method == "PERCENTILE" and additional_percent is None:
173 raise ValueError("Must specify additional_percent")
175 # Retain only common time points between different models if multiple model inputs.
176 if isinstance(cubes, iris.cube.CubeList) and len(cubes) > 1:
177 logger.debug("Extracting common time points as multiple model inputs detected.")
178 for cube in cubes:
179 cube.coord("forecast_reference_time").bounds = None
180 cube.coord("forecast_period").bounds = None
181 cubes = cubes.extract_overlapping(
182 ["forecast_reference_time", "forecast_period"]
183 )
184 if len(cubes) == 0:
185 raise ValueError("No overlapping times detected in input cubes.")
187 collapsed_cubes = iris.cube.CubeList([])
188 for cube in iter_maybe(cubes):
189 # Ensure hour coordinate in each input is sorted, and data adjusted if needed.
190 sorted_cube = iris.cube.CubeList()
191 for fcst_slice in cube.slices_over(["forecast_reference_time"]):
192 # Categorise the time coordinate by hour of the day.
193 fcst_slice = add_hour_coordinate(fcst_slice)
194 if method == "PERCENTILE":
195 by_hour = fcst_slice.aggregated_by(
196 "hour", getattr(iris.analysis, method), percent=additional_percent
197 )
198 else:
199 by_hour = fcst_slice.aggregated_by(
200 "hour", getattr(iris.analysis, method)
201 )
202 # Compute if data needs sorting to lie in increasing order [0..23].
203 # Note multiple forecasts can sit in same cube spanning different
204 # initialisation times and data ranges.
205 time_points = by_hour.coord("hour").points
206 time_points_sorted = np.sort(by_hour.coord("hour").points)
207 if time_points[0] != time_points_sorted[0]: 207 ↛ 215line 207 didn't jump to line 215 because the condition on line 207 was always true
208 nroll = time_points[0] / (time_points[1] - time_points[0])
209 # Shift hour coordinate and data cube to be in time of day order.
210 by_hour.coord("hour").points = np.roll(time_points, nroll, 0)
211 by_hour.data = np.roll(by_hour.data, nroll, axis=0)
213 # Remove unnecessary time coordinate.
214 # "hour" and "forecast_period" remain as AuxCoord.
215 by_hour.remove_coord("time")
217 sorted_cube.append(by_hour)
219 # Recombine cube slices.
220 cube = sorted_cube.merge_cube()
222 # Apply a mask to check for invalid data, this will allow NaNs to
223 # be ignored.
224 cube.data = np.ma.masked_invalid(cube.data)
226 if cube.coords("forecast_reference_time", dim_coords=True):
227 # Collapse by forecast reference time to get a single cube.
228 cube = collapse(
229 cube,
230 "forecast_reference_time",
231 method,
232 additional_percent=additional_percent,
233 )
234 else:
235 # Or remove forecast reference time if a single case, as collapse
236 # will have effectively done this.
237 cube.remove_coord("forecast_reference_time")
239 # Promote "hour" to dim_coord.
240 iris.util.promote_aux_coord_to_dim_coord(cube, "hour")
241 collapsed_cubes.append(cube)
243 if len(collapsed_cubes) == 1:
244 return collapsed_cubes[0]
245 else:
246 return collapsed_cubes
249def collapse_by_validity_time(
250 cubes: iris.cube.Cube | iris.cube.CubeList,
251 method: str,
252 additional_percent: float | None = None,
253 **kwargs,
254) -> iris.cube.Cube:
255 """Collapse a cube around validity time for multiple cases.
257 First checks if the data can be aggregated easily. Then creates a new cube
258 by slicing over the time dimensions, removing the time dimensions,
259 re-merging the data, and creating a new time coordinate. It then collapses
260 by the new time coordinate for a specified method using the collapse
261 function.
263 Arguments
264 ---------
265 cubes: iris.cube.Cube | iris.cube.CubeList
266 Cube to collapse by validity time or CubeList that will be converted
267 to a cube before collapsing by validity time.
268 method: str
269 Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN',
270 'PERCENTILE'. For 'PERCENTILE' the additional_percent must be specified.
272 Returns
273 -------
274 cube: iris.cube.Cube | iris.cube.CubeList
275 Single variable collapsed by lead time based on chosen method.
277 Raises
278 ------
279 ValueError
280 If additional_percent wasn't supplied while using PERCENTILE method.
281 """
282 if method == "PERCENTILE" and additional_percent is None:
283 raise ValueError("Must specify additional_percent")
285 collapsed_cubes = iris.cube.CubeList([])
286 for cube in iter_maybe(cubes):
287 # Slice over cube by both time dimensions to create a CubeList.
288 new_cubelist = iris.cube.CubeList(
289 cube.slices_over(["forecast_period", "forecast_reference_time"])
290 )
291 for sub_cube in new_cubelist:
292 # Reconstruct the time coordinate if it is missing.
293 if "time" not in [coord.name() for coord in sub_cube.coords()]:
294 ref_time_coord = sub_cube.coord("forecast_reference_time")
295 ref_units = ref_time_coord.units
296 ref_time = ref_units.num2date(ref_time_coord.points)
297 period_coord = sub_cube.coord("forecast_period")
298 period_units = period_coord.units
299 # Given how we are slicing there will only be one point.
300 period_seconds = period_units.convert(period_coord.points[0], "seconds")
301 period_duration = datetime.timedelta(seconds=period_seconds)
302 time = ref_time + period_duration
303 time_points = ref_units.date2num(time)
304 time_coord = iris.coords.AuxCoord(
305 points=time_points, standard_name="time", units=ref_units
306 )
307 sub_cube.add_aux_coord(time_coord)
308 # Remove forecast_period and forecast_reference_time coordinates.
309 sub_cube.remove_coord("forecast_period")
310 sub_cube.remove_coord("forecast_reference_time")
311 # Create new CubeList by merging with unique = False to produce a validity
312 # time cube.
313 merged_list_1 = new_cubelist.merge(unique=False)
314 # Create a new "fake" coordinate and apply to each remaining cube to allow
315 # final merging to take place into a single cube.
316 equalised_validity_time = iris.coords.AuxCoord(
317 points=0, long_name="equalised_validity_time", units="1"
318 )
319 for sub_cube, eq_valid_time in zip(
320 merged_list_1, range(len(merged_list_1)), strict=True
321 ):
322 sub_cube.add_aux_coord(equalised_validity_time.copy(points=eq_valid_time))
324 # Merge CubeList to create final cube.
325 final_cube = merged_list_1.merge_cube()
326 logger.debug("Pre-collapse validity time cube:\n%s", final_cube)
328 # Apply a mask to check for invalid data, this will allow NaNs to
329 # be ignored.
330 final_cube.data = np.ma.masked_invalid(final_cube.data)
332 # Collapse over equalised_validity_time as a proxy for equal validity
333 # time.
334 try:
335 collapsed_cube = collapse(
336 final_cube,
337 "equalised_validity_time",
338 method,
339 additional_percent=additional_percent,
340 )
341 except iris.exceptions.CoordinateCollapseError as err:
342 raise ValueError(
343 "Cubes do not overlap therefore cannot collapse across validity time."
344 ) from err
345 collapsed_cube.remove_coord("equalised_validity_time")
346 collapsed_cubes.append(collapsed_cube)
348 if len(collapsed_cubes) == 1:
349 return collapsed_cubes[0]
350 else:
351 return collapsed_cubes
354def proportion(
355 cubes: iris.cube.Cube | iris.cube.CubeList,
356 coordinate: str | list[str],
357 condition: str,
358 threshold: float,
359 **kwargs,
360) -> iris.cube.Cube | iris.cube.CubeList:
361 """Find the proportion of an event for all cubes.
363 Find the proportion of points at a specified threhsold in each cube into a
364 cube collapsing around the specified coordinate(s).
366 Arguments
367 ---------
368 cubes: iris.cube.Cube | iris.cube.CubeList
369 Cube or CubeList to collapse and iterate over one dimension
370 coordinate: str | list[str]
371 Coordinate(s) to collapse over e.g. 'time', 'longitude', 'latitude',
372 'model_level_number', 'realization'. A list of multiple coordinates can
373 be given.
374 condition: str
375 The condition for the event. Expected arguments are eq, ne, lt, gt, le, ge.
376 The letters correspond to the following conditions
377 eq: equal to;
378 ne: not equal to;
379 lt: less than;
380 gt: greater than;
381 le: less than or equal to;
382 ge: greater than or equal to.
383 threshold: float
384 The value for the event.
386 Returns
387 -------
388 collapsed_cubes: iris.cube.Cube | iris.cube.CubeList
389 The proportion of the event.
390 """
391 # Set method
392 method = "PROPORTION"
393 # Retain only common time points between different models if multiple model inputs.
394 if isinstance(cubes, iris.cube.CubeList) and len(cubes) > 1: 394 ↛ 395line 394 didn't jump to line 395 because the condition on line 394 was never true
395 logger.debug("Extracting common time points as multiple model inputs detected.")
396 for cube in cubes:
397 cube.coord("forecast_reference_time").bounds = None
398 cube.coord("forecast_period").bounds = None
399 cubes = cubes.extract_overlapping(
400 ["forecast_reference_time", "forecast_period"]
401 )
402 if len(cubes) == 0:
403 raise ValueError("No overlapping times detected in input cubes.")
405 collapsed_cubes = iris.cube.CubeList([])
406 with warnings.catch_warnings():
407 warnings.filterwarnings(
408 "ignore", "Cannot check if coordinate is contiguous", UserWarning
409 )
410 warnings.filterwarnings(
411 "ignore", "Collapsing spatial coordinate.+without weighting", UserWarning
412 )
413 for cube in iter_maybe(cubes):
414 # Apply a mask to check for invalid data, this will allow NaNs to
415 # be ignored.
416 cube.data = np.ma.masked_invalid(cube.data)
417 match condition:
418 case "eq":
419 new_cube = cube.collapsed(
420 coordinate,
421 getattr(iris.analysis, method),
422 function=lambda values: values == threshold,
423 )
424 case "ne":
425 new_cube = cube.collapsed(
426 coordinate,
427 getattr(iris.analysis, method),
428 function=lambda values: values != threshold,
429 )
430 case "gt":
431 new_cube = cube.collapsed(
432 coordinate,
433 getattr(iris.analysis, method),
434 function=lambda values: values > threshold,
435 )
436 case "ge":
437 new_cube = cube.collapsed(
438 coordinate,
439 getattr(iris.analysis, method),
440 function=lambda values: values >= threshold,
441 )
442 case "lt":
443 new_cube = cube.collapsed(
444 coordinate,
445 getattr(iris.analysis, method),
446 function=lambda values: values < threshold,
447 )
448 case "le":
449 new_cube = cube.collapsed(
450 coordinate,
451 getattr(iris.analysis, method),
452 function=lambda values: values <= threshold,
453 )
454 case _:
455 raise ValueError(
456 """Unexpected value for condition. Expected eq, ne, gt, ge, lt, le. Got {condition}."""
457 )
458 name = cube.long_name if cube.long_name else cube.name()
459 new_cube.rename(f"probability_of_{name}_{condition}_{threshold}")
460 new_cube.units = "1"
461 collapsed_cubes.append(new_cube)
463 if len(collapsed_cubes) == 1: 463 ↛ 466line 463 didn't jump to line 466 because the condition on line 463 was always true
464 return collapsed_cubes[0]
465 else:
466 return collapsed_cubes
469# TODO
470# Collapse function that calculates means, medians etc across members of an
471# ensemble or stratified groups. Need to allow collapse over realisation
472# dimension for fixed time. Hence will require reading in of CubeList