Adds operator for collapse_by_lead_time

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,

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,
+    )