Implement unique selection reduction

climada/hazard/forecast.py

@@ -25,7 +25,7 @@
 import scipy.sparse as sparse
 
 from ..util.checker import size
-from ..util.forecast import Forecast
+from ..util.forecast import Forecast, reduce_unique_selection
 from .base import Hazard
 
 LOGGER = logging.getLogger(__name__)
@@ -106,7 +106,7 @@ def _check_sizes(self):
         size(exp_len=num_entries, var=self.member, var_name="Forecast.member")
         size(exp_len=num_entries, var=self.lead_time, var_name="Forecast.lead_time")
 
-    def _reduce_attrs(self, event_name: str):
+    def _reduce_attrs(self, event_name: str, **attrs):
         """
         Reduce the attributes of a HazardForecast to a single value.
 
@@ -123,32 +123,45 @@ def _reduce_attrs(self, event_name: str):
         event_name : str
             The event_name given to the reduced data.
         """
+
+        def unique_or_default(attr, default):
+            if len(unique := np.unique(getattr(self, attr))) == 1:
+                return unique.item(0)
+            return default
+
         reduced_attrs = {
-            "lead_time": np.array([np.timedelta64("NaT")]),
-            "member": np.array([-1]),
-            "event_id": np.array([0]),
-            "event_name": np.array([event_name]),
-            "date": np.array([0]),
-            "frequency": np.array([1]),
-            "orig": np.array([True]),
-        }
+            "lead_time": unique_or_default("lead_time", np.timedelta64("NaT")),
+            "member": unique_or_default("member", -1),
+            "event_id": unique_or_default("event_id", 1),
+            "event_name": unique_or_default("event_name", event_name),
+            "date": unique_or_default("date", 0),
+            "frequency": 1,
+            "orig": unique_or_default("orig", True),
+        } | attrs
+        reduced_attrs = {key: np.array([value]) for key, value in reduced_attrs.items()}
 
         return reduced_attrs
 
-    def min(self):
+    def min(self, dim=None):
         """
         Reduce the intensity and fraction of a HazardForecast to the minimum
         value.
 
         Parameters
         ----------
-        None
+        dim : str | None
+            Dimension to reduce over. If None, reduce over all data.
 
         Returns
         -------
         HazardForecast
             A HazardForecast object with the min intensity and fraction.
         """
+        if dim is not None:
+            return reduce_unique_selection(
+                self, values=getattr(self, dim), select=dim, reduce_attr="min"
+            )
+
         red_intensity = self.intensity.min(axis=0).tocsr()
         red_fraction = self.fraction.min(axis=0).tocsr()
         return HazardForecast(
@@ -162,20 +175,26 @@ def min(self):
             **self._reduce_attrs("min"),
         )
 
-    def max(self):
+    def max(self, dim=None):
         """
         Reduce the intensity and fraction of a HazardForecast to the maximum
         value.
 
         Parameters
         ----------
-        None
+        dim : str | None
+            Dimension to reduce over. If None, reduce over all data.
 
         Returns
         -------
         HazardForecast
             A HazardForecast object with the min intensity and fraction.
         """
+        if dim is not None:
+            return reduce_unique_selection(
+                self, values=getattr(self, dim), select=dim, reduce_attr="max"
+            )
+
         red_intensity = self.intensity.max(axis=0).tocsr()
         red_fraction = self.fraction.max(axis=0).tocsr()
         return HazardForecast(
@@ -189,19 +208,25 @@ def max(self):
             **self._reduce_attrs("max"),
         )
 
-    def mean(self):
+    def mean(self, dim=None):
         """
         Reduce the intensity and fraction of a HazardForecast to the mean value.
 
         Parameters
         ----------
-        None
+        dim : str | None
+            Dimension to reduce over. If None, reduce over all data.
 
         Returns
         -------
         HazardForecast
             A HazardForecast object with the min intensity and fraction.
         """
+        if dim is not None:
+            return reduce_unique_selection(
+                self, values=getattr(self, dim), select=dim, reduce_attr="mean"
+            )
+
         red_intensity = sparse.csr_matrix(self.intensity.mean(axis=0))
         red_fraction = sparse.csr_matrix(self.fraction.mean(axis=0))
         return HazardForecast(

climada/hazard/test/test_forecast.py

@@ -30,7 +30,7 @@
 from climada.hazard.test.test_base import hazard_kwargs
 
 
-@pytest.fixture
+@pytest.fixture(scope="module")
 def haz_kwargs():
     return hazard_kwargs()
 
@@ -40,14 +40,18 @@ def hazard(haz_kwargs):
     return Hazard(**haz_kwargs)
 
 
+# TODO: params should only apply to dimension reduction test
 @pytest.fixture
-def lead_time(haz_kwargs):
-    return pd.timedelta_range("1h", periods=len(haz_kwargs["event_id"])).to_numpy()
+def lead_time(request, haz_kwargs):
+    return pd.timedelta_range("1h", periods=len(haz_kwargs["event_id"]))
+    return np.tile(base_range.to_numpy(), request.param)
 
 
+# TODO: params should only apply to dimension reduction test
 @pytest.fixture
-def member(haz_kwargs):
+def member(request, haz_kwargs):
     return np.arange(len(haz_kwargs["event_id"]))
+    return np.tile(base_range, request.param)
 
 
 @pytest.fixture
@@ -254,7 +258,7 @@ def test_hazard_forecast_mean_min_max(haz_fc, attr):
     npt.assert_array_equal(np.isnat(haz_fcst_reduced.lead_time), [True])
     npt.assert_array_equal(haz_fcst_reduced.member, [-1])
     npt.assert_array_equal(haz_fcst_reduced.event_name, [attr])
-    npt.assert_array_equal(haz_fcst_reduced.event_id, [0])
+    npt.assert_array_equal(haz_fcst_reduced.event_id, [1])
     npt.assert_array_equal(haz_fcst_reduced.frequency, [1])
     npt.assert_array_equal(haz_fcst_reduced.date, [0])
     npt.assert_array_equal(haz_fcst_reduced.orig, [True])
@@ -284,7 +288,7 @@ def test_hazard_forecast_quantile(haz_fc, quantile):
     npt.assert_array_equal(
         haz_fcst_quantile.event_name, np.array([f"quantile_{quantile}"])
     )
-    npt.assert_array_equal(haz_fcst_quantile.event_id, np.array([0]))
+    npt.assert_array_equal(haz_fcst_quantile.event_id, np.array([1]))
     npt.assert_array_equal(haz_fcst_quantile.frequency, np.array([1]))
     npt.assert_array_equal(haz_fcst_quantile.date, np.array([0]))
     npt.assert_array_equal(haz_fcst_quantile.orig, np.array([True]))
@@ -311,8 +315,69 @@ def test_median(haz_fc):
     # check that attributes where reduced correctly
     npt.assert_array_equal(haz_fcst_median.member, np.array([-1]))
     npt.assert_array_equal(haz_fcst_median.lead_time, np.array([np.timedelta64("NaT")]))
-    npt.assert_array_equal(haz_fcst_median.event_id, np.array([0]))
+    npt.assert_array_equal(haz_fcst_median.event_id, np.array([1]))
     npt.assert_array_equal(haz_fcst_median.event_name, np.array(["median"]))
     npt.assert_array_equal(haz_fcst_median.frequency, np.array([1]))
     npt.assert_array_equal(haz_fcst_median.date, np.array([0]))
     npt.assert_array_equal(haz_fcst_median.orig, np.array([True]))
+
+
+@pytest.mark.parametrize("attr", ["min", "mean", "max"])
+def test_hazard_forecast_mean_min_max_member(haz_fc, attr):
+    """Check mean, min, and max methods for HazardForecast with dim argument"""
+
+    for dim, unique_vals in zip(
+        ["member", "lead_time"],
+        [np.unique(haz_fc.member), np.unique(haz_fc.lead_time)],
+    ):
+        haz_fcst_reduced = getattr(haz_fc, attr)(dim=dim)
+        # Assert sparse matrices
+        expected_intensity = []
+        expected_fraction = []
+        for val in unique_vals:
+            mask = getattr(haz_fc, dim) == val
+            expected_intensity.append(
+                getattr(haz_fc.intensity.todense()[mask], attr)(axis=0)
+            )
+            expected_fraction.append(
+                getattr(haz_fc.fraction.todense()[mask], attr)(axis=0)
+            )
+        npt.assert_array_equal(
+            haz_fcst_reduced.intensity.todense(),
+            np.vstack(expected_intensity),
+        )
+        npt.assert_array_equal(
+            haz_fcst_reduced.fraction.todense(),
+            np.vstack(expected_fraction),
+        )
+        # Check that attributes where reduced correctly
+        if dim == "lead_time":
+            npt.assert_array_equal(haz_fcst_reduced.member, np.unique(haz_fc.member))
+            npt.assert_array_equal(
+                haz_fcst_reduced.lead_time,
+                np.array([np.timedelta64("NaT")] * len(unique_vals)),
+            )
+        else:  # dim == "member"
+            npt.assert_array_equal(
+                haz_fcst_reduced.lead_time, np.unique(haz_fc.lead_time)
+            )
+            npt.assert_array_equal(
+                haz_fcst_reduced.member,
+                np.array([-1] * len(unique_vals)),
+            )
+        npt.assert_array_equal(
+            haz_fcst_reduced.event_name,
+            np.array([attr] * len(unique_vals)),
+        )
+        npt.assert_array_equal(haz_fcst_reduced.event_id, [0] * len(unique_vals))
+        npt.assert_array_equal(haz_fcst_reduced.frequency, [1] * len(unique_vals))
+        npt.assert_array_equal(haz_fcst_reduced.date, [0] * len(unique_vals))
+        npt.assert_array_equal(haz_fcst_reduced.orig, [True] * len(unique_vals))
+    # TODO add test in case no reduction happens (e.g., all values along dim are unique)
+
+
+def test_hazard_forecast_mean_max_min_dim_error(haz_fc):
+    """Check mean, min, and max methods for ImpactForecast with dim argument"""
+    for attr in ["min", "mean", "max"]:
+        with pytest.raises(AttributeError):
+            getattr(haz_fc, attr)(dim="invalid_dim")

climada/util/forecast.py

@@ -92,3 +92,33 @@ def idx_lead_time(self, lead_time: np.ndarray) -> np.ndarray:
         """
 
         return np.isin(self.lead_time, lead_time)
+
+
+def reduce_unique_selection(forecast, values, select, reduce_attr):
+    """
+    Reduce an attribute of a forecast object by selecting unique values
+    and performing an attribute reduction method.
+
+    Parameters
+    ----------
+    forecast : HazardForecast | ImpactForecast
+        Forecast object to reduce.
+    values : np.ndarray
+        Array of values for which to select and reduce the attribute.
+    select : str
+        Name of the attribute to select on (e.g. 'lead_time', 'member').
+    reduce_attr : str
+        Name of the attribute reduction method to call (e.g. 'min', 'mean').
+
+    Returns
+    -------
+    Forecast
+        Forecast object with the attribute reduced by the reduction method
+        and selected by the unique values.
+    """
+    return forecast.concat(
+        [
+            getattr(forecast.select(**{select: [val]}), reduce_attr)(dim=None)
+            for val in np.unique(values)
+        ]
+    )