Add forecast tutorial (#1193)

* Add impact forecast tutorial
* Add 'dim' argument to reductions
* Update tests

---------

Co-authored-by: Lukas Riedel <[email protected]>
Co-authored-by: luseverin <[email protected]>
Co-authored-by: Evelyn Mühlhofer <[email protected]>
Co-authored-by: Chahan Kropf <[email protected]>

Author: 5 people

climada/engine/impact_calc.py

@@ -234,7 +234,7 @@ def _return_impact(self, imp_mat_gen, save_mat):
             )
             if isinstance(self.hazard, HazardForecast):
                 eai_exp = np.full_like(eai_exp, np.nan, dtype=eai_exp.dtype)
-                aai_agg = np.full_like(aai_agg, np.nan, dtype=aai_agg.dtype)
+                aai_agg = np.nan
                 LOGGER.warning(
                     "eai_exp and aai_agg are undefined with forecasts. "
                     "Setting them to NaN arrays."

climada/engine/impact_forecast.py

@@ -21,14 +21,14 @@
 
 import logging
 from pathlib import Path
-from typing import Union
+from typing import Literal, Union
 
 import numpy as np
-import scipy.sparse as sparse
+from scipy import sparse
 
 from ..util import log_level
 from ..util.checker import size
-from ..util.forecast import Forecast
+from ..util.forecast import Forecast, reduce_unique_selection
 from .impact import Impact
 
 LOGGER = logging.getLogger(__name__)
@@ -244,48 +244,83 @@ 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):
-        """
-        Reduce the attributes of an ImpactForecast to a single value.
-
-        Attributes are modified as follows:
-        - lead_time: set to NaT
-        - member: set to -1
-        - event_id: set to 0
-        - event_name: set to the name of the reduction method (default)
-        - date: set to 0
-        - frequency: set to 1
+    def select(
+        self,
+        event_ids=None,
+        event_names=None,
+        dates=None,
+        coord_exp=None,
+        reset_frequency=False,
+        member=None,
+        lead_time=None,
+    ):
+        """Select entries based on the parameters and return a new instance.
+        The selection will contain the intersection of all given parameters.
 
         Parameters
         ----------
-        event_name : str
-            The event name given to the reduced data.
+        member : Sequence of ints
+            Ensemble members to select
+        lead_time : Sequence of numpy.timedelta64
+            Lead times to select
+
+        See Also
+        --------
+        :py:meth:`~climada.engine.impact.Impact.select`
         """
-        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]),
-        }
-
-        return reduced_attrs
-
-    def min(self):
+        if member is not None or lead_time is not None:
+            mask_member = (
+                self.idx_member(member)
+                if member is not None
+                else np.full_like(self.member, True, dtype=bool)
+            )
+            mask_lead_time = (
+                self.idx_lead_time(lead_time)
+                if lead_time is not None
+                else np.full_like(self.lead_time, True, dtype=bool)
+            )
+            event_id_from_forecast_mask = np.asarray(self.event_id)[
+                (mask_member & mask_lead_time)
+            ]
+            event_ids = (
+                np.intersect1d(event_ids, event_id_from_forecast_mask)
+                if event_ids is not None
+                else event_id_from_forecast_mask
+            )
+
+        return super().select(
+            event_ids=event_ids,
+            event_names=event_names,
+            dates=dates,
+            coord_exp=coord_exp,
+            reset_frequency=reset_frequency,
+        )
+
+    def min(self, dim: Literal["member", "lead_time"] | None = None):
         """
         Reduce the impact matrix and at_event of an ImpactForecast to the minimum
         value.
 
         Parameters
         ----------
-        None
+        dim : "member", "lead_time", or None
+            Dimension to reduce over. If None, reduce over all data.
 
         Returns
         -------
         ImpactForecast
             An ImpactForecast object with the min impact matrix and at_event.
         """
+        if dim is not None:
+            rdim = self._reduce_iter_dim(dim)
+            return reduce_unique_selection(
+                self,
+                values=getattr(self, rdim),
+                select=rdim,
+                reduce_attr="min",
+                concat_kws={"reset_event_ids": True},
+            )
+
         red_imp_mat = self.imp_mat.min(axis=0).tocsr()
         red_at_event = np.array([red_imp_mat.sum()])
         return ImpactForecast(
@@ -302,20 +337,31 @@ def min(self):
             **self._reduce_attrs("min"),
         )
 
-    def max(self):
+    def max(self, dim: Literal["member", "lead_time"] | None = None):
         """
         Reduce the impact matrix and at_event of an ImpactForecast to the maximum
         value.
 
         Parameters
         ----------
-        None
+        dim : "member", "lead_time", or None
+            Dimension to reduce over. If None, reduce over all data.
 
         Returns
         -------
         ImpactForecast
             An ImpactForecast object with the max impact matrix and at_event.
         """
+        if dim is not None:
+            rdim = self._reduce_iter_dim(dim)
+            return reduce_unique_selection(
+                self,
+                values=getattr(self, rdim),
+                select=rdim,
+                reduce_attr="max",
+                concat_kws={"reset_event_ids": True},
+            )
+
         red_imp_mat = self.imp_mat.max(axis=0).tocsr()
         red_at_event = np.array([red_imp_mat.sum()])
         return ImpactForecast(
@@ -332,19 +378,30 @@ def max(self):
             **self._reduce_attrs("max"),
         )
 
-    def mean(self):
+    def mean(self, dim: Literal["member", "lead_time"] | None = None):
         """
         Reduce the impact matrix and at_event of an ImpactForecast to the mean value.
 
         Parameters
         ----------
-        None
+        dim : "member", "lead_time", or None
+            Dimension to reduce over. If None, reduce over all data.
 
         Returns
         -------
         ImpactForecast
             An ImpactForecast object with the mean impact matrix and at_event.
         """
+        if dim is not None:
+            rdim = self._reduce_iter_dim(dim)
+            return reduce_unique_selection(
+                self,
+                values=getattr(self, rdim),
+                select=rdim,
+                reduce_attr="mean",
+                concat_kws={"reset_event_ids": True},
+            )
+
         red_imp_mat = sparse.csr_matrix(self.imp_mat.mean(axis=0))
         red_at_event = np.array([red_imp_mat.sum()])
         return ImpactForecast(
@@ -361,62 +418,26 @@ def mean(self):
             **self._reduce_attrs("mean"),
         )
 
-    def select(
+    def _quantile(
         self,
-        event_ids=None,
-        event_names=None,
-        dates=None,
-        coord_exp=None,
-        reset_frequency=False,
-        member=None,
-        lead_time=None,
+        q: float,
+        dim: Literal["member", "lead_time"] | None = None,
+        event_name: str | None = None,
     ):
-        """Select entries based on the parameters and return a new instance.
-        The selection will contain the intersection of all given parameters.
-
-        Parameters
-        ----------
-        member : Sequence of ints
-            Ensemble members to select
-        lead_time : Sequence of numpy.timedelta64
-            Lead times to select
-
-        See Also
-        --------
-        :py:meth:`~climada.engine.impact.Impact.select`
-        """
-        if member is not None or lead_time is not None:
-            mask_member = (
-                self.idx_member(member)
-                if member is not None
-                else np.full_like(self.member, True, dtype=bool)
-            )
-            mask_lead_time = (
-                self.idx_lead_time(lead_time)
-                if lead_time is not None
-                else np.full_like(self.lead_time, True, dtype=bool)
-            )
-            event_id_from_forecast_mask = np.asarray(self.event_id)[
-                (mask_member & mask_lead_time)
-            ]
-            event_ids = (
-                np.intersect1d(event_ids, event_id_from_forecast_mask)
-                if event_ids is not None
-                else event_id_from_forecast_mask
-            )
-
-        return super().select(
-            event_ids=event_ids,
-            event_names=event_names,
-            dates=dates,
-            coord_exp=coord_exp,
-            reset_frequency=reset_frequency,
-        )
-
-    def _quantile(self, q: float, event_name: str | None = None):
         """
         Reduce the impact matrix and at_event of an ImpactForecast to the quantile value.
         """
+        if dim is not None:
+            rdim = self._reduce_iter_dim(dim)
+            return reduce_unique_selection(
+                self,
+                values=getattr(self, rdim),
+                select=rdim,
+                reduce_attr="quantile",
+                q=q,
+                concat_kws={"reset_event_ids": True},
+            )
+
         red_imp_mat = sparse.csr_matrix(np.quantile(self.imp_mat.toarray(), q, axis=0))
         red_at_event = np.array([red_imp_mat.sum()])
         if event_name is None:
@@ -435,33 +456,38 @@ def _quantile(self, q: float, event_name: str | None = None):
             **self._reduce_attrs(event_name),
         )
 
-    def quantile(self, q: float):
+    def quantile(self, q: float, dim: Literal["member", "lead_time"] | None = None):
         """
         Reduce the impact matrix and at_event of an ImpactForecast to the quantile value.
 
         Parameters
         ----------
         q : float
             The quantile to compute, which must be between 0 and 1.
+        dim : "member", "lead_time", or None
+            The dimension to reduce when computing the quantile. If ``None`` (default),
+            this computes the centroid-wise quantile over the entire matrix.
 
         Returns
         -------
         ImpactForecast
             An ImpactForecast object with the quantile impact matrix and at_event.
         """
-        return self._quantile(q=q)
+        return self._quantile(q=q, dim=dim)
 
-    def median(self):
+    def median(self, dim: Literal["member", "lead_time"] | None = None):
         """
         Reduce the impact matrix and at_event of an ImpactForecast to the median value.
 
         Parameters
         ----------
-        None
+        dim : "member", "lead_time", or None
+            The dimension to reduce when computing the median. If ``None`` (default),
+            this computes the centroid-wise median over the entire matrix.
 
         Returns
         -------
         ImpactForecast
             An ImpactForecast object with the median impact matrix and at_event.
         """
-        return self._quantile(q=0.5, event_name="median")
+        return self._quantile(q=0.5, dim=dim, event_name="median")