Coverage for src/CSET/operators/collapse.py: 65%
212 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-15 08: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
34def _coord_is_effectively_scalar(coord):
35 """True if coordinate is scalar or all points have the same value."""
36 points = np.asarray(coord.points)
38 if points.size <= 1:
39 return True
41 return np.all(points == points.flat[0])
44def _fix_analysis_forecasttime(cubes: iris.cube.CubeList):
46 analysis_cubes = iris.cube.CubeList()
47 forecast_cubes = iris.cube.CubeList()
49 for cube in cubes:
50 fp = cube.coord("forecast_period")
52 if fp.points.max() == 0: 52 ↛ 53line 52 didn't jump to line 53 because the condition on line 52 was never true
53 analysis_cubes.append(cube)
54 else:
55 forecast_cubes.append(cube)
57 longest_forecast_period = max(
58 cube.coord("forecast_period").points.max() for cube in forecast_cubes
59 )
61 frts = {cube.coord("forecast_reference_time").points[0] for cube in forecast_cubes}
63 # We might have multiple cubes, with different reference times,
64 # so we need to go through all possibilities and ensure analysis
65 # matches all of these. Maybe a list of forecast reference times and leadtimes
66 # to construct the appropriate cubes?
67 if len(frts) > 1:
68 raise ValueError("Forecast cubes have different forecast_reference_times")
70 frt_coord = forecast_cubes[0].coord("forecast_reference_time")
72 forecast_reference_time = frt_coord.units.num2date(frt_coord.points[0])
74 longest_cube = max(
75 forecast_cubes, key=lambda c: c.coord("forecast_period").points.max()
76 )
78 forecast_reference_time = longest_cube.coord(
79 "forecast_reference_time"
80 ).units.num2date(longest_cube.coord("forecast_reference_time").points[0])
82 forecast_end_time = longest_cube.coord("time").units.num2date(
83 longest_cube.coord("time").points.max()
84 )
86 time_constraint = iris.Constraint(
87 time=lambda cell: forecast_reference_time <= cell.point <= forecast_end_time
88 )
90 # manually slicing needed rather than just assigning cube = cube.extract...
91 for i, cube in enumerate(analysis_cubes): 91 ↛ 92line 91 didn't jump to line 92 because the loop on line 91 never started
92 analysis_cubes[i] = cube.extract(time_constraint)
94 fixed_cubes = iris.cube.CubeList()
96 for cube in forecast_cubes:
97 fixed_cubes.append(cube)
99 print(analysis_cubes)
100 print(analysis_cubes[0])
101 print(analysis_cubes[0].coord("time"))
102 print(analysis_cubes[0].coord("forecast_reference_time"))
103 print(analysis_cubes[0].coord("forecast_period"))
104 print("pre")
106 for cube in analysis_cubes:
107 # Get coords
108 time_coord = cube.coord("time")
109 frt_coord = cube.coord("forecast_reference_time")
110 fp_coord = cube.coord("forecast_period")
112 if _coord_is_effectively_scalar(frt_coord) and _coord_is_effectively_scalar(
113 fp_coord
114 ):
115 print("forecast_reference_time and forecast_period already static")
116 return cube
118 #
119 frt0_hours = frt_coord.points[0]
121 fp_hours = time_coord.points - frt0_hours
123 # Remove old coords
124 cube.remove_coord("forecast_reference_time")
125 cube.remove_coord("forecast_period")
127 # Add constant FRT back as scalar
128 new_frt = iris.coords.AuxCoord(
129 frt0_hours,
130 standard_name="forecast_reference_time",
131 units=time_coord.units,
132 )
133 cube.add_aux_coord(new_frt)
135 # Add forecast period varying along time dimension
136 time_dim = cube.coord_dims("time")[0]
138 new_fp = iris.coords.AuxCoord(
139 fp_hours,
140 standard_name="forecast_period",
141 units="hours",
142 )
144 cube.add_aux_coord(new_fp, data_dims=(time_dim,))
146 fixed_cubes.append(cube)
148 print(fixed_cubes)
149 print(fixed_cubes[0])
150 print(fixed_cubes[0].coord("time"))
151 print(fixed_cubes[0].coord("forecast_reference_time"))
152 print(fixed_cubes[0].coord("forecast_period"))
154 # something to check we are matching input, i.e. if cube, only return cube, not cubelist.
156 return fixed_cubes
159def collapse(
160 cubes: iris.cube.Cube | iris.cube.CubeList,
161 coordinate: str | list[str],
162 method: str,
163 additional_percent: float = None,
164 **kwargs,
165) -> iris.cube.Cube | iris.cube.CubeList:
166 """Collapse coordinate(s) of a single cube or of every cube in a cube list.
168 Collapses similar fields in each cube into a cube collapsing around the
169 specified coordinate(s) and method. This could be a (weighted) mean or
170 percentile.
172 Arguments
173 ---------
174 cubes: iris.cube.Cube | iris.cube.CubeList
175 Cube or CubeList to collapse and iterate over one dimension
176 coordinate: str | list[str]
177 Coordinate(s) to collapse over e.g. 'time', 'longitude', 'latitude',
178 'model_level_number', 'realization'. A list of multiple coordinates can
179 be given.
180 method: str
181 Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN',
182 'PERCENTILE' getattr creates iris.analysis.MEAN, etc. For PERCENTILE YAML
183 file requires i.e. method: 'PERCENTILE' additional_percent: 90.
184 additional_percent: float, optional
185 Required for the PERCENTILE method. This is a number between 0 and 100.
187 Returns
188 -------
189 collapsed_cubes: iris.cube.Cube | iris.cube.CubeList
190 Single variable but several methods of aggregation
192 Raises
193 ------
194 ValueError
195 If additional_percent wasn't supplied while using PERCENTILE method.
196 """
197 if method == "SEQ" or method == "" or method is None: 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 return cubes
199 if method == "PERCENTILE" and additional_percent is None:
200 raise ValueError("Must specify additional_percent")
202 # Retain only common time points between different models if multiple model inputs.
204 cubes = _fix_analysis_forecasttime(cubes)
206 if isinstance(cubes, iris.cube.CubeList) and len(cubes) > 1:
207 logging.debug(
208 "Extracting common time points as multiple model inputs detected."
209 )
210 for cube in cubes:
211 cube.coord("forecast_reference_time").bounds = None
212 cube.coord("forecast_period").bounds = None
213 cubes = cubes.extract_overlapping(
214 ["forecast_reference_time", "forecast_period"]
215 )
216 if len(cubes) == 0:
217 raise ValueError("No overlapping times detected in input cubes.")
219 collapsed_cubes = iris.cube.CubeList([])
220 with warnings.catch_warnings():
221 warnings.filterwarnings(
222 "ignore", "Cannot check if coordinate is contiguous", UserWarning
223 )
224 warnings.filterwarnings(
225 "ignore", "Collapsing spatial coordinate.+without weighting", UserWarning
226 )
227 for cube in iter_maybe(cubes):
228 # Apply a mask to check for invalid data, this will allow NaNs to
229 # be ignored.
230 cube.data = np.ma.masked_invalid(cube.data)
231 if method == "PERCENTILE":
232 collapsed_cubes.append(
233 cube.collapsed(
234 coordinate,
235 getattr(iris.analysis, method),
236 percent=additional_percent,
237 )
238 )
239 elif method == "RANGE":
240 cube_max = cube.collapsed(coordinate, iris.analysis.MAX)
241 cube_min = cube.collapsed(coordinate, iris.analysis.MIN)
242 collapsed_cubes.append(cube_max - cube_min)
243 else:
244 collapsed_cubes.append(
245 cube.collapsed(coordinate, getattr(iris.analysis, method))
246 )
247 if len(collapsed_cubes) == 1:
248 return collapsed_cubes[0]
249 else:
250 return collapsed_cubes
253def collapse_by_hour_of_day(
254 cubes: iris.cube.Cube | iris.cube.CubeList,
255 method: str,
256 additional_percent: float = None,
257 **kwargs,
258) -> iris.cube.Cube:
259 """Collapse a cube by hour of the day.
261 Collapses a cube by hour of the day in the time coordinates provided by the
262 model. It is useful for creating diurnal cycle plots. It aggregates all 00
263 UTC together regardless of lead time.
265 Arguments
266 ---------
267 cubes: iris.cube.Cube | iris.cube.CubeList
268 Cube to collapse and iterate over one dimension or CubeList to convert
269 to a cube and then collapse prior to aggregating by hour. If a CubeList
270 is provided each cube is handled separately.
271 method: str
272 Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN',
273 'PERCENTILE'. For 'PERCENTILE' the additional_percent must be specified.
275 Returns
276 -------
277 cube: iris.cube.Cube
278 Single variable but several methods of aggregation.
280 Raises
281 ------
282 ValueError
283 If additional_percent wasn't supplied while using PERCENTILE method.
285 Notes
286 -----
287 Collapsing of the cube is around the 'time' coordinate. The coordinates are
288 first grouped by the hour of day, and then aggregated by the hour of day to
289 create a diurnal cycle. This operator is applicable for both single
290 forecasts and for multiple forecasts. The hour used is based on the units of
291 the time coordinate. If the time coordinate is in UTC, hour will be in UTC.
293 To apply this operator successfully there must only be one time dimension.
294 Should a MultiDim exception be raised the user first needs to apply the
295 collapse operator to reduce the time dimensions before applying this
296 operator. A cube containing the two time dimensions
297 'forecast_reference_time' and 'forecast_period' will be automatically
298 collapsed by lead time before being being collapsed by hour of day.
299 """
300 if method == "PERCENTILE" and additional_percent is None:
301 raise ValueError("Must specify additional_percent")
303 # Retain only common time points between different models if multiple model inputs.
304 if isinstance(cubes, iris.cube.CubeList) and len(cubes) > 1:
305 logging.debug(
306 "Extracting common time points as multiple model inputs detected."
307 )
308 for cube in cubes:
309 cube.coord("forecast_reference_time").bounds = None
310 cube.coord("forecast_period").bounds = None
311 cubes = cubes.extract_overlapping(
312 ["forecast_reference_time", "forecast_period"]
313 )
314 if len(cubes) == 0:
315 raise ValueError("No overlapping times detected in input cubes.")
317 collapsed_cubes = iris.cube.CubeList([])
318 for cube in iter_maybe(cubes):
319 # Ensure hour coordinate in each input is sorted, and data adjusted if needed.
320 sorted_cube = iris.cube.CubeList()
321 for fcst_slice in cube.slices_over(["forecast_reference_time"]):
322 # Categorise the time coordinate by hour of the day.
323 fcst_slice = add_hour_coordinate(fcst_slice)
324 if method == "PERCENTILE":
325 by_hour = fcst_slice.aggregated_by(
326 "hour", getattr(iris.analysis, method), percent=additional_percent
327 )
328 else:
329 by_hour = fcst_slice.aggregated_by(
330 "hour", getattr(iris.analysis, method)
331 )
332 # Compute if data needs sorting to lie in increasing order [0..23].
333 # Note multiple forecasts can sit in same cube spanning different
334 # initialisation times and data ranges.
335 time_points = by_hour.coord("hour").points
336 time_points_sorted = np.sort(by_hour.coord("hour").points)
337 if time_points[0] != time_points_sorted[0]: 337 ↛ 345line 337 didn't jump to line 345 because the condition on line 337 was always true
338 nroll = time_points[0] / (time_points[1] - time_points[0])
339 # Shift hour coordinate and data cube to be in time of day order.
340 by_hour.coord("hour").points = np.roll(time_points, nroll, 0)
341 by_hour.data = np.roll(by_hour.data, nroll, axis=0)
343 # Remove unnecessary time coordinate.
344 # "hour" and "forecast_period" remain as AuxCoord.
345 by_hour.remove_coord("time")
347 sorted_cube.append(by_hour)
349 # Recombine cube slices.
350 cube = sorted_cube.merge_cube()
352 # Apply a mask to check for invalid data, this will allow NaNs to
353 # be ignored.
354 cube.data = np.ma.masked_invalid(cube.data)
356 if cube.coords("forecast_reference_time", dim_coords=True):
357 # Collapse by forecast reference time to get a single cube.
358 cube = collapse(
359 cube,
360 "forecast_reference_time",
361 method,
362 additional_percent=additional_percent,
363 )
364 else:
365 # Or remove forecast reference time if a single case, as collapse
366 # will have effectively done this.
367 cube.remove_coord("forecast_reference_time")
369 # Promote "hour" to dim_coord.
370 iris.util.promote_aux_coord_to_dim_coord(cube, "hour")
371 collapsed_cubes.append(cube)
373 if len(collapsed_cubes) == 1:
374 return collapsed_cubes[0]
375 else:
376 return collapsed_cubes
379def collapse_by_validity_time(
380 cubes: iris.cube.Cube | iris.cube.CubeList,
381 method: str,
382 additional_percent: float = None,
383 **kwargs,
384) -> iris.cube.Cube:
385 """Collapse a cube around validity time for multiple cases.
387 First checks if the data can be aggregated easily. Then creates a new cube
388 by slicing over the time dimensions, removing the time dimensions,
389 re-merging the data, and creating a new time coordinate. It then collapses
390 by the new time coordinate for a specified method using the collapse
391 function.
393 Arguments
394 ---------
395 cubes: iris.cube.Cube | iris.cube.CubeList
396 Cube to collapse by validity time or CubeList that will be converted
397 to a cube before collapsing by validity time.
398 method: str
399 Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN',
400 'PERCENTILE'. For 'PERCENTILE' the additional_percent must be specified.
402 Returns
403 -------
404 cube: iris.cube.Cube | iris.cube.CubeList
405 Single variable collapsed by lead time based on chosen method.
407 Raises
408 ------
409 ValueError
410 If additional_percent wasn't supplied while using PERCENTILE method.
411 """
412 if method == "PERCENTILE" and additional_percent is None:
413 raise ValueError("Must specify additional_percent")
415 collapsed_cubes = iris.cube.CubeList([])
416 for cube in iter_maybe(cubes): 416 ↛ 478line 416 didn't jump to line 478 because the loop on line 416 didn't complete
417 # Slice over cube by both time dimensions to create a CubeList.
418 new_cubelist = iris.cube.CubeList(
419 cube.slices_over(["forecast_period", "forecast_reference_time"])
420 )
421 for sub_cube in new_cubelist:
422 # Reconstruct the time coordinate if it is missing.
423 if "time" not in [coord.name() for coord in sub_cube.coords()]:
424 ref_time_coord = sub_cube.coord("forecast_reference_time")
425 ref_units = ref_time_coord.units
426 ref_time = ref_units.num2date(ref_time_coord.points)
427 period_coord = sub_cube.coord("forecast_period")
428 period_units = period_coord.units
429 # Given how we are slicing there will only be one point.
430 period_seconds = period_units.convert(period_coord.points[0], "seconds")
431 period_duration = datetime.timedelta(seconds=period_seconds)
432 time = ref_time + period_duration
433 time_points = ref_units.date2num(time)
434 time_coord = iris.coords.AuxCoord(
435 points=time_points, standard_name="time", units=ref_units
436 )
437 sub_cube.add_aux_coord(time_coord)
438 # Remove forecast_period and forecast_reference_time coordinates.
439 sub_cube.remove_coord("forecast_period")
440 sub_cube.remove_coord("forecast_reference_time")
441 # Create new CubeList by merging with unique = False to produce a validity
442 # time cube.
443 merged_list_1 = new_cubelist.merge(unique=False)
444 # Create a new "fake" coordinate and apply to each remaining cube to allow
445 # final merging to take place into a single cube.
446 equalised_validity_time = iris.coords.AuxCoord(
447 points=0, long_name="equalised_validity_time", units="1"
448 )
449 for sub_cube, eq_valid_time in zip(
450 merged_list_1, range(len(merged_list_1)), strict=True
451 ):
452 sub_cube.add_aux_coord(equalised_validity_time.copy(points=eq_valid_time))
454 # Merge CubeList to create final cube.
455 final_cube = merged_list_1.merge_cube()
456 logging.debug("Pre-collapse validity time cube:\n%s", final_cube)
458 # Apply a mask to check for invalid data, this will allow NaNs to
459 # be ignored.
460 final_cube.data = np.ma.masked_invalid(final_cube.data)
462 # Collapse over equalised_validity_time as a proxy for equal validity
463 # time.
464 try:
465 collapsed_cube = collapse(
466 final_cube,
467 "equalised_validity_time",
468 method,
469 additional_percent=additional_percent,
470 )
471 except iris.exceptions.CoordinateCollapseError as err:
472 raise ValueError(
473 "Cubes do not overlap therefore cannot collapse across validity time."
474 ) from err
475 collapsed_cube.remove_coord("equalised_validity_time")
476 collapsed_cubes.append(collapsed_cube)
478 if len(collapsed_cubes) == 1:
479 return collapsed_cubes[0]
480 else:
481 return collapsed_cubes
484def proportion(
485 cubes: iris.cube.Cube | iris.cube.CubeList,
486 coordinate: str | list[str],
487 condition: str,
488 threshold: float,
489 **kwargs,
490) -> iris.cube.Cube | iris.cube.CubeList:
491 """Find the proportion of an event for all cubes.
493 Find the proportion of points at a specified threhsold in each cube into a
494 cube collapsing around the specified coordinate(s).
496 Arguments
497 ---------
498 cubes: iris.cube.Cube | iris.cube.CubeList
499 Cube or CubeList to collapse and iterate over one dimension
500 coordinate: str | list[str]
501 Coordinate(s) to collapse over e.g. 'time', 'longitude', 'latitude',
502 'model_level_number', 'realization'. A list of multiple coordinates can
503 be given.
504 condition: str
505 The condition for the event. Expected arguments are eq, ne, lt, gt, le, ge.
506 The letters correspond to the following conditions
507 eq: equal to;
508 ne: not equal to;
509 lt: less than;
510 gt: greater than;
511 le: less than or equal to;
512 ge: greater than or equal to.
513 threshold: float
514 The value for the event.
516 Returns
517 -------
518 collapsed_cubes: iris.cube.Cube | iris.cube.CubeList
519 The proportion of the event.
520 """
521 # Set method
522 method = "PROPORTION"
523 # Retain only common time points between different models if multiple model inputs.
524 if isinstance(cubes, iris.cube.CubeList) and len(cubes) > 1: 524 ↛ 525line 524 didn't jump to line 525 because the condition on line 524 was never true
525 logging.debug(
526 "Extracting common time points as multiple model inputs detected."
527 )
528 for cube in cubes:
529 cube.coord("forecast_reference_time").bounds = None
530 cube.coord("forecast_period").bounds = None
531 cubes = cubes.extract_overlapping(
532 ["forecast_reference_time", "forecast_period"]
533 )
534 if len(cubes) == 0:
535 raise ValueError("No overlapping times detected in input cubes.")
537 collapsed_cubes = iris.cube.CubeList([])
538 with warnings.catch_warnings():
539 warnings.filterwarnings(
540 "ignore", "Cannot check if coordinate is contiguous", UserWarning
541 )
542 warnings.filterwarnings(
543 "ignore", "Collapsing spatial coordinate.+without weighting", UserWarning
544 )
545 for cube in iter_maybe(cubes):
546 # Apply a mask to check for invalid data, this will allow NaNs to
547 # be ignored.
548 cube.data = np.ma.masked_invalid(cube.data)
549 match condition:
550 case "eq":
551 new_cube = cube.collapsed(
552 coordinate,
553 getattr(iris.analysis, method),
554 function=lambda values: values == threshold,
555 )
556 case "ne":
557 new_cube = cube.collapsed(
558 coordinate,
559 getattr(iris.analysis, method),
560 function=lambda values: values != threshold,
561 )
562 case "gt":
563 new_cube = cube.collapsed(
564 coordinate,
565 getattr(iris.analysis, method),
566 function=lambda values: values > threshold,
567 )
568 case "ge":
569 new_cube = cube.collapsed(
570 coordinate,
571 getattr(iris.analysis, method),
572 function=lambda values: values >= threshold,
573 )
574 case "lt":
575 new_cube = cube.collapsed(
576 coordinate,
577 getattr(iris.analysis, method),
578 function=lambda values: values < threshold,
579 )
580 case "le":
581 new_cube = cube.collapsed(
582 coordinate,
583 getattr(iris.analysis, method),
584 function=lambda values: values <= threshold,
585 )
586 case _:
587 raise ValueError(
588 """Unexpected value for condition. Expected eq, ne, gt, ge, lt, le. Got {condition}."""
589 )
590 name = cube.long_name if cube.long_name else cube.name()
591 new_cube.rename(f"probability_of_{name}_{condition}_{threshold}")
592 new_cube.units = "1"
593 collapsed_cubes.append(new_cube)
595 if len(collapsed_cubes) == 1: 595 ↛ 598line 595 didn't jump to line 598 because the condition on line 595 was always true
596 return collapsed_cubes[0]
597 else:
598 return collapsed_cubes
601# TODO
602# Collapse function that calculates means, medians etc across members of an
603# ensemble or stratified groups. Need to allow collapse over realisation
604# dimension for fixed time. Hence will require reading in of CubeList