Coverage for src/CSET/operators/aggregate.py: 83%
102 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-11 09:22 +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 is_time_aggregatable
31logger = logging.getLogger(__name__)
34def time_aggregate(
35 cubes: iris.cube.Cube | iris.cube.CubeList,
36 method: str,
37 interval_iso: str,
38 **kwargs,
39) -> iris.cube.Cube | iris.cube.CubeList:
40 """Aggregate cube/cubes by its time coordinate.
42 Aggregates similar (stash) fields in a cube or cube list for the specified coordinate and
43 using the method supplied. The aggregated cube/cubelist will keep the coordinate and
44 add a coordinate with the aggregated end time points.
46 Can also handle multiple forecast reference times.
48 Examples are: 1. Generating hourly or 6-hourly precipitation accumulations
49 given an interval for the new time coordinate.
51 We use the isodate class to convert ISO 8601 durations into time intervals
52 for creating a new time coordinate for aggregation.
54 We use the lambda function to pass coord and interval into the callable
55 category function in add_categorised to allow users to define their own
56 sub-daily intervals for the new time coordinate.
58 Arguments
59 ---------
60 cubes: iris.cube.Cube | iris.cube.CubeList
61 Cube or CubeList to aggregate and iterate over one dimension
62 method: str
63 Type of aggregate i.e. method: 'SUM', getattr creates
64 iris.analysis.SUM, etc.
65 interval_iso: isodate timedelta ISO 8601 object i.e PT6H (6 hours), PT30M (30 mins)
66 Interval to aggregate over.
68 Returns
69 -------
70 resampled_cubes: iris.cube.Cube | iris.cube.CubeList
71 Single variable but several methods of aggregation
73 Raises
74 ------
75 ValueError
76 If the constraint doesn't produce a single cube containing a field.
77 """
78 if interval_iso == "0": 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true
79 return cubes
81 if isinstance(cubes, iris.cube.Cube): 81 ↛ 84line 81 didn't jump to line 84 because the condition on line 81 was always true
82 cubes = iris.cube.CubeList([cubes])
84 resampled_cubes = iris.cube.CubeList()
86 timedelta = isodate.parse_duration(interval_iso)
87 interval = int(timedelta.total_seconds() / 3600)
89 for cube in cubes:
90 # Handle cubes with multiple forecast cycles.
91 if cube.coord("forecast_reference_time").shape[0] > 1: 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true
92 aggregated_cube = _aggregate_multi_frt_cube(cube, method, interval)
93 else:
94 aggregated_cube = _aggregate_by_interval(cube, method, interval)
96 resampled_cubes.append(aggregated_cube)
97 if len(resampled_cubes) == 1: 97 ↛ 99line 97 didn't jump to line 99 because the condition on line 97 was always true
98 return resampled_cubes[0]
99 return resampled_cubes
102def ensure_aggregatable_across_cases(
103 cubes: iris.cube.Cube | iris.cube.CubeList,
104) -> iris.cube.CubeList:
105 """Ensure a Cube or CubeList can be aggregated across multiple cases.
107 The cubes are grouped into buckets of compatible cubes, then each bucket is
108 converted into a single aggregatable cube with ``forecast_period`` and
109 ``forecast_reference_time`` dimension coordinates.
111 Arguments
112 ---------
113 cubes: iris.cube.Cube | iris.cube.CubeList
114 Each cube is checked to determine if it has the necessary
115 dimensional coordinates to be aggregatable, being processed if needed.
117 Returns
118 -------
119 cubes: iris.cube.CubeList
120 A CubeList of time aggregatable cubes.
122 Raises
123 ------
124 ValueError
125 If any of the provided cubes cannot be made aggregatable.
127 Notes
128 -----
129 This is a simple operator designed to ensure that a Cube is aggregatable
130 across cases. If a CubeList is presented it will create an aggregatable Cube
131 from that list. Its functionality is for case study (or trial) aggregation
132 to ensure that the full dataset can be loaded as a single cube. This
133 functionality is particularly useful for percentiles, Q-Q plots, and
134 histograms.
136 The necessary dimension coordinates for a cube to be aggregatable are
137 ``forecast_period`` and ``forecast_reference_time``.
138 """
140 # Group compatible cubes.
141 class Buckets:
142 def __init__(self):
143 self.buckets = []
145 def add(self, cube: iris.cube.Cube):
146 """Add a cube into a bucket.
148 If the cube is compatible with an existing bucket it is added there.
149 Otherwise it gets its own bucket.
150 """
151 for bucket in self.buckets:
152 if bucket[0].is_compatible(cube):
153 bucket.append(cube)
154 return
155 self.buckets.append(iris.cube.CubeList([cube]))
157 def get_buckets(self) -> list[iris.cube.CubeList]:
158 return self.buckets
160 b = Buckets()
161 for cube in iter_maybe(cubes):
162 b.add(cube)
163 buckets = b.get_buckets()
165 logger.debug("Buckets:\n%s", "\n---\n".join(str(b) for b in buckets))
167 # Ensure each bucket is a single aggregatable cube.
168 aggregatable_cubes = iris.cube.CubeList()
169 for bucket in buckets:
170 # Single cubes that are already aggregatable won't need processing.
171 if len(bucket) == 1 and is_time_aggregatable(bucket[0]):
172 aggregatable_cube = bucket[0]
173 aggregatable_cube = _add_nref(aggregatable_cube)
174 aggregatable_cubes.append(aggregatable_cube)
175 continue
177 # Create an aggregatable cube from the provided CubeList.
178 to_merge = iris.cube.CubeList()
179 for cube in bucket:
180 try:
181 to_merge.extend(
182 cube.slices_over(["forecast_period", "forecast_reference_time"])
183 )
184 except iris.exceptions.CoordinateNotFoundError as err:
185 raise ValueError(
186 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.",
187 cube,
188 ) from err
189 aggregatable_cube = to_merge.merge_cube()
191 # Add attribute on number of forecast_reference_times
192 aggregatable_cube = _add_nref(aggregatable_cube)
194 # Verify cube is now aggregatable.
195 if not is_time_aggregatable(aggregatable_cube): 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 raise ValueError(
197 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.",
198 aggregatable_cube,
199 )
200 aggregatable_cubes.append(aggregatable_cube)
202 return aggregatable_cubes
205def add_hour_coordinate(
206 cubes: iris.cube.Cube | iris.cube.CubeList,
207) -> iris.cube.Cube | iris.cube.CubeList:
208 """Add a category coordinate of hour of day to a Cube or CubeList.
210 Arguments
211 ---------
212 cubes: iris.cube.Cube | iris.cube.CubeList
213 Cube of any variable that has a time coordinate.
214 Note input Cube or CubeList items should only have 1 time dimension.
216 Returns
217 -------
218 cube: iris.cube.Cube
219 A Cube with an additional auxiliary coordinate of hour.
221 Notes
222 -----
223 This is a simple operator designed to be used prior to case aggregation for
224 histograms, Q-Q plots, and percentiles when aggregated by hour of day.
225 """
226 new_cubelist = iris.cube.CubeList()
227 for cube in iter_maybe(cubes):
228 # Add a category coordinate of hour into each cube.
229 iris.util.promote_aux_coord_to_dim_coord(cube, "time")
230 iris.coord_categorisation.add_hour(cube, "time", name="hour")
231 cube.coord("hour").units = "hours"
232 new_cubelist.append(cube)
234 if len(new_cubelist) == 1:
235 return new_cubelist[0]
236 else:
237 return new_cubelist
240def rolling_window_time_aggregation(
241 cubes: iris.cube.Cube | iris.cube.CubeList, method: str, window: int
242) -> iris.cube.Cube | iris.cube.CubeList:
243 """Aggregate a cube along the time dimension using a rolling window.
245 Arguments
246 ---------
247 cubes: iris.cube.Cube | iris.cube.CubeList
248 Cube or Cubelist of any variable to be aggregated over a rolling window
249 in time.
250 method: str
251 Type of aggregate i.e. method: 'MAX', getattr creates
252 iris.analysis.MAX, etc.
253 window: int
254 The rolling window size.
256 Returns
257 -------
258 cube: iris.cube.Cube | iris.cube.CubeList
259 A Cube or Cubelist of the rolling window aggregate. The Cubes will have
260 a time dimension that is reduced in size to the original cube by the
261 window size.
263 Notes
264 -----
265 This operator is designed to be used to help create daily maxima and minima
266 for any variable.
267 """
268 new_cubelist = iris.cube.CubeList()
269 for cube in iter_maybe(cubes):
270 # Use a rolling window in time to applied specified aggregation method
271 # over a specified window length.
272 window_cube = cube.rolling_window(
273 "time", getattr(iris.analysis, method), window
274 )
275 new_cubelist.append(window_cube)
277 if len(new_cubelist) == 1:
278 return new_cubelist[0]
279 else:
280 return new_cubelist
283def _add_nref(cube: iris.cube.Cube):
284 """Retain information on number of forecast_reference_time inputs.
286 This preserves information on number of aggregated cases that can
287 otherwise be lost on subsequent calls to collapse functions.
288 """
289 nref = np.size(cube.coord("forecast_reference_time").points)
290 cube.coord("time").attributes["number_reference_times"] = nref
291 return cube
294def _aggregate_multi_frt_cube(
295 cube: iris.cube.Cube, method: str, interval: int
296) -> iris.cube.Cube:
297 """Aggregate a cube with multiple forecast reference times.
299 Aggregates each forecast cycle separately, then concatenates the results
300 back into a single cube along forecast_reference_time.
301 """
302 aggregated_cycles = iris.cube.CubeList()
303 for frt_cube in cube.slices_over("forecast_reference_time"):
304 iris.coord_categorisation.add_categorised_coord(
305 frt_cube,
306 "interval",
307 "time",
308 lambda coord, cell: cell // interval * interval,
309 )
310 agg = frt_cube.aggregated_by(
311 "interval",
312 getattr(iris.analysis, method),
313 )
314 agg.remove_coord("interval")
315 agg = iris.util.new_axis(
316 agg,
317 agg.coord("forecast_reference_time"),
318 )
319 aggregated_cycles.append(agg)
320 # 2d auxtime causes issues concatenating. Solution is to nuke it. Do we need it downstream?
321 # as we can construct it if needed.
322 for cb in aggregated_cycles:
323 cb.remove_coord("time")
324 return aggregated_cycles.concatenate_cube()
327def _aggregate_by_interval(cube: iris.cube.Cube, method: str, interval: int):
328 iris.coord_categorisation.add_categorised_coord(
329 cube,
330 "interval",
331 "time",
332 lambda coord, cell: cell // interval * interval,
333 )
334 aggregated_cube = cube.aggregated_by(
335 "interval",
336 getattr(iris.analysis, method),
337 )
338 aggregated_cube.remove_coord("interval")
339 return aggregated_cube