Coverage for src/CSET/operators/aggregate.py: 94%
184 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-19 12:33 +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
18from typing import NamedTuple
20import iris
21import iris.analysis
22import iris.coord_categorisation
23import iris.cube
24import iris.exceptions
25import iris.util
26import isodate
27import numpy as np
28from iris.coords import AuxCoord, DimCoord
29from iris.cube import Cube, CubeList
31from CSET._common import iter_maybe
32from CSET.operators._utils import is_time_aggregatable
34logger = logging.getLogger(__name__)
37def time_aggregate(
38 cubes: iris.cube.Cube | iris.cube.CubeList,
39 method: str,
40 interval_iso: str,
41 **kwargs,
42) -> iris.cube.Cube | iris.cube.CubeList:
43 """Aggregate cube/cubes by its time coordinate.
45 Aggregates fields in a cube or cube list for the specified coordinate and
46 using the method supplied. The aggregated cube/cubelist will keep the coordinate and
47 add a coordinate with the aggregated end time points.
49 Can also handle multiple forecast reference times.
51 Examples include generating hourly or 6-hourly precipitation accumulations for precipitation, or maximum screen level temperature every 3 hours.
53 resampled_cubes = iris.cube.CubeList()
55 timedelta = isodate.parse_duration(interval_iso)
56 interval = int(timedelta.total_seconds() / 3600)
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.
67 interval_iso: str
68 A string containing a datetime timedelta for resampling over in hours, i.e. PT3H, PT24H.
70 Returns
71 -------
72 resampled_cubes: iris.cube.Cube | iris.cube.CubeList
73 Cube or CubeList containing aggregated cubes.
75 Raises
76 ------
77 ValueError
78 If the constraint doesn't produce a single cube containing a field.
79 """
80 # Return unchanged cubes if interval_iso is 0, to allow this operator to be used across multiple fields where only some will need resampling.
81 if interval_iso == "0": 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true
82 return cubes
84 cubes = iter_maybe(cubes)
86 resampled_cubes = iris.cube.CubeList()
88 timedelta = isodate.parse_duration(interval_iso)
89 interval = int(timedelta.total_seconds() / 3600)
91 for cube in cubes:
92 if cube.coord("forecast_reference_time").shape[0] > 1:
93 # Handle cubes with multiple forecast cycles.
94 aggregated_cube = _aggregate_in_time_multiple_frt(cube, method, interval)
95 else:
96 aggregated_cube = _aggregate_in_time_single_frt(cube, method, interval)
98 resampled_cubes.append(aggregated_cube)
99 if len(resampled_cubes) == 1: 99 ↛ 101line 99 didn't jump to line 101 because the condition on line 99 was always true
100 return resampled_cubes[0]
101 return resampled_cubes
104def ensure_aggregatable_across_cases(
105 cubes: iris.cube.Cube | iris.cube.CubeList,
106) -> iris.cube.CubeList:
107 """Ensure a Cube or CubeList can be aggregated across multiple cases.
109 The cubes are grouped into buckets of compatible cubes, then each bucket is
110 converted into a single aggregatable cube with ``forecast_period`` and
111 ``forecast_reference_time`` dimension coordinates.
113 Arguments
114 ---------
115 cubes: iris.cube.Cube | iris.cube.CubeList
116 Each cube is checked to determine if it has the necessary
117 dimensional coordinates to be aggregatable, being processed if needed.
119 Returns
120 -------
121 cubes: iris.cube.CubeList
122 A CubeList of time aggregatable cubes.
124 Raises
125 ------
126 ValueError
127 If any of the provided cubes cannot be made aggregatable.
129 Notes
130 -----
131 This is a simple operator designed to ensure that a Cube is aggregatable
132 across cases. If a CubeList is presented it will create an aggregatable Cube
133 from that list. Its functionality is for case study (or trial) aggregation
134 to ensure that the full dataset can be loaded as a single cube. This
135 functionality is particularly useful for percentiles, Q-Q plots, and
136 histograms.
138 The necessary dimension coordinates for a cube to be aggregatable are
139 ``forecast_period`` and ``forecast_reference_time``.
140 """
142 # Group compatible cubes.
143 class Buckets:
144 def __init__(self):
145 self.buckets = []
147 def add(self, cube: iris.cube.Cube):
148 """Add a cube into a bucket.
150 If the cube is compatible with an existing bucket it is added there.
151 Otherwise it gets its own bucket.
152 """
153 for bucket in self.buckets:
154 if bucket[0].is_compatible(cube):
155 bucket.append(cube)
156 return
157 self.buckets.append(iris.cube.CubeList([cube]))
159 def get_buckets(self) -> list[iris.cube.CubeList]:
160 return self.buckets
162 b = Buckets()
163 for cube in iter_maybe(cubes):
164 b.add(cube)
165 buckets = b.get_buckets()
167 logger.debug("Buckets:\n%s", "\n---\n".join(str(b) for b in buckets))
169 # Ensure each bucket is a single aggregatable cube.
170 aggregatable_cubes = iris.cube.CubeList()
171 for bucket in buckets:
172 # Single cubes that are already aggregatable won't need processing.
173 if len(bucket) == 1 and is_time_aggregatable(bucket[0]):
174 aggregatable_cube = bucket[0]
175 aggregatable_cube = _add_nref(aggregatable_cube)
176 aggregatable_cubes.append(aggregatable_cube)
177 continue
179 # Create an aggregatable cube from the provided CubeList.
180 to_merge = iris.cube.CubeList()
181 for cube in bucket:
182 try:
183 to_merge.extend(
184 cube.slices_over(["forecast_period", "forecast_reference_time"])
185 )
186 except iris.exceptions.CoordinateNotFoundError as err:
187 raise ValueError(
188 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.",
189 cube,
190 ) from err
191 aggregatable_cube = to_merge.merge_cube()
193 # Add attribute on number of forecast_reference_times
194 aggregatable_cube = _add_nref(aggregatable_cube)
196 # Verify cube is now aggregatable.
197 if not is_time_aggregatable(aggregatable_cube): 197 ↛ 198line 197 didn't jump to line 198 because the condition on line 197 was never true
198 raise ValueError(
199 "Cube should have 'forecast_period' and 'forecast_reference_time' dimension coordinates.",
200 aggregatable_cube,
201 )
202 aggregatable_cubes.append(aggregatable_cube)
204 return aggregatable_cubes
207def combine_obs_across_forecasts(cubes: CubeList) -> Cube:
208 """
209 Combine observation cubes from multiple forecast_reference_times.
211 Input:
212 CubeList of cubes with dimensions
214 (time, station)
216 Output:
217 Cube with dimensions
219 (forecast_reference_time,
220 forecast_period,
221 station)
223 where
225 time
227 becomes a 2D auxiliary coordinate attached to
229 (forecast_reference_time, forecast_period)
231 Only stations present in every forecast are retained.
232 All station metadata coordinates are preserved.
233 """
234 if len(cubes) < 2: 234 ↛ 235line 234 didn't jump to line 235 because the condition on line 234 was never true
235 raise ValueError("Need at least two cubes")
237 common_stations = _get_common_stations(cubes)
238 station_lookup = _build_station_lookup(cubes, common_stations)
239 fp_hours = _generate_forecast_period(cubes)
240 return _make_aggregated_obs_cube(cubes, station_lookup, common_stations, fp_hours)
243def add_hour_coordinate(
244 cubes: iris.cube.Cube | iris.cube.CubeList,
245) -> iris.cube.Cube | iris.cube.CubeList:
246 """Add a category coordinate of hour of day to a Cube or CubeList.
248 Arguments
249 ---------
250 cubes: iris.cube.Cube | iris.cube.CubeList
251 Cube of any variable that has a time coordinate.
252 Note input Cube or CubeList items should only have 1 time dimension.
254 Returns
255 -------
256 cube: iris.cube.Cube
257 A Cube with an additional auxiliary coordinate of hour.
259 Notes
260 -----
261 This is a simple operator designed to be used prior to case aggregation for
262 histograms, Q-Q plots, and percentiles when aggregated by hour of day.
263 """
264 new_cubelist = iris.cube.CubeList()
265 for cube in iter_maybe(cubes):
266 # Add a category coordinate of hour into each cube.
267 iris.util.promote_aux_coord_to_dim_coord(cube, "time")
268 iris.coord_categorisation.add_hour(cube, "time", name="hour")
269 cube.coord("hour").units = "hours"
270 new_cubelist.append(cube)
272 if len(new_cubelist) == 1:
273 return new_cubelist[0]
274 else:
275 return new_cubelist
278def rolling_window_time_aggregation(
279 cubes: iris.cube.Cube | iris.cube.CubeList, method: str, window: int
280) -> iris.cube.Cube | iris.cube.CubeList:
281 """Aggregate a cube along the time dimension using a rolling window.
283 Arguments
284 ---------
285 cubes: iris.cube.Cube | iris.cube.CubeList
286 Cube or Cubelist of any variable to be aggregated over a rolling window
287 in time.
288 method: str
289 Type of aggregate i.e. method: 'MAX', getattr creates
290 iris.analysis.MAX, etc.
291 window: int
292 The rolling window size.
294 Returns
295 -------
296 cube: iris.cube.Cube | iris.cube.CubeList
297 A Cube or Cubelist of the rolling window aggregate. The Cubes will have
298 a time dimension that is reduced in size to the original cube by the
299 window size.
301 Notes
302 -----
303 This operator is designed to be used to help create daily maxima and minima
304 for any variable.
305 """
306 new_cubelist = iris.cube.CubeList()
307 for cube in iter_maybe(cubes):
308 # Use a rolling window in time to applied specified aggregation method
309 # over a specified window length.
310 window_cube = cube.rolling_window(
311 "time", getattr(iris.analysis, method), window
312 )
313 new_cubelist.append(window_cube)
315 if len(new_cubelist) == 1:
316 return new_cubelist[0]
317 else:
318 return new_cubelist
321def _get_common_stations(cubes: CubeList) -> list[str]:
322 """Find stations common to all cubes with complete data."""
323 if len(cubes) < 2:
324 raise ValueError(
325 f"Need at least two cubes to find common stations, but got {len(cubes)}"
326 )
328 valid_station_sets = []
330 for cb in cubes:
331 names = cb.coord("Station_Name").points
332 data = cb.data
334 # Handle masked and unmasked arrays
335 if np.ma.isMaskedArray(data): 335 ↛ 336line 335 didn't jump to line 336 because the condition on line 335 was never true
336 mask = np.ma.getmaskarray(data)
337 station_valid = ~np.any(mask, axis=0) & np.all(
338 np.isfinite(data.filled(np.nan)), axis=0
339 )
340 else:
341 station_valid = np.all(np.isfinite(data), axis=0)
342 valid_station_sets.append(set(names[station_valid]))
344 common_stations = sorted(set.intersection(*valid_station_sets))
346 logger.info(
347 "Retaining %s stations with complete observations", len(common_stations)
348 )
350 if not common_stations: 350 ↛ 351line 350 didn't jump to line 351 because the condition on line 350 was never true
351 raise ValueError(
352 "No stations with complete data in all forecast_reference_times"
353 )
354 return common_stations
357class StationLookup(NamedTuple):
358 """Container for station data, forecast reference times and validity times."""
360 data: list[np.ma.MaskedArray | np.ndarray]
361 frt_points: list[float]
362 time_points: list[np.ndarray]
365def _build_station_lookup(cubes: CubeList, common_stations: list[str]) -> StationLookup:
366 """Build station lookup table.
368 This makes a data structure which contains
370 data_common_stations:
371 the data, for example observed_temperature_at_screen_level, but only from the common stations.
373 frt_points:
374 the forecast reference times.
376 time_points:
377 the time points.
379 """
380 data_common_stations = []
381 frt_points = []
382 time_points = []
384 for cube in cubes:
385 # get the station names
386 station_names = cube.coord("Station_Name").points
388 # dict where {station_name0: idx0, station_name1: idx1, ... }
389 lookup = {station_name: idx for idx, station_name in enumerate(station_names)}
391 # the station indices of the common stations
392 station_indices = [lookup[name] for name in common_stations]
394 # get cube with only the common stations
395 cube_common_stations = cube[:, station_indices]
397 # get only the data from common stations cube
398 data_common_stations.append(cube_common_stations.data)
400 # get the forecast reference times
401 frt_points.append(cube.coord("forecast_reference_time").points[0])
403 # get the time points
404 time_points.append(cube.coord("time").points)
406 # Get the time axis length
407 ntime = len(time_points[0])
409 for t in time_points[1:]:
410 if len(t) != ntime:
411 raise ValueError("Forecasts have different numbers of lead times")
412 return StationLookup(data_common_stations, frt_points, time_points)
415def _generate_forecast_period(cubes) -> np.ndarray:
416 """Generate forecast period values from time and forecast reference time.
418 Calculates the forecast period in hours for each time point, relative to
419 the first forecast reference time of the first cube.
420 """
421 time_coord = cubes[0].coord("time")
422 frt_coord = cubes[0].coord("forecast_reference_time")
424 frt_date = frt_coord.units.num2date(frt_coord.points[0])
426 fp_hours = []
427 for dt in time_coord.units.num2date(time_coord.points):
428 fp_hours.append((dt - frt_date).total_seconds() / 3600)
430 fp_hours = np.asarray(fp_hours)
432 return fp_hours
435def _make_aggregated_obs_cube(
436 cubes: CubeList,
437 station_lookup: StationLookup,
438 common_stations: list[str],
439 fp_hours: np.ndarray,
440) -> Cube:
441 time_coord = cubes[0].coord("time")
442 """Create new aggregated obs cube by rebuilding."""
443 # Stack data
444 data = np.stack(station_lookup.data, axis=0)
446 # shape (forecast_reference_time, forecast_period, station)
447 # Output coordinates
448 frt_out = DimCoord(
449 station_lookup.frt_points,
450 standard_name="forecast_reference_time",
451 units=cubes[0].coord("forecast_reference_time").units,
452 )
454 fp_out = DimCoord(
455 fp_hours,
456 standard_name="forecast_period",
457 units="hours",
458 )
460 station_out = DimCoord(
461 np.arange(len(common_stations)),
462 long_name="station",
463 )
465 cube_out = Cube(
466 data,
467 standard_name=cubes[0].standard_name,
468 long_name=cubes[0].long_name,
469 var_name=cubes[0].var_name,
470 units=cubes[0].units,
471 attributes=cubes[0].attributes.copy(),
472 dim_coords_and_dims=[
473 (frt_out, 0),
474 (fp_out, 1),
475 (station_out, 2),
476 ],
477 )
479 # Preserve station metadata coordinates
480 ref_cube = cubes[0]
482 ref_names = ref_cube.coord("Station_Name").points
484 ref_lookup = {name: idx for idx, name in enumerate(ref_names)}
486 common_idx = [ref_lookup[name] for name in common_stations]
488 # skip coords as we have awkward station and station_0 arbitrary monotonic arrays.
489 for coord in ref_cube.aux_coords:
490 dims = ref_cube.coord_dims(coord)
492 # only coords attached solely to station axis
493 if dims != (1,):
494 continue
496 values = coord.points[common_idx]
498 # verify same in every cube
499 for cube in cubes[1:]:
500 cube_names = cube.coord("Station_Name").points
502 cube_lookup = {name: idx for idx, name in enumerate(cube_names)}
504 idx = [cube_lookup[name] for name in common_stations]
506 other_values = cube.coord(coord.name()).points[idx]
508 if not np.array_equal( 508 ↛ 512line 508 didn't jump to line 512 because the condition on line 508 was never true
509 values,
510 other_values,
511 ):
512 raise ValueError(f"Station metadata differs for coord '{coord.name()}'")
514 aux = AuxCoord(
515 values,
516 standard_name=coord.standard_name,
517 long_name=coord.long_name,
518 var_name=coord.var_name,
519 units=coord.units,
520 attributes=coord.attributes.copy(),
521 )
523 cube_out.add_aux_coord(aux, (2,))
525 # Add valid-time auxiliary coord
526 time_2d = np.vstack(station_lookup.time_points)
528 cube_out.add_aux_coord(
529 AuxCoord(
530 time_2d,
531 standard_name="time",
532 units=time_coord.units,
533 ),
534 (0, 1),
535 )
537 return cube_out
540def _add_nref(cube: iris.cube.Cube):
541 """Retain information on number of forecast_reference_time inputs.
543 This preserves information on number of aggregated cases that can
544 otherwise be lost on subsequent calls to collapse functions.
545 """
546 nref = np.size(cube.coord("forecast_reference_time").points)
547 cube.coord("time").attributes["number_reference_times"] = nref
548 return cube
551def _aggregate_in_time_multiple_frt(
552 cube: iris.cube.Cube, method: str, interval: int
553) -> iris.cube.Cube:
554 """Aggregate a cube with multiple forecast reference times.
556 Aggregates each forecast cycle separately, then concatenates the results
557 back into a single cube along forecast_reference_time.
558 """
559 aggregated_cycles = iris.cube.CubeList()
560 for frt_cube in cube.slices_over("forecast_reference_time"):
561 iris.coord_categorisation.add_categorised_coord(
562 frt_cube,
563 "interval",
564 "time",
565 lambda coord, cell: cell // interval * interval,
566 )
567 agg = frt_cube.aggregated_by(
568 "interval",
569 getattr(iris.analysis, method),
570 )
571 agg.remove_coord("interval")
572 agg = iris.util.new_axis(
573 agg,
574 agg.coord("forecast_reference_time"),
575 )
576 aggregated_cycles.append(agg)
577 # Current approach is to remove time 2d auxcoord as it causes issues concatenating into a single cube.
578 # Time can be reconstructed from forecast_reference_time and forecast_period later on.
579 # as we can construct it if needed.
580 for cb in aggregated_cycles:
581 cb.remove_coord("time")
582 return aggregated_cycles.concatenate_cube()
585def _aggregate_in_time_single_frt(cube: iris.cube.Cube, method: str, interval: int):
586 """Aggregate a cube with one forecast reference time."""
587 iris.coord_categorisation.add_categorised_coord(
588 cube,
589 "interval",
590 "time",
591 lambda coord, cell: cell // interval * interval,
592 )
593 aggregated_cube = cube.aggregated_by(
594 "interval",
595 getattr(iris.analysis, method),
596 )
597 aggregated_cube.remove_coord("interval")
598 return aggregated_cube