diff --git a/src/CSET/operators/collapse.py b/src/CSET/operators/collapse.py index 71bb03a8c..d536f7b9e 100644 --- a/src/CSET/operators/collapse.py +++ b/src/CSET/operators/collapse.py @@ -21,6 +21,8 @@ import iris.coord_categorisation import iris.cube +from CSET.operators._utils import ensure_aggregatable_across_cases + def collapse( cube: iris.cube.Cube, @@ -76,6 +78,54 @@ def collapse( return collapsed_cube +def collapse_by_lead_time( + cube: iris.cube.Cube | iris.cube.CubeList, + method: str, + additional_percent: float = None, + **kwargs, +) -> iris.cube.Cube: + """Collapse a cube around lead time for multiple cases. + + First checks if the data can be aggregated by lead time easily. Then + collapses by lead time for a specified method using the collapse function. + + Arguments + --------- + cube: iris.cube.Cube | iris.cube.CubeList + Cube to collapse by lead time or CubeList that will be converted + to a cube before collapsing by lead time. + method: str + Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN', + 'PERCENTILE' getattr creates iris.analysis.MEAN, etc For PERCENTILE + YAML file requires i.e. method: 'PERCENTILE' additional_percent: 90 + + Returns + ------- + cube: iris.cube.Cube + Single variable collapsed by lead time based on chosen method. + + Raises + ------ + ValueError + If additional_percent wasn't supplied while using PERCENTILE method. + """ + if method == "PERCENTILE" and additional_percent is None: + raise ValueError("Must specify additional_percent") + # Ensure the cube can be aggregated over mutlipe cases. + cube_to_collapse = ensure_aggregatable_across_cases(cube) + # Collapse by lead time. + if method == "PERCENTILE": + collapsed_cube = collapse( + cube_to_collapse, + "forecast_period", + method, + additional_percent=additional_percent, + ) + else: + collapsed_cube = collapse(cube_to_collapse, "forecast_period", method) + return collapsed_cube + + def collapse_by_hour_of_day( cube: iris.cube.Cube, method: str, diff --git a/tests/operators/test_collapse.py b/tests/operators/test_collapse.py index 7b28660a2..dcd7f528f 100644 --- a/tests/operators/test_collapse.py +++ b/tests/operators/test_collapse.py @@ -16,6 +16,7 @@ import iris import iris.cube +import numpy as np import pytest from CSET.operators import collapse @@ -29,6 +30,22 @@ def long_forecast() -> iris.cube.Cube: ) +@pytest.fixture() +def long_forecast_multi_day() -> iris.cube.Cube: + """Get long_forecast_multi_day to run tests on.""" + return iris.load_cube( + "tests/test_data/long_forecast_air_temp_multi_day.nc", "air_temperature" + ) + + +@pytest.fixture() +def long_forecast_many_cubes() -> iris.cube.Cube: + """Get long_forecast_may_cubes to run tests on.""" + return iris.load( + "tests/test_data/long_forecast_air_temp_fcst_*.nc", "air_temperature" + ) + + def test_collapse(cube): """Reduces dimension of cube.""" # Test collapsing a single coordinate. @@ -79,3 +96,67 @@ def test_collapse_by_hour_of_day_percentile(long_forecast): ) expected_cube = "<iris 'Cube' of air_temperature / (K) (percentile_over_hour: 2; -- : 24; grid_latitude: 3; grid_longitude: 3)>" assert repr(collapsed_cube) == expected_cube + + +def test_collapse_by_lead_time_single_cube(long_forecast_multi_day): + """Check cube collapse by lead time.""" + calculated_cube = collapse.collapse( + long_forecast_multi_day, "forecast_period", "MEAN" + ) + assert np.allclose( + calculated_cube.data, + collapse.collapse_by_lead_time(long_forecast_multi_day, "MEAN").data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_collapse_by_lead_time_cube_list( + long_forecast_multi_day, long_forecast_many_cubes +): + """Check CubeList is made into an aggregatable cube and collapses by lead time.""" + calculated_cube = collapse.collapse( + long_forecast_multi_day, "forecast_period", "MEAN" + ) + assert np.allclose( + calculated_cube.data, + collapse.collapse_by_lead_time(long_forecast_many_cubes, "MEAN").data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_collapse_by_lead_time_single_cube_percentile(long_forecast_multi_day): + """Check Cube collapse by lead time with percentiles.""" + calculated_cube = collapse.collapse( + long_forecast_multi_day, "forecast_period", "PERCENTILE", additional_percent=75 + ) + with pytest.raises(ValueError): + collapse.collapse_by_lead_time(long_forecast_multi_day, "PERCENTILE") + assert np.allclose( + calculated_cube.data, + collapse.collapse_by_lead_time( + long_forecast_multi_day, "PERCENTILE", additional_percent=75 + ).data, + rtol=1e-06, + atol=1e-02, + ) + + +def test_collapse_by_lead_time_cube_list_percentile( + long_forecast_multi_day, long_forecast_many_cubes +): + """Check CubeList is made into an aggregatable cube and collapses by lead time with percentiles.""" + calculated_cube = collapse.collapse( + long_forecast_multi_day, "forecast_period", "PERCENTILE", additional_percent=75 + ) + with pytest.raises(ValueError): + collapse.collapse_by_lead_time(long_forecast_many_cubes, "PERCENTILE") + assert np.allclose( + calculated_cube.data, + collapse.collapse_by_lead_time( + long_forecast_many_cubes, "PERCENTILE", additional_percent=75 + ).data, + rtol=1e-06, + atol=1e-02, + )