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