Coverage for src/CSET/operators/aggregate.py: 91%
184 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-11 15:11 +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)
203 return aggregatable_cubes
206def combine_obs_across_forecasts(cubes: CubeList) -> Cube:
207 """
208 Combine observation cubes from multiple forecast_reference_times.
210 Input:
211 CubeList of cubes with dimensions
213 (time, station)
215 Output:
216 Cube with dimensions
218 (forecast_reference_time,
219 forecast_period,
220 station)
222 where
224 time
226 becomes a 2D auxiliary coordinate attached to
228 (forecast_reference_time, forecast_period)
230 Only stations present in every forecast are retained.
231 All station metadata coordinates are preserved.
232 """
233 if len(cubes) < 2:
234 raise ValueError("Need at least two cubes")
236 common_stations = _get_common_stations(cubes)
237 station_lookup = _build_station_lookup(cubes, common_stations)
238 fp_hours = _generate_forecast_period(cubes)
239 return _make_aggregated_obs_cube(cubes, station_lookup, common_stations, fp_hours)
242def add_hour_coordinate(
243 cubes: iris.cube.Cube | iris.cube.CubeList,
244) -> iris.cube.Cube | iris.cube.CubeList:
245 """Add a category coordinate of hour of day to a Cube or CubeList.
247 Arguments
248 ---------
249 cubes: iris.cube.Cube | iris.cube.CubeList
250 Cube of any variable that has a time coordinate.
251 Note input Cube or CubeList items should only have 1 time dimension.
253 Returns
254 -------
255 cube: iris.cube.Cube
256 A Cube with an additional auxiliary coordinate of hour.
258 Notes
259 -----
260 This is a simple operator designed to be used prior to case aggregation for
261 histograms, Q-Q plots, and percentiles when aggregated by hour of day.
262 """
263 new_cubelist = iris.cube.CubeList()
264 for cube in iter_maybe(cubes):
265 # Add a category coordinate of hour into each cube.
266 iris.util.promote_aux_coord_to_dim_coord(cube, "time")
267 iris.coord_categorisation.add_hour(cube, "time", name="hour")
268 cube.coord("hour").units = "hours"
269 new_cubelist.append(cube)
271 if len(new_cubelist) == 1:
272 return new_cubelist[0]
273 else:
274 return new_cubelist
277def rolling_window_time_aggregation(
278 cubes: iris.cube.Cube | iris.cube.CubeList, method: str, window: int
279) -> iris.cube.Cube | iris.cube.CubeList:
280 """Aggregate a cube along the time dimension using a rolling window.
282 Arguments
283 ---------
284 cubes: iris.cube.Cube | iris.cube.CubeList
285 Cube or Cubelist of any variable to be aggregated over a rolling window
286 in time.
287 method: str
288 Type of aggregate i.e. method: 'MAX', getattr creates
289 iris.analysis.MAX, etc.
290 window: int
291 The rolling window size.
293 Returns
294 -------
295 cube: iris.cube.Cube | iris.cube.CubeList
296 A Cube or Cubelist of the rolling window aggregate. The Cubes will have
297 a time dimension that is reduced in size to the original cube by the
298 window size.
300 Notes
301 -----
302 This operator is designed to be used to help create daily maxima and minima
303 for any variable.
304 """
305 new_cubelist = iris.cube.CubeList()
306 for cube in iter_maybe(cubes):
307 # Use a rolling window in time to applied specified aggregation method
308 # over a specified window length.
309 window_cube = cube.rolling_window(
310 "time", getattr(iris.analysis, method), window
311 )
312 new_cubelist.append(window_cube)
314 if len(new_cubelist) == 1:
315 return new_cubelist[0]
316 else:
317 return new_cubelist
320def _get_common_stations(cubes: CubeList) -> list[str]:
321 """Find stations common to all cubes with complete data."""
322 if len(cubes) < 2:
323 raise ValueError(
324 f"Need at least two cubes to find common stations, but got {len(cubes)}"
325 )
327 valid_station_sets = []
329 for cb in cubes:
330 names = cb.coord("Station_Name").points
331 data = cb.data
333 # Handle masked and unmasked arrays
334 if np.ma.isMaskedArray(data): 334 ↛ 335line 334 didn't jump to line 335 because the condition on line 334 was never true
335 mask = np.ma.getmaskarray(data)
336 station_valid = ~np.any(mask, axis=0) & np.all(
337 np.isfinite(data.filled(np.nan)), axis=0
338 )
339 else:
340 station_valid = np.all(np.isfinite(data), axis=0)
341 valid_station_sets.append(set(names[station_valid]))
343 common_stations = sorted(set.intersection(*valid_station_sets))
345 logger.info(
346 "Retaining %s stations with complete observations", len(common_stations)
347 )
349 if not common_stations: 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true
350 raise ValueError(
351 "No stations with complete data in all forecast_reference_times"
352 )
353 return common_stations
356class StationLookup(NamedTuple):
357 """Container for station data, forecast reference times and validity times."""
359 data: list[np.ma.MaskedArray | np.ndarray]
360 frt_points: list[float]
361 time_points: list[np.ndarray]
364def _build_station_lookup(cubes: CubeList, common_stations: list[str]) -> StationLookup:
365 """Build station lookup table.
367 This makes a data structure which contains
369 data_common_stations:
370 the data, for example observed_temperature_at_screen_level, but only from the common stations.
372 frt_points:
373 the forecast reference times.
375 time_points:
376 the time points.
378 """
379 data_common_stations = []
380 frt_points = []
381 time_points = []
383 for cube in cubes:
384 # get the station names
385 station_names = cube.coord("Station_Name").points
387 # dict where {station_name0: idx0, station_name1: idx1, ... }
388 lookup = {station_name: idx for idx, station_name in enumerate(station_names)}
390 # the station indices of the common stations
391 station_indices = [lookup[name] for name in common_stations]
393 # get cube with only the common stations
394 cube_common_stations = cube[:, station_indices]
396 # get only the data from common stations cube
397 data_common_stations.append(cube_common_stations.data)
399 # get the forecast reference times
400 frt_points.append(cube.coord("forecast_reference_time").points[0])
402 # get the time points
403 time_points.append(cube.coord("time").points)
405 # Get the time axis length
406 ntime = len(time_points[0])
408 for t in time_points[1:]:
409 if len(t) != ntime: 409 ↛ 410line 409 didn't jump to line 410 because the condition on line 409 was never true
410 raise ValueError("Forecasts have different numbers of lead times")
411 return StationLookup(data_common_stations, frt_points, time_points)
414def _generate_forecast_period(cubes) -> np.ndarray:
415 """Generate forecast period values from time and forecast reference time.
417 Calculates the forecast period in hours for each time point, relative to
418 the first forecast reference time of the first cube.
419 """
420 time_coord = cubes[0].coord("time")
421 frt_coord = cubes[0].coord("forecast_reference_time")
423 frt_date = frt_coord.units.num2date(frt_coord.points[0])
425 fp_hours = []
426 for dt in time_coord.units.num2date(time_coord.points):
427 fp_hours.append((dt - frt_date).total_seconds() / 3600)
429 fp_hours = np.asarray(fp_hours)
431 return fp_hours
434def _make_aggregated_obs_cube(
435 cubes: CubeList,
436 station_lookup: StationLookup,
437 common_stations: list[str],
438 fp_hours: np.ndarray,
439) -> Cube:
440 time_coord = cubes[0].coord("time")
441 """Create new aggregated obs cube by rebuilding."""
442 # Stack data
443 data = np.stack(station_lookup.data, axis=0)
445 # shape (forecast_reference_time, forecast_period, station)
446 # Output coordinates
447 frt_out = DimCoord(
448 station_lookup.frt_points,
449 standard_name="forecast_reference_time",
450 units=cubes[0].coord("forecast_reference_time").units,
451 )
453 fp_out = DimCoord(
454 fp_hours,
455 standard_name="forecast_period",
456 units="hours",
457 )
459 station_out = DimCoord(
460 np.arange(len(common_stations)),
461 long_name="station",
462 )
464 cube_out = Cube(
465 data,
466 standard_name=cubes[0].standard_name,
467 long_name=cubes[0].long_name,
468 var_name=cubes[0].var_name,
469 units=cubes[0].units,
470 attributes=cubes[0].attributes.copy(),
471 dim_coords_and_dims=[
472 (frt_out, 0),
473 (fp_out, 1),
474 (station_out, 2),
475 ],
476 )
478 # Preserve station metadata coordinates
479 ref_cube = cubes[0]
481 ref_names = ref_cube.coord("Station_Name").points
483 ref_lookup = {name: idx for idx, name in enumerate(ref_names)}
485 common_idx = [ref_lookup[name] for name in common_stations]
487 # skip coords as we have awkward station and station_0 arbitrary monotonic arrays.
488 for coord in ref_cube.aux_coords:
489 dims = ref_cube.coord_dims(coord)
491 # only coords attached solely to station axis
492 if dims != (1,):
493 continue
495 values = coord.points[common_idx]
497 # verify same in every cube
498 for cube in cubes[1:]:
499 cube_names = cube.coord("Station_Name").points
501 cube_lookup = {name: idx for idx, name in enumerate(cube_names)}
503 idx = [cube_lookup[name] for name in common_stations]
505 other_values = cube.coord(coord.name()).points[idx]
507 if not np.array_equal( 507 ↛ 511line 507 didn't jump to line 511 because the condition on line 507 was never true
508 values,
509 other_values,
510 ):
511 raise ValueError(f"Station metadata differs for coord '{coord.name()}'")
513 aux = AuxCoord(
514 values,
515 standard_name=coord.standard_name,
516 long_name=coord.long_name,
517 var_name=coord.var_name,
518 units=coord.units,
519 attributes=coord.attributes.copy(),
520 )
522 cube_out.add_aux_coord(aux, (2,))
524 # Add valid-time auxiliary coord
525 time_2d = np.vstack(station_lookup.time_points)
527 cube_out.add_aux_coord(
528 AuxCoord(
529 time_2d,
530 standard_name="time",
531 units=time_coord.units,
532 ),
533 (0, 1),
534 )
536 return cube_out
539def _add_nref(cube: iris.cube.Cube):
540 """Retain information on number of forecast_reference_time inputs.
542 This preserves information on number of aggregated cases that can
543 otherwise be lost on subsequent calls to collapse functions.
544 """
545 nref = np.size(cube.coord("forecast_reference_time").points)
546 cube.coord("time").attributes["number_reference_times"] = nref
547 return cube
550def _aggregate_in_time_multiple_frt(
551 cube: iris.cube.Cube, method: str, interval: int
552) -> iris.cube.Cube:
553 """Aggregate a cube with multiple forecast reference times.
555 Aggregates each forecast cycle separately, then concatenates the results
556 back into a single cube along forecast_reference_time.
557 """
558 aggregated_cycles = iris.cube.CubeList()
559 for frt_cube in cube.slices_over("forecast_reference_time"):
560 iris.coord_categorisation.add_categorised_coord(
561 frt_cube,
562 "interval",
563 "time",
564 lambda coord, cell: cell // interval * interval,
565 )
566 agg = frt_cube.aggregated_by(
567 "interval",
568 getattr(iris.analysis, method),
569 )
570 agg.remove_coord("interval")
571 agg = iris.util.new_axis(
572 agg,
573 agg.coord("forecast_reference_time"),
574 )
575 aggregated_cycles.append(agg)
576 # Current approach is to remove time 2d auxcoord as it causes issues concatenating into a single cube.
577 # Time can be reconstructed from forecast_reference_time and forecast_period later on.
578 # as we can construct it if needed.
579 for cb in aggregated_cycles:
580 cb.remove_coord("time")
581 return aggregated_cycles.concatenate_cube()
584def _aggregate_in_time_single_frt(cube: iris.cube.Cube, method: str, interval: int):
585 """Aggregate a cube with one forecast reference time."""
586 iris.coord_categorisation.add_categorised_coord(
587 cube,
588 "interval",
589 "time",
590 lambda coord, cell: cell // interval * interval,
591 )
592 aggregated_cube = cube.aggregated_by(
593 "interval",
594 getattr(iris.analysis, method),
595 )
596 aggregated_cube.remove_coord("interval")
597 return aggregated_cube