Merge pull request #1067 from MetOffice/1030_collapse_by_validity_time

Collapse by validity time

Author: daflack

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

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 = "<iris 'Cube' of air_temperature / (K) (time: 145; grid_latitude: 3; grid_longitude: 3)>"
+    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 = "<iris 'Cube' of air_temperature / (K) (time: 145; grid_latitude: 3; grid_longitude: 3)>"
+    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 = "<iris 'Cube' of air_temperature / (K) (percentile_over_equalised_validity_time: 2; time: 145; grid_latitude: 3; grid_longitude: 3)>"
+    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 = "<iris 'Cube' of air_temperature / (K) (percentile_over_equalised_validity_time: 2; time: 145; grid_latitude: 3; grid_longitude: 3)>"
+    assert repr(collapsed_cube) == expected_cube