Coverage for src/CSET/operators/aggregate.py: 90%
121 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-22 13:57 +0000
1# © Crown copyright, Met Office (2022-2024) 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 aggregate across either 1 or 2 dimensions."""
17import logging
19import iris
20import iris.analysis
21import iris.coord_categorisation
22import iris.cube
23import iris.exceptions
24import iris.util
25import isodate
26import numpy as np
28from CSET._common import iter_maybe
29from CSET.operators._utils import (
30 identify_unique_times,
31 is_time_aggregatable,
32 is_time_aux_coord,
33 remove_cell_method,
34 remove_duplicates,
35)
37logger = logging.getLogger(__name__)
40def _add_nref(cube: iris.cube.Cube):
41 """Retain information on number of forecast_reference_time inputs.
43 This preserves information on number of aggregated cases that can
44 otherwise be lost on subsequent calls to collapse functions.
45 """
46 nref = np.size(cube.coord("forecast_reference_time").points)
47 cube.coord("time").attributes["number_reference_times"] = nref
48 return cube
51def time_aggregate(
52 cube: iris.cube.Cube,
53 method: str,
54 interval_iso: str,
55 **kwargs,
56) -> iris.cube.Cube:
57 """Aggregate cube by its time coordinate.
59 Aggregates similar (stash) fields in a cube for the specified coordinate and
60 using the method supplied. The aggregated cube will keep the coordinate and
61 add a further coordinate with the aggregated end time points.
63 Examples are: 1. Generating hourly or 6-hourly precipitation accumulations
64 given an interval for the new time coordinate.
66 We use the isodate class to convert ISO 8601 durations into time intervals
67 for creating a new time coordinate for aggregation.
69 We use the lambda function to pass coord and interval into the callable
70 category function in add_categorised to allow users to define their own
71 sub-daily intervals for the new time coordinate.
73 Arguments
74 ---------
75 cube: iris.cube.Cube
76 Cube to aggregate and iterate over one dimension
77 coordinate: str
78 Coordinate to aggregate over i.e. 'time', 'longitude',
79 'latitude','model_level_number'.
80 method: str
81 Type of aggregate i.e. method: 'SUM', getattr creates
82 iris.analysis.SUM, etc.
83 interval_iso: isodate timedelta ISO 8601 object i.e PT6H (6 hours), PT30M (30 mins)
84 Interval to aggregate over.
86 Returns
87 -------
88 cube: iris.cube.Cube
89 Single variable but several methods of aggregation
91 Raises
92 ------
93 ValueError
94 If the constraint doesn't produce a single cube containing a field.
95 """
96 # Duration of ISO timedelta.
97 timedelta = isodate.parse_duration(interval_iso)
99 # Convert interval format to whole hours.
100 interval = int(timedelta.total_seconds() / 3600)
102 # Add time categorisation overwriting hourly increment via lambda coord.
103 # https://scitools-iris.readthedocs.io/en/latest/_modules/iris/coord_categorisation.html
104 iris.coord_categorisation.add_categorised_coord(
105 cube, "interval", "time", lambda coord, cell: cell // interval * interval
106 )
108 # Aggregate cube using supplied method.
109 aggregated_cube = cube.aggregated_by("interval", getattr(iris.analysis, method))
110 aggregated_cube.remove_coord("interval")
111 return aggregated_cube
114def ensure_aggregatable_across_cases(
115 cubes: iris.cube.Cube | iris.cube.CubeList,
116 time_coord_name: str,
117) -> iris.cube.CubeList:
118 """Ensure a Cube or CubeList can be aggregated across multiple cases.
120 The cubes are grouped into buckets of compatible cubes. Each bucket is then
121 processed to create aggregatable cubes. The function handles two types of
122 time coordinates:
124 - Dimension coordinates: Cubes with ``forecast_period`` and
125 ``forecast_reference_time`` as dimension coordinates are sliced and merged.
126 - Auxiliary coordinates: Cubes with time as an auxiliary coordinate are
127 aggregated at each unique time point and then merged.
129 Arguments
130 ---------
131 cubes: iris.cube.Cube | iris.cube.CubeList
132 Each cube is checked to determine if it has the necessary time
133 coordinates (either as dimension or auxiliary coordinates) to be
134 aggregatable, being processed if needed.
135 time_coord_name: str
136 Name of the time coordinate to aggregate over, typically "time" or
137 "forecast_period".
139 Returns
140 -------
141 cubes: iris.cube.CubeList
142 A CubeList of time aggregatable cubes. Each cube will have a
143 ``number_reference_times`` attribute on its time coordinate indicating
144 the number of forecast reference times aggregated.
146 Raises
147 ------
148 ValueError
149 If any of the provided cubes cannot be made aggregatable (i.e., missing
150 required time coordinates).
152 Notes
153 -----
154 This operator is designed to ensure that cubes can be aggregated across
155 multiple cases (e.g., different model runs or forecast times). Its
156 functionality is particularly useful for case study or trial aggregation
157 when computing statistics such as percentiles, Q-Q plots, and histograms.
159 For cubes to be aggregatable, they must have either:
161 - ``forecast_period`` and ``forecast_reference_time`` as dimension
162 or auxiliary coordinates
163 """
165 # Group compatible cubes.
166 class Buckets:
167 def __init__(self):
168 self.buckets = []
170 def add(self, cube: iris.cube.Cube):
171 """Add a cube into a bucket.
173 If the cube is compatible with an existing bucket it is added there.
174 Otherwise it gets its own bucket.
175 """
176 for bucket in self.buckets:
177 if bucket[0].is_compatible(cube):
178 bucket.append(cube)
179 return
180 self.buckets.append(iris.cube.CubeList([cube]))
182 def get_buckets(self) -> list[iris.cube.CubeList]:
183 return self.buckets
185 b = Buckets()
186 for cube in iter_maybe(cubes):
187 b.add(cube)
188 buckets = b.get_buckets()
190 logger.debug("Buckets:\n%s", "\n---\n".join(str(b) for b in buckets))
192 # Ensure each bucket is a single aggregatable cube.
193 aggregatable_cubes = iris.cube.CubeList()
194 for bucket in buckets:
195 # Single cubes that are already aggregatable won't need processing.
196 if len(bucket) == 1 and is_time_aggregatable(bucket[0]):
197 aggregatable_cube = bucket[0]
198 aggregatable_cube = _add_nref(aggregatable_cube)
199 aggregatable_cubes.append(aggregatable_cube)
200 continue
202 # Check if all cubes in bucket are time_aggregatable
203 if all(is_time_aggregatable(cube) for cube in bucket):
204 to_merge = iris.cube.CubeList()
205 for cube in bucket:
206 try:
207 to_merge.extend(
208 cube.slices_over(["forecast_period", "forecast_reference_time"])
209 )
210 except iris.exceptions.CoordinateNotFoundError as err:
211 raise ValueError(
212 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.",
213 cube,
214 ) from err
215 aggregatable_cube = to_merge.merge_cube()
217 # Add attribute on number of forecast_reference_times
218 aggregatable_cube = _add_nref(aggregatable_cube)
220 aggregatable_cubes.append(aggregatable_cube)
222 elif all(is_time_aux_coord(cube) for cube in bucket): 222 ↛ 223line 222 didn't jump to line 223 because the condition on line 222 was never true
223 times = identify_unique_times(bucket, time_coord_name)
225 aggregated_list = []
226 for time in times:
227 for cube in bucket:
228 aggregated_list.extend(
229 _aggregate_without_time_dimcoords(
230 cube, time, iris.analysis.MEAN, None
231 )
232 )
234 cubes_to_merge = iris.cube.CubeList(aggregated_list)
236 aggregatable_cube = cubes_to_merge.merge_cube()
237 aggregatable_cube = _add_nref(aggregatable_cube)
238 aggregatable_cubes.append(aggregatable_cube)
240 else:
241 raise ValueError(
242 "Cube is missing required dimension or auxiliary time coordinates",
243 bucket,
244 )
246 return aggregatable_cubes
249def add_hour_coordinate(
250 cubes: iris.cube.Cube | iris.cube.CubeList,
251) -> iris.cube.Cube | iris.cube.CubeList:
252 """Add a category coordinate of hour of day to a Cube or CubeList.
254 Arguments
255 ---------
256 cubes: iris.cube.Cube | iris.cube.CubeList
257 Cube of any variable that has a time coordinate.
258 Note input Cube or CubeList items should only have 1 time dimension.
260 Returns
261 -------
262 cube: iris.cube.Cube
263 A Cube with an additional auxiliary coordinate of hour.
265 Notes
266 -----
267 This is a simple operator designed to be used prior to case aggregation for
268 histograms, Q-Q plots, and percentiles when aggregated by hour of day.
269 """
270 new_cubelist = iris.cube.CubeList()
271 for cube in iter_maybe(cubes):
272 # Add a category coordinate of hour into each cube.
273 iris.util.promote_aux_coord_to_dim_coord(cube, "time")
274 iris.coord_categorisation.add_hour(cube, "time", name="hour")
275 cube.coord("hour").units = "hours"
276 new_cubelist.append(cube)
278 if len(new_cubelist) == 1:
279 return new_cubelist[0]
280 else:
281 return new_cubelist
284def rolling_window_time_aggregation(
285 cubes: iris.cube.Cube | iris.cube.CubeList, method: str, window: int
286) -> iris.cube.Cube | iris.cube.CubeList:
287 """Aggregate a cube along the time dimension using a rolling window.
289 Arguments
290 ---------
291 cubes: iris.cube.Cube | iris.cube.CubeList
292 Cube or Cubelist of any variable to be aggregated over a rolling window
293 in time.
294 method: str
295 Type of aggregate i.e. method: 'MAX', getattr creates
296 iris.analysis.MAX, etc.
297 window: int
298 The rolling window size.
300 Returns
301 -------
302 cube: iris.cube.Cube | iris.cube.CubeList
303 A Cube or Cubelist of the rolling window aggregate. The Cubes will have
304 a time dimension that is reduced in size to the original cube by the
305 window size.
307 Notes
308 -----
309 This operator is designed to be used to help create daily maxima and minima
310 for any variable.
311 """
312 new_cubelist = iris.cube.CubeList()
313 for cube in iter_maybe(cubes):
314 # Use a rolling window in time to applied specified aggregation method
315 # over a specified window length.
316 window_cube = cube.rolling_window(
317 "time", getattr(iris.analysis, method), window
318 )
319 new_cubelist.append(window_cube)
321 if len(new_cubelist) == 1:
322 return new_cubelist[0]
323 else:
324 return new_cubelist
327def _aggregate_without_time_dimcoords(cubes, time_coord, aggregator, percentile):
328 """Aggregate cubes at a specific time without time dimension coordinates.
330 Arguments
331 ---------
332 cubes: iris.cube.CubeList
333 Cubes to aggregate.
334 time_coord: iris.coords.Coord
335 Single time coordinate point to aggregate at.
336 aggregator: iris.analysis.Aggregator
337 Aggregation method (e.g., iris.analysis.MEAN).
338 percentile: float, optional
339 Percentile value if using PERCENTILE aggregator.
340 """
341 # Handle single cube input
342 if isinstance(cubes, iris.cube.Cube):
343 cubes = iris.cube.CubeList([cubes])
345 # Check the supplied time coordinate to make sure it corresponds to a
346 # single time only
347 if len(time_coord.points) != 1:
348 raise ValueError("Time coordinate should specify a single time only")
350 # Remove any duplicate cubes in the input
351 cubes = remove_duplicates(cubes)
353 time_coord_name = time_coord.name()
355 time_constraint = iris.Constraint(
356 coord_values={time_coord_name: lambda cell: cell.point in time_coord.cells()}
357 )
358 cubes_at_time = cubes.extract(time_constraint)
360 # Add a temporary "number" coordinate to uniquely label the different
361 # data points at this time
362 number = 0
363 numbered_cubes = iris.cube.CubeList()
364 for cube in cubes_at_time:
365 for slc in cube.slices_over(time_coord_name):
366 number_coord = iris.coords.AuxCoord(number, long_name="number")
367 slc.add_aux_coord(number_coord)
368 numbered_cubes.append(slc)
369 number += 1
370 cubes_at_time = numbered_cubes
372 cubes_at_time = cubes_at_time.merge()
374 aggregated_cubes = iris.cube.CubeList()
375 for cube in cubes_at_time:
376 # If there was only a single data point at this time, then "number"
377 # will be a scalar coordinate. If so, make it a dimension coordinate
378 # to allow collapsing
379 if not cube.coord_dims("number"):
380 cube = iris.util.new_axis(cube, scalar_coord="number")
382 num_cases = cube.coord("number").points.size
383 num_cases_coord = iris.coords.AuxCoord(num_cases, long_name="num_cases")
384 cube.add_aux_coord(num_cases_coord)
386 # Do aggregation across the temporary "number" coordinate
387 if isinstance(aggregator, type(iris.analysis.PERCENTILE)):
388 cube = cube.collapsed("number", aggregator, percent=percentile)
389 else:
390 cube = cube.collapsed("number", aggregator)
392 # Now remove the "number" coordinate and its cell method
393 cube.remove_coord("number")
394 cell_method = iris.coords.CellMethod(aggregator.name(), coords="number")
395 cube = remove_cell_method(cube, cell_method)
397 aggregated_cubes.append(cube)
399 return aggregated_cubes