Coverage for src/CSET/operators/aggregate.py: 98%
77 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-04 08:32 +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 _add_nref(cube: iris.cube.Cube):
35 """Retain information on number of forecast_reference_time inputs.
37 This preserves information on number of aggregated cases that can
38 otherwise be lost on subsequent calls to collapse functions.
39 """
40 nref = np.size(cube.coord("forecast_reference_time").points)
41 cube.coord("time").attributes["number_reference_times"] = nref
42 return cube
45def time_aggregate(
46 cube: iris.cube.Cube,
47 method: str,
48 interval_iso: str,
49 **kwargs,
50) -> iris.cube.Cube:
51 """Aggregate cube by its time coordinate.
53 Aggregates similar (stash) fields in a cube for the specified coordinate and
54 using the method supplied. The aggregated cube will keep the coordinate and
55 add a further coordinate with the aggregated end time points.
57 Examples are: 1. Generating hourly or 6-hourly precipitation accumulations
58 given an interval for the new time coordinate.
60 We use the isodate class to convert ISO 8601 durations into time intervals
61 for creating a new time coordinate for aggregation.
63 We use the lambda function to pass coord and interval into the callable
64 category function in add_categorised to allow users to define their own
65 sub-daily intervals for the new time coordinate.
67 Arguments
68 ---------
69 cube: iris.cube.Cube
70 Cube to aggregate and iterate over one dimension
71 coordinate: str
72 Coordinate to aggregate over i.e. 'time', 'longitude',
73 'latitude','model_level_number'.
74 method: str
75 Type of aggregate i.e. method: 'SUM', getattr creates
76 iris.analysis.SUM, etc.
77 interval_iso: isodate timedelta ISO 8601 object i.e PT6H (6 hours), PT30M (30 mins)
78 Interval to aggregate over.
80 Returns
81 -------
82 cube: iris.cube.Cube
83 Single variable but several methods of aggregation
85 Raises
86 ------
87 ValueError
88 If the constraint doesn't produce a single cube containing a field.
89 """
90 # Duration of ISO timedelta.
91 timedelta = isodate.parse_duration(interval_iso)
93 # Convert interval format to whole hours.
94 interval = int(timedelta.total_seconds() / 3600)
96 # Add time categorisation overwriting hourly increment via lambda coord.
97 # https://scitools-iris.readthedocs.io/en/latest/_modules/iris/coord_categorisation.html
98 iris.coord_categorisation.add_categorised_coord(
99 cube, "interval", "time", lambda coord, cell: cell // interval * interval
100 )
102 # Aggregate cube using supplied method.
103 aggregated_cube = cube.aggregated_by("interval", getattr(iris.analysis, method))
104 aggregated_cube.remove_coord("interval")
105 return aggregated_cube
108def ensure_aggregatable_across_cases(
109 cubes: iris.cube.Cube | iris.cube.CubeList,
110) -> iris.cube.CubeList:
111 """Ensure a Cube or CubeList can be aggregated across multiple cases.
113 The cubes are grouped into buckets of compatible cubes, then each bucket is
114 converted into a single aggregatable cube with ``forecast_period`` and
115 ``forecast_reference_time`` dimension coordinates.
117 Arguments
118 ---------
119 cubes: iris.cube.Cube | iris.cube.CubeList
120 Each cube is checked to determine if it has the the necessary
121 dimensional coordinates to be aggregatable, being processed if needed.
123 Returns
124 -------
125 cubes: iris.cube.CubeList
126 A CubeList of time aggregatable cubes.
128 Raises
129 ------
130 ValueError
131 If any of the provided cubes cannot be made aggregatable.
133 Notes
134 -----
135 This is a simple operator designed to ensure that a Cube is aggregatable
136 across cases. If a CubeList is presented it will create an aggregatable Cube
137 from that list. Its functionality is for case study (or trial) aggregation
138 to ensure that the full dataset can be loaded as a single cube. This
139 functionality is particularly useful for percentiles, Q-Q plots, and
140 histograms.
142 The necessary dimension coordinates for a cube to be aggregatable are
143 ``forecast_period`` and ``forecast_reference_time``.
144 """
146 # Group compatible cubes.
147 class Buckets:
148 def __init__(self):
149 self.buckets = []
151 def add(self, cube: iris.cube.Cube):
152 """Add a cube into a bucket.
154 If the cube is compatible with an existing bucket it is added there.
155 Otherwise it gets its own bucket.
156 """
157 for bucket in self.buckets:
158 if bucket[0].is_compatible(cube):
159 bucket.append(cube)
160 return
161 self.buckets.append(iris.cube.CubeList([cube]))
163 def get_buckets(self) -> list[iris.cube.CubeList]:
164 return self.buckets
166 b = Buckets()
167 for cube in iter_maybe(cubes):
168 b.add(cube)
169 buckets = b.get_buckets()
171 logger.debug("Buckets:\n%s", "\n---\n".join(str(b) for b in buckets))
173 # Ensure each bucket is a single aggregatable cube.
174 aggregatable_cubes = iris.cube.CubeList()
175 for bucket in buckets:
176 # Single cubes that are already aggregatable won't need processing.
177 if len(bucket) == 1 and is_time_aggregatable(bucket[0]):
178 aggregatable_cube = bucket[0]
179 aggregatable_cube = _add_nref(aggregatable_cube)
180 aggregatable_cubes.append(aggregatable_cube)
181 continue
183 # Create an aggregatable cube from the provided CubeList.
184 to_merge = iris.cube.CubeList()
185 for cube in bucket:
186 try:
187 to_merge.extend(
188 cube.slices_over(["forecast_period", "forecast_reference_time"])
189 )
190 except iris.exceptions.CoordinateNotFoundError as err:
191 raise ValueError(
192 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.",
193 cube,
194 ) from err
195 aggregatable_cube = to_merge.merge_cube()
197 # Add attribute on number of forecast_reference_times
198 aggregatable_cube = _add_nref(aggregatable_cube)
200 # Verify cube is now aggregatable.
201 if not is_time_aggregatable(aggregatable_cube): 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 raise ValueError(
203 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.",
204 aggregatable_cube,
205 )
206 aggregatable_cubes.append(aggregatable_cube)
208 return aggregatable_cubes
211def add_hour_coordinate(
212 cubes: iris.cube.Cube | iris.cube.CubeList,
213) -> iris.cube.Cube | iris.cube.CubeList:
214 """Add a category coordinate of hour of day to a Cube or CubeList.
216 Arguments
217 ---------
218 cubes: iris.cube.Cube | iris.cube.CubeList
219 Cube of any variable that has a time coordinate.
220 Note input Cube or CubeList items should only have 1 time dimension.
222 Returns
223 -------
224 cube: iris.cube.Cube
225 A Cube with an additional auxiliary coordinate of hour.
227 Notes
228 -----
229 This is a simple operator designed to be used prior to case aggregation for
230 histograms, Q-Q plots, and percentiles when aggregated by hour of day.
231 """
232 new_cubelist = iris.cube.CubeList()
233 for cube in iter_maybe(cubes):
234 # Add a category coordinate of hour into each cube.
235 iris.util.promote_aux_coord_to_dim_coord(cube, "time")
236 iris.coord_categorisation.add_hour(cube, "time", name="hour")
237 cube.coord("hour").units = "hours"
238 new_cubelist.append(cube)
240 if len(new_cubelist) == 1:
241 return new_cubelist[0]
242 else:
243 return new_cubelist
246def rolling_window_time_aggregation(
247 cubes: iris.cube.Cube | iris.cube.CubeList, method: str, window: int
248) -> iris.cube.Cube | iris.cube.CubeList:
249 """Aggregate a cube along the time dimension using a rolling window.
251 Arguments
252 ---------
253 cubes: iris.cube.Cube | iris.cube.CubeList
254 Cube or Cubelist of any variable to be aggregated over a rolling window
255 in time.
256 method: str
257 Type of aggregate i.e. method: 'MAX', getattr creates
258 iris.analysis.MAX, etc.
259 window: int
260 The rolling window size.
262 Returns
263 -------
264 cube: iris.cube.Cube | iris.cube.CubeList
265 A Cube or Cubelist of the rolling window aggregate. The Cubes will have
266 a time dimension that is reduced in size to the original cube by the
267 window size.
269 Notes
270 -----
271 This operator is designed to be used to help create daily maxima and minima
272 for any variable.
273 """
274 new_cubelist = iris.cube.CubeList()
275 for cube in iter_maybe(cubes):
276 # Use a rolling window in time to applied specified aggregation method
277 # over a specified window length.
278 window_cube = cube.rolling_window(
279 "time", getattr(iris.analysis, method), window
280 )
281 new_cubelist.append(window_cube)
283 if len(new_cubelist) == 1:
284 return new_cubelist[0]
285 else:
286 return new_cubelist