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