Coverage for src/CSET/operators/aggregate.py: 55%
141 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-24 08:19 +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
211import iris
212import numpy as np
214from iris.coords import AuxCoord, DimCoord
215from iris.cube import Cube
218def combine_obs_across_forecasts(cubes):
219 """
220 Combine observation cubes from multiple forecast_reference_times.
222 Input:
223 CubeList of cubes with dimensions
225 (time, station)
227 Output:
228 Cube with dimensions
230 (forecast_reference_time,
231 forecast_period,
232 station)
234 where
236 time
238 becomes a 2D auxiliary coordinate attached to
240 (forecast_reference_time, forecast_period)
242 Only stations present in every forecast are retained.
243 All station metadata coordinates are preserved.
244 """
246 if len(cubes) < 2:
247 raise ValueError("Need at least two cubes")
249 # --------------------------------------------------------------
250 # Find common stations
251 # --------------------------------------------------------------
253 station_sets = []
255 for cube in cubes:
256 station_sets.append(
257 set(cube.coord("Station_Name").points)
258 )
260 common_stations = sorted(set.intersection(*station_sets))
262 if not common_stations:
263 raise ValueError(
264 "No stations common to all forecast_reference_times"
265 )
267 # --------------------------------------------------------------
268 # Build station lookup for every cube
269 # --------------------------------------------------------------
271 subset_data = []
272 frt_points = []
273 time_points = []
275 for cube in cubes:
277 names = cube.coord("Station_Name").points
279 lookup = {
280 name: idx
281 for idx, name in enumerate(names)
282 }
284 station_indices = [
285 lookup[name]
286 for name in common_stations
287 ]
289 subcube = cube[:, station_indices]
291 subset_data.append(subcube.data)
293 frt_points.append(
294 cube.coord("forecast_reference_time").points[0]
295 )
297 time_points.append(
298 cube.coord("time").points
299 )
301 # --------------------------------------------------------------
302 # Check all cubes have same time axis length
303 # --------------------------------------------------------------
305 ntime = len(time_points[0])
307 for t in time_points[1:]:
308 if len(t) != ntime:
309 raise ValueError(
310 "Forecasts have different numbers of lead times"
311 )
313 # --------------------------------------------------------------
314 # Generate forecast period
315 # --------------------------------------------------------------
317 time_coord = cubes[0].coord("time")
318 frt_coord = cubes[0].coord("forecast_reference_time")
320 frt_date = frt_coord.units.num2date(
321 frt_coord.points[0]
322 )
324 fp_hours = []
326 for dt in time_coord.units.num2date(
327 time_coord.points
328 ):
329 fp_hours.append(
330 (dt - frt_date).total_seconds() / 3600
331 )
333 fp_hours = np.asarray(fp_hours)
335 # --------------------------------------------------------------
336 # Stack data
337 # --------------------------------------------------------------
339 data = np.stack(subset_data, axis=0)
341 # shape:
342 #
343 # (forecast_reference_time,
344 # forecast_period,
345 # station)
347 # --------------------------------------------------------------
348 # Output coordinates
349 # --------------------------------------------------------------
351 frt_out = DimCoord(
352 frt_points,
353 standard_name="forecast_reference_time",
354 units=cubes[0].coord(
355 "forecast_reference_time"
356 ).units,
357 )
359 fp_out = DimCoord(
360 fp_hours,
361 standard_name="forecast_period",
362 units="hours",
363 )
365 station_out = DimCoord(
366 np.arange(len(common_stations)),
367 long_name="station",
368 )
370 cube_out = Cube(
371 data,
372 standard_name=cubes[0].standard_name,
373 long_name=cubes[0].long_name,
374 var_name=cubes[0].var_name,
375 units=cubes[0].units,
376 attributes=cubes[0].attributes.copy(),
377 dim_coords_and_dims=[
378 (frt_out, 0),
379 (fp_out, 1),
380 (station_out, 2),
381 ],
382 )
384 # --------------------------------------------------------------
385 # Preserve station metadata coordinates
386 # --------------------------------------------------------------
388 ref_cube = cubes[0]
390 ref_names = ref_cube.coord("Station_Name").points
392 ref_lookup = {
393 name: idx
394 for idx, name in enumerate(ref_names)
395 }
397 common_idx = [
398 ref_lookup[name]
399 for name in common_stations
400 ]
402 # skip coords as we have awkward station and station_0 arbritary monotonic arrays.
403 for coord in ref_cube.aux_coords:
405 try:
406 dims = ref_cube.coord_dims(coord)
407 except Exception:
408 continue
410 # only coords attached solely to station axis
411 if dims != (1,):
412 continue
414 values = coord.points[common_idx]
416 # verify same in every cube
417 for cube in cubes[1:]:
419 cube_names = cube.coord(
420 "Station_Name"
421 ).points
423 cube_lookup = {
424 name: idx
425 for idx, name in enumerate(cube_names)
426 }
428 idx = [
429 cube_lookup[name]
430 for name in common_stations
431 ]
433 other_values = cube.coord(
434 coord.name()
435 ).points[idx]
437 if not np.array_equal(
438 values,
439 other_values,
440 ):
441 raise ValueError(
442 f"Station metadata differs for "
443 f"coord '{coord.name()}'"
444 )
446 aux = AuxCoord(
447 values,
448 standard_name=coord.standard_name,
449 long_name=coord.long_name,
450 var_name=coord.var_name,
451 units=coord.units,
452 attributes=coord.attributes.copy(),
453 )
455 cube_out.add_aux_coord(aux, (2,))
457 # --------------------------------------------------------------
458 # Add valid-time auxiliary coord
459 # --------------------------------------------------------------
461 time_2d = np.vstack(time_points)
463 cube_out.add_aux_coord(
464 AuxCoord(
465 time_2d,
466 standard_name="time",
467 units=time_coord.units,
468 ),
469 (0, 1),
470 )
472 return cube_out
475def add_hour_coordinate(
476 cubes: iris.cube.Cube | iris.cube.CubeList,
477) -> iris.cube.Cube | iris.cube.CubeList:
478 """Add a category coordinate of hour of day to a Cube or CubeList.
480 Arguments
481 ---------
482 cubes: iris.cube.Cube | iris.cube.CubeList
483 Cube of any variable that has a time coordinate.
484 Note input Cube or CubeList items should only have 1 time dimension.
486 Returns
487 -------
488 cube: iris.cube.Cube
489 A Cube with an additional auxiliary coordinate of hour.
491 Notes
492 -----
493 This is a simple operator designed to be used prior to case aggregation for
494 histograms, Q-Q plots, and percentiles when aggregated by hour of day.
495 """
496 new_cubelist = iris.cube.CubeList()
497 for cube in iter_maybe(cubes):
498 # Add a category coordinate of hour into each cube.
499 iris.util.promote_aux_coord_to_dim_coord(cube, "time")
500 iris.coord_categorisation.add_hour(cube, "time", name="hour")
501 cube.coord("hour").units = "hours"
502 new_cubelist.append(cube)
504 if len(new_cubelist) == 1:
505 return new_cubelist[0]
506 else:
507 return new_cubelist
510def rolling_window_time_aggregation(
511 cubes: iris.cube.Cube | iris.cube.CubeList, method: str, window: int
512) -> iris.cube.Cube | iris.cube.CubeList:
513 """Aggregate a cube along the time dimension using a rolling window.
515 Arguments
516 ---------
517 cubes: iris.cube.Cube | iris.cube.CubeList
518 Cube or Cubelist of any variable to be aggregated over a rolling window
519 in time.
520 method: str
521 Type of aggregate i.e. method: 'MAX', getattr creates
522 iris.analysis.MAX, etc.
523 window: int
524 The rolling window size.
526 Returns
527 -------
528 cube: iris.cube.Cube | iris.cube.CubeList
529 A Cube or Cubelist of the rolling window aggregate. The Cubes will have
530 a time dimension that is reduced in size to the original cube by the
531 window size.
533 Notes
534 -----
535 This operator is designed to be used to help create daily maxima and minima
536 for any variable.
537 """
538 new_cubelist = iris.cube.CubeList()
539 for cube in iter_maybe(cubes):
540 # Use a rolling window in time to applied specified aggregation method
541 # over a specified window length.
542 window_cube = cube.rolling_window(
543 "time", getattr(iris.analysis, method), window
544 )
545 new_cubelist.append(window_cube)
547 if len(new_cubelist) == 1:
548 return new_cubelist[0]
549 else:
550 return new_cubelist