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