diff --git a/src/CSET/operators/collapse.py b/src/CSET/operators/collapse.py index 75a4566ef..5b4bb435e 100644 --- a/src/CSET/operators/collapse.py +++ b/src/CSET/operators/collapse.py @@ -183,6 +183,80 @@ def collapse_by_hour_of_day( return collapsed_cube +def collapse_by_validity_time( + cube: iris.cube.Cube | iris.cube.CubeList, + method: str, + additional_percent: float = None, + **kwargs, +) -> iris.cube.Cube: + """Collapse a cube around validity time for multiple cases. + + First checks if the data can be aggregated easily. Then creates a new cube + by slicing over the time dimensions, removing the time dimensions, + re-merging the data, and creating a new time coordinate. It then collapses + by the new time coordinate for a specified method using the collapse + function. + + Arguments + --------- + cube: iris.cube.Cube | iris.cube.CubeList + Cube to collapse by validity time or CubeList that will be converted + to a cube before collapsing by validity time. + method: str + Type of collapse i.e. method: 'MEAN', 'MAX', 'MIN', 'MEDIAN', + 'PERCENTILE'. For 'PERCENTILE' the additional_percent must be specified. + + 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 multiple times. + cube_to_collapse = ensure_aggregatable_across_cases(cube) + # Convert to a cube that is split by validity time. + # Slice over cube by both time dimensions to create a CubeList. + new_cubelist = iris.cube.CubeList( + cube_to_collapse.slices_over(["forecast_period", "forecast_reference_time"]) + ) + # Remove forecast_period and forecast_reference_time coordinates. + for sub_cube in new_cubelist: + sub_cube.remove_coord("forecast_period") + sub_cube.remove_coord("forecast_reference_time") + # Create new CubeList by merging with unique = False to produce a validity + # time cube. + merged_list_1 = new_cubelist.merge(unique=False) + # Create a new "fake" coordinate and apply to each remaining cube to allow + # final merging to take place into a single cube. + equalised_validity_time = iris.coords.AuxCoord( + points=0, long_name="equalised_validity_time", units="1" + ) + for sub_cube, eq_valid_time in zip( + merged_list_1, range(len(merged_list_1)), strict=True + ): + sub_cube.add_aux_coord(equalised_validity_time.copy(points=eq_valid_time)) + + # Merge CubeList to create final cube. + final_cube = merged_list_1.merge_cube() + # Collapse over fake_time_coord to represent collapsing over validity time. + if method == "PERCENTILE": + collapsed_cube = collapse( + final_cube, + "equalised_validity_time", + method, + additional_percent=additional_percent, + ) + else: + collapsed_cube = collapse(final_cube, "equalised_validity_time", method) + return collapsed_cube + + # TODO # Collapse function that calculates means, medians etc across members of an # ensemble or stratified groups. Need to allow collapse over realisation diff --git a/tests/operators/test_collapse.py b/tests/operators/test_collapse.py index dcd7f528f..9938c3d6c 100644 --- a/tests/operators/test_collapse.py +++ b/tests/operators/test_collapse.py @@ -160,3 +160,44 @@ def test_collapse_by_lead_time_cube_list_percentile( rtol=1e-06, atol=1e-02, ) + + +def test_collapse_by_validity_time(long_forecast_multi_day): + """Reduce a dimension of a cube by validity time.""" + collapsed_cube = collapse.collapse_by_validity_time(long_forecast_multi_day, "MEAN") + expected_cube = "" + assert repr(collapsed_cube) == expected_cube + + +def test_collapse_by_validity_time_cubelist(long_forecast_many_cubes): + """Convert to cube and reduce a dimension by validity time.""" + collapsed_cube = collapse.collapse_by_validity_time( + long_forecast_many_cubes, "MEAN" + ) + expected_cube = "" + assert repr(collapsed_cube) == expected_cube + + +def test_collapse_by_validity_time_percentile(long_forecast_multi_day): + """Reduce by validity time with percentiles.""" + # Test successful collapsing by validity time. + collapsed_cube = collapse.collapse_by_validity_time( + long_forecast_multi_day, "PERCENTILE", additional_percent=[25, 75] + ) + expected_cube = "" + assert repr(collapsed_cube) == expected_cube + + +def test_collapse_by_validity_time_percentile_fail(long_forecast_multi_day): + """Test not specifying additional percent fails.""" + with pytest.raises(ValueError): + collapse.collapse_by_validity_time(long_forecast_multi_day, "PERCENTILE") + + +def test_collapse_by_validity_time_cubelist_percentile(long_forecast_many_cubes): + """Convert to cube and reduce by validity time with percentiles.""" + collapsed_cube = collapse.collapse_by_validity_time( + long_forecast_many_cubes, "PERCENTILE", additional_percent=[25, 75] + ) + expected_cube = "" + assert repr(collapsed_cube) == expected_cube