Support for mixed precip/noprecip input and AR order 1

.pre-commit-config.yaml

@@ -4,3 +4,10 @@ repos:
     hooks:
     - id: black
       language_version: python3
+
+-   repo: https://github.com/PyCQA/bandit
+    rev: 1.9.4
+    hooks:
+    - id: bandit
+      args: ["-c", "bandit.yaml"]
+      additional_dependencies: ["bandit[yaml]"]

bandit.yaml

@@ -0,0 +1,3 @@
+# skip check for assert usage in the tests
+assert_used:
+  skips: ['*test_*.py']

pysteps/blending/skill_scores.py

@@ -136,7 +136,9 @@ def lt_dependent_cor_nwp(lt, correlations, outdir_path, n_model=0, skill_kwargs=
     return rho
 
 
-def lt_dependent_cor_extrapolation(PHI, correlations=None, correlations_prev=None):
+def lt_dependent_cor_extrapolation(
+    PHI, correlations=None, correlations_prev=None, ar_order=2
+):
     """Determine the correlation of the extrapolation (nowcast) component for
     lead time lt and cascade k, by assuming that the correlation determined at
     t=0 regresses towards the climatological values.
@@ -153,6 +155,8 @@ def lt_dependent_cor_extrapolation(PHI, correlations=None, correlations_prev=Non
         be found from the AR-2 model.
     correlations_prev : array-like, optional
         Similar to correlations, but from the timestep before that.
+    ar_order : int, optional
+        The order of the autoregressive model to use. Must be >= 1.
 
     Returns
     -------
@@ -170,12 +174,22 @@ def lt_dependent_cor_extrapolation(PHI, correlations=None, correlations_prev=Non
     if correlations_prev is None:
         correlations_prev = np.repeat(1.0, PHI.shape[0])
     # Same for correlations at first time step, we set it to
-    # phi1 / (1 - phi2), see BPS2004
-    if correlations is None:
-        correlations = PHI[:, 0] / (1.0 - PHI[:, 1])
-
-    # Calculate the correlation for lead time lt
-    rho = PHI[:, 0] * correlations + PHI[:, 1] * correlations_prev
+    if ar_order == 1:
+        # Yule-Walker equation for AR-1 model -> just PHI
+        if correlations is None:
+            correlations = PHI[:, 0]
+        # Calculate the correlation for lead time lt (AR-1)
+        rho = PHI[:, 0] * correlations
+    elif ar_order == 2:
+        # phi1 / (1 - phi2), see BPS2004
+        if correlations is None:
+            correlations = PHI[:, 0] / (1.0 - PHI[:, 1])
+        # Calculate the correlation for lead time lt (AR-2)
+        rho = PHI[:, 0] * correlations + PHI[:, 1] * correlations_prev
+    else:
+        raise ValueError(
+            "Autoregression of higher order than 2 is not defined. Please set ar_order = 2."
+        )
 
     # Finally, set the current correlations array as the previous one for the
     # next time step

pysteps/blending/steps.py

@@ -1196,12 +1196,22 @@ def transform_to_lagrangian(precip, i):
             self.__precip_models = np.stack(temp_precip_models)
 
         # Check for zero input fields in the radar, nowcast and NWP data.
+        # decide based on latest input file only
         self.__params.zero_precip_radar = check_norain(
-            self.__precip,
+            self.__precip[-1],
             self.__params.precip_threshold,
             self.__config.norain_threshold,
             self.__params.noise_kwargs["win_fun"],
         )
+        # Set all inputs to 0 (also previous times)
+        # not strictly necessary but to be consistent with checking only
+        # latest observation rain field to set zero_precip_radar
+        if self.__params.zero_precip_radar:
+            self.__precip = np.where(
+                np.isfinite(self.__precip),
+                np.ones(self.__precip.shape) * np.nanmin(self.__precip),
+                self.__precip,
+            )
 
         # The norain fraction threshold used for nwp is the default value of 0.0,
         # since nwp does not suffer from clutter.
@@ -1481,6 +1491,8 @@ def __estimate_ar_parameters_radar(self):
 
             # Finally, transpose GAMMA to ensure that the shape is the same as np.empty((n_cascade_levels, ar_order))
             GAMMA = GAMMA.transpose()
+            if len(GAMMA.shape) == 1:
+                GAMMA = GAMMA.reshape((GAMMA.size, 1))
             assert GAMMA.shape == (
                 self.__config.n_cascade_levels,
                 self.__config.ar_order,
@@ -1629,13 +1641,22 @@ def __prepare_forecast_loop(self):
             )
 
         # initizalize the current and previous extrapolation forecast scale for the nowcasting component
-        # phi1 / (1 - phi2), see BPS2004
+        # ar_order=1: rho = phi1
+        # ar_order=2: rho = phi1 / (1 - phi2)
+        #   -> see Hamilton1994, Pulkkinen2019, BPS2004 (Yule-Walker equations)
         self.__state.rho_extrap_cascade_prev = np.repeat(
             1.0, self.__params.PHI.shape[0]
         )
-        self.__state.rho_extrap_cascade = self.__params.PHI[:, 0] / (
-            1.0 - self.__params.PHI[:, 1]
-        )
+        if self.__config.ar_order == 1:
+            self.__state.rho_extrap_cascade = self.__params.PHI[:, 0]
+        elif self.__config.ar_order == 2:
+            self.__state.rho_extrap_cascade = self.__params.PHI[:, 0] / (
+                1.0 - self.__params.PHI[:, 1]
+            )
+        else:
+            raise ValueError(
+                "Autoregression of higher order than 2 is not defined. Please set ar_order = 2."
+            )
 
     def __initialize_noise_cascades(self):
         """
@@ -2037,6 +2058,7 @@ def __determine_skill_for_current_timestep(self, t):
                 PHI=self.__params.PHI,
                 correlations=self.__state.rho_extrap_cascade,
                 correlations_prev=self.__state.rho_extrap_cascade_prev,
+                ar_order=self.__config.ar_order,
             )
 
     def __determine_NWP_skill_for_next_timestep(self, t, j, worker_state):
@@ -2228,9 +2250,12 @@ def __regress_extrapolation_and_noise_cascades(self, j, worker_state, t):
                         )
                     )
                     # Renormalize the cascade
-                    worker_state.precip_cascades[j][i][1] /= np.std(
-                        worker_state.precip_cascades[j][i][1]
-                    )
+                    for nl, obs_cascade_level in enumerate(
+                        worker_state.precip_cascades[j][i]
+                    ):
+                        worker_state.precip_cascades[j][i][nl] /= np.std(
+                            obs_cascade_level
+                        )
                 else:
                     # use the deterministic AR(p) model computed above if
                     # perturbations are disabled
@@ -3625,6 +3650,11 @@ def forecast(
     turns out to be a warranted functionality.
     """
 
+    # Check the input precip and ar_order to be consistent
+    # zero-precip in previous time steps has to be removed
+    # (constant field causes autoregression to fail)
+    precip, ar_order = check_previous_radar_obs(precip, ar_order)
+
     blending_config = StepsBlendingConfig(
         n_ens_members=n_ens_members,
         n_cascade_levels=n_cascade_levels,
@@ -3686,6 +3716,65 @@ def forecast(
     return forecast_steps_nowcast
 
 
+# TODO: Where does this piece of code best fit: in utils or inside the class?
+def check_previous_radar_obs(precip, ar_order):
+    """Check all radar time steps and remove zero precipitation time steps
+
+    Parameters
+    ----------
+    precip : array-like
+      Array of shape (ar_order+1,m,n) containing the input precipitation fields
+      ordered by timestamp from oldest to newest. The time steps between the
+      inputs are assumed to be regular.
+    ar_order : int
+      The order of the autoregressive model to use. Must be >= 1.
+
+    Returns
+    -------
+    precip : numpy array
+      Array of shape (ar_order+1,m,n) containing the modified array with
+      input precipitation fields ordered by timestamp from oldest to newest.
+      The time steps between the inputs are assumed to be regular.
+    ar_order : int
+      The order of the autoregressive model to use. Must be >= 1.
+      Adapted to match with precip.shape equal (ar_order+1,m,n).
+    """
+    # Quick check radar input - at least 2 time steps
+    if not precip.shape[0] >= 2:
+        raise ValueError(
+            "Wrong precip shape. The radar input must have at least 2 time steps."
+        )
+
+    # Check all time steps for zero-precip (constant field, minimum==maximum)
+    zero_precip = [np.nanmin(obs) == np.nanmax(obs) for obs in precip]
+    if zero_precip[-1] or ~np.any(zero_precip):
+        # Unchanged if no rain in latest time step -> will be processed as zero_precip_radar=True
+        # or Unchanged if all time steps contain rain
+        return precip, ar_order
+    elif zero_precip[-2]:
+        # try to use a previous time step
+        if not np.all(zero_precip[:-2]):
+            # find latest non-zero precip
+            # ATTENTION: This changes the time between precip[-2] and precip[-1] from initial 5min to a longer period
+            print(
+                "[WARNING] Radar input time steps adapted and ar_order set to 1. Input delta time changed."
+            )
+            prev = np.arange(len(zero_precip[:-2]))[~np.array(zero_precip[:-2])][-1]
+            return precip[[prev, -1]], 1
+        raise ValueError(
+            "Precipitation in latest but no previous time step. Not possible to calculate autoregression."
+        )
+    else:
+        # Keep the latest time steps that do all contain precip
+        precip = precip[np.max(np.arange(len(zero_precip))[zero_precip]) + 1 :].copy()
+        if precip.shape[0] - 1 < ar_order:
+            # Give a warning
+            print(
+                f"[WARNING] Remove zero-precip time steps from radar input.\nRemaining non-zero time steps: {precip.shape[0]}, ar_order set to {precip.shape[0]-1}."
+            )
+        return precip, min(ar_order, precip.shape[0] - 1)
+
+
 # TODO: Where does this piece of code best fit: in utils or inside the class?
 def calculate_ratios(correlations):
     """Calculate explained variance ratios from correlation.

pysteps/nowcasts/steps.py

@@ -358,11 +358,19 @@ def compute_forecast(self):
         self.__precip = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy()
         self.__initialize_nowcast_components()
         if check_norain(
-            self.__precip,
+            self.__precip[-1],
             self.__config.precip_threshold,
             self.__config.norain_threshold,
             self.__params.noise_kwargs["win_fun"],
         ):
+            # Set all to inputs to 0 (also previous times) as we just check latest input in check_norain
+            self.__precip = self.__precip.copy()
+            self.__precip = np.where(
+                np.isfinite(self.__precip),
+                np.ones(self.__precip.shape) * np.nanmin(self.__precip),
+                self.__precip,
+            )
+
             return zero_precipitation_forecast(
                 self.__config.n_ens_members,
                 self.__time_steps,
@@ -1495,6 +1503,11 @@ def forecast(
     :cite:`Seed2003`, :cite:`BPS2006`, :cite:`SPN2013`, :cite:`PCH2019b`
     """
 
+    # Check the input precip and ar_order to be consistent
+    # zero-precip in previous time steps has to be removed
+    # (constant field causes autoregression to fail)
+    precip, ar_order = check_previous_radar_obs(precip, ar_order)
+
     nowcaster_config = StepsNowcasterConfig(
         n_ens_members=n_ens_members,
         n_cascade_levels=n_cascade_levels,
@@ -1534,3 +1547,61 @@ def forecast(
     nowcaster.reset_states_and_params()
     # Call the appropriate methods within the class
     return forecast_steps_nowcast
+
+
+# TODO: Where does this piece of code best fit: in utils or inside the class?
+def check_previous_radar_obs(precip, ar_order):
+    """Check all radar time steps and remove zero precipitation time steps
+
+    Parameters
+    ----------
+    precip : array-like
+      Array of shape (ar_order+1,m,n) containing the input precipitation fields
+      ordered by timestamp from oldest to newest. The time steps between the
+      inputs are assumed to be regular.
+    ar_order : int
+      The order of the autoregressive model to use. Must be >= 1.
+
+    Returns
+    -------
+    precip : numpy array
+      Array of shape (ar_order+1,m,n) containing the modified array with
+      input precipitation fields ordered by timestamp from oldest to newest.
+      The time steps between the inputs are assumed to be regular.
+    ar_order : int
+      The order of the autoregressive model to use. Must be >= 1.
+      Adapted to match with precip.shape equal (ar_order+1,m,n).
+    """
+    if not precip.shape[0] >= 2:
+        raise ValueError(
+            "Wrong precip shape. The radar input must have at least 2 time steps."
+        )
+
+    # Check all time steps for zero-precip (constant field, minimum==maximum)
+    zero_precip = [np.nanmin(obs) == np.nanmax(obs) for obs in precip]
+    if zero_precip[-1] or ~np.any(zero_precip):
+        # Unchanged if no rain in latest time step -> will be processed as zero_precip_radar=True
+        # or Unchanged if all time steps contain rain
+        return precip, ar_order
+    elif zero_precip[-2]:
+        # try to use a previous time step
+        if not np.all(zero_precip[:-2]):
+            # find latest non-zero precip
+            # ATTENTION: This changes the time between precip[-2] and precip[-1] from initial 5min to a longer period
+            print(
+                "[WARNING] Radar input time steps adapted and ar_order set to 1. Input delta time changed."
+            )
+            prev = np.arange(len(zero_precip[:-2]))[~np.array(zero_precip[:-2])][-1]
+            return precip[[prev, -1]], 1
+        raise ValueError(
+            "Precipitation in latest but no previous time step. Not possible to calculate autoregression."
+        )
+    else:
+        # Keep the latest time steps that do all contain precip
+        precip = precip[np.max(np.arange(len(zero_precip))[zero_precip]) + 1 :].copy()
+        if precip.shape[0] - 1 < ar_order:
+            # Give a warning
+            print(
+                f"[WARNING] Radar input only with {precip.shape[0]} non-zero time steps and ar_order set to {precip.shape[0]-1}."
+            )
+        return precip, min(ar_order, precip.shape[0] - 1)