diff --git a/src/CSET/operators/aggregate.py b/src/CSET/operators/aggregate.py index 451bb0eed..31be60c22 100644 --- a/src/CSET/operators/aggregate.py +++ b/src/CSET/operators/aggregate.py @@ -31,31 +31,21 @@ logger = logging.getLogger(__name__) -def _add_nref(cube: iris.cube.Cube): - """Retain information on number of forecast_reference_time inputs. - - This preserves information on number of aggregated cases that can - otherwise be lost on subsequent calls to collapse functions. - """ - nref = np.size(cube.coord("forecast_reference_time").points) - cube.coord("time").attributes["number_reference_times"] = nref - return cube - - def time_aggregate( - cube: iris.cube.Cube, + cubes: iris.cube.Cube | iris.cube.CubeList, method: str, interval_iso: str, **kwargs, -) -> iris.cube.Cube: - """Aggregate cube by its time coordinate. +) -> iris.cube.Cube | iris.cube.CubeList: + """Aggregate cube/cubes by its time coordinate. - Aggregates similar (stash) fields in a cube for the specified coordinate and - using the method supplied. The aggregated cube will keep the coordinate and - add a further coordinate with the aggregated end time points. + Aggregates fields in a cube or cube list for the specified coordinate and + using the method supplied. The aggregated cube/cubelist will keep the coordinate and + add a coordinate with the aggregated end time points. - Examples are: 1. Generating hourly or 6-hourly precipitation accumulations - given an interval for the new time coordinate. + Can also handle multiple forecast reference times. + + Examples include generating hourly or 6-hourly precipitation accumulations for precipitation, or maximum screen level temperature every 3 hours. We use the isodate class to convert ISO 8601 durations into time intervals for creating a new time coordinate for aggregation. @@ -66,43 +56,50 @@ def time_aggregate( Arguments --------- - cube: iris.cube.Cube - Cube to aggregate and iterate over one dimension - coordinate: str - Coordinate to aggregate over i.e. 'time', 'longitude', - 'latitude','model_level_number'. + cubes: iris.cube.Cube | iris.cube.CubeList + Cube or CubeList to aggregate and iterate over one dimension method: str Type of aggregate i.e. method: 'SUM', getattr creates iris.analysis.SUM, etc. interval_iso: isodate timedelta ISO 8601 object i.e PT6H (6 hours), PT30M (30 mins) Interval to aggregate over. + interval_iso: str + A string containing a datetime timedelta for resampling over in hours, i.e. PT3H, PT24H. Returns ------- - cube: iris.cube.Cube - Single variable but several methods of aggregation + resampled_cubes: iris.cube.Cube | iris.cube.CubeList + Cube or CubeList containing aggregated cubes. Raises ------ ValueError If the constraint doesn't produce a single cube containing a field. """ - # Duration of ISO timedelta. - timedelta = isodate.parse_duration(interval_iso) + # Return unchanged cubes if interval_iso is 0, to allow this operator to be used across multiple fields where only some will need resampling. + if interval_iso == "0": + return cubes + + cubes = iter_maybe(cubes) + + resampled_cubes = iris.cube.CubeList() - # Convert interval format to whole hours. + timedelta = isodate.parse_duration(interval_iso) interval = int(timedelta.total_seconds() / 3600) - # Add time categorisation overwriting hourly increment via lambda coord. - # https://scitools-iris.readthedocs.io/en/latest/_modules/iris/coord_categorisation.html - iris.coord_categorisation.add_categorised_coord( - cube, "interval", "time", lambda coord, cell: cell // interval * interval - ) + for cube in cubes: + if cube.coord("forecast_reference_time").shape[0] > 1: + # Handle cubes with multiple forecast cycles. + aggregated_cube = _aggregate_in_time_multiple_frt_cube( + cube, method, interval + ) + else: + aggregated_cube = _aggregate_in_time_single_frt(cube, method, interval) - # Aggregate cube using supplied method. - aggregated_cube = cube.aggregated_by("interval", getattr(iris.analysis, method)) - aggregated_cube.remove_coord("interval") - return aggregated_cube + resampled_cubes.append(aggregated_cube) + if len(resampled_cubes) == 1: + return resampled_cubes[0] + return resampled_cubes def ensure_aggregatable_across_cases( @@ -117,7 +114,7 @@ def ensure_aggregatable_across_cases( Arguments --------- cubes: iris.cube.Cube | iris.cube.CubeList - Each cube is checked to determine if it has the the necessary + Each cube is checked to determine if it has the necessary dimensional coordinates to be aggregatable, being processed if needed. Returns @@ -284,3 +281,65 @@ def rolling_window_time_aggregation( return new_cubelist[0] else: return new_cubelist + + +def _add_nref(cube: iris.cube.Cube): + """Retain information on number of forecast_reference_time inputs. + + This preserves information on number of aggregated cases that can + otherwise be lost on subsequent calls to collapse functions. + """ + nref = np.size(cube.coord("forecast_reference_time").points) + cube.coord("time").attributes["number_reference_times"] = nref + return cube + + +def _aggregate_in_time_multiple_frt_cube( + cube: iris.cube.Cube, method: str, interval: int +) -> iris.cube.Cube: + """Aggregate a cube with multiple forecast reference times. + + Aggregates each forecast cycle separately, then concatenates the results + back into a single cube along forecast_reference_time. + """ + aggregated_cycles = iris.cube.CubeList() + for frt_cube in cube.slices_over("forecast_reference_time"): + iris.coord_categorisation.add_categorised_coord( + frt_cube, + "interval", + "time", + lambda coord, cell: cell // interval * interval, + ) + agg = frt_cube.aggregated_by( + "interval", + getattr(iris.analysis, method), + ) + agg.remove_coord("interval") + agg = iris.util.new_axis( + agg, + agg.coord("forecast_reference_time"), + ) + aggregated_cycles.append(agg) + # Current approach is to remove time 2d auxcoord as it causes issues concatenating into a single cube. + # Time can be reconstructed from forecast_reference_time and forecast_period later on. + # as we can construct it if needed. + for cb in aggregated_cycles: + cb.remove_coord("time") + return aggregated_cycles.concatenate_cube() + + +def _aggregate_in_time_single_frt(cube: iris.cube.Cube, method: str, interval: int): + # Aggregate a cube with one forecast reference time. + + iris.coord_categorisation.add_categorised_coord( + cube, + "interval", + "time", + lambda coord, cell: cell // interval * interval, + ) + aggregated_cube = cube.aggregated_by( + "interval", + getattr(iris.analysis, method), + ) + aggregated_cube.remove_coord("interval") + return aggregated_cube diff --git a/tests/operators/test_aggregate.py b/tests/operators/test_aggregate.py index ce3e97f7f..d7af3948e 100644 --- a/tests/operators/test_aggregate.py +++ b/tests/operators/test_aggregate.py @@ -16,13 +16,14 @@ import iris import iris.cube +import isodate import numpy as np import pytest from CSET.operators import aggregate -def test_aggregate(cube): +def test_aggregate(cube, long_forecast_multi_day): """Aggregate time to 2 hour intervals.""" # Set test interval to 2 hours. interval = "PT2H" @@ -41,6 +42,52 @@ def test_aggregate(cube): ) +def test_aggregate_multi_frt(long_forecast_multi_day): + """Aggregate time to 1 hour intervals.""" + # Set test interval to 1 hours. + + interval = "PT1H" + + aggregated_cube = aggregate.time_aggregate( + long_forecast_multi_day.copy(), method="SUM", interval_iso=interval + ) + # Check if number of coords on aggregated cube one less than original cube, + # as the time aux coord is removed. + assert len(aggregated_cube.coords()) + 1 == len(long_forecast_multi_day.coords()) + + +def test_aggregate_in_time_single_frt(cube): + """Aggregate time to 2 hour intervals.""" + # Set test interval to 2 hours. + interval_iso = "PT2H" + timedelta = isodate.parse_duration(interval_iso) + interval = int(timedelta.total_seconds() / 3600) + + aggregated_cube = aggregate._aggregate_in_time_single_frt( + cube.copy(), method="SUM", interval=interval + ) + + # Check if number of coords on aggregated cube is same as original cube. + assert len(aggregated_cube.coords()) == len(cube.coords()), ( + "aggregated cube does not have additional aux coordinate" + ) + + +def test_aggregate_in_time_multi_frt(long_forecast_multi_day): + """Aggregate time to 1 hour intervals.""" + # Set test interval to 1 hour. + interval_iso = "PT1H" + timedelta = isodate.parse_duration(interval_iso) + interval = int(timedelta.total_seconds() / 3600) + + aggregated_cube = aggregate._aggregate_in_time_multiple_frt_cube( + long_forecast_multi_day.copy(), method="SUM", interval=interval + ) + # Check if number of coords on aggregated cube one less than original cube, + # as the time aux coord is removed. + assert len(aggregated_cube.coords()) + 1 == len(long_forecast_multi_day.coords()) + + def test_ensure_aggregatable_across_cases_true_aggregatable_cube( long_forecast_multi_day, ):