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

@@ -58,7 +58,7 @@
 from pysteps.nowcasts import utils as nowcast_utils
 from pysteps.postprocessing import probmatching
 from pysteps.timeseries import autoregression, correlation
-from pysteps.utils.check_norain import check_norain
+from pysteps.utils.check_norain import check_norain, check_previous_radar_obs
 
 try:
     import dask
@@ -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,10 +1491,15 @@ 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()
-            assert GAMMA.shape == (
+            if len(GAMMA.shape) == 1:
+                GAMMA = GAMMA.reshape((GAMMA.size, 1))
+            if not GAMMA.shape == (
                 self.__config.n_cascade_levels,
                 self.__config.ar_order,
-            )
+            ):
+                raise ValueError(
+                    f"GAMMA shape {GAMMA.shape} does not match {self.__config.n_cascade_levels} cascade levels and ar_order of {self.__config.ar_order}"
+                )
 
         # Print the GAMMA value
         nowcast_utils.print_corrcoefs(GAMMA)
@@ -1629,13 +1644,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 +2061,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 +2253,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 nt, obs_cascade_timestep in enumerate(
+                        worker_state.precip_cascades[j][i]
+                    ):
+                        worker_state.precip_cascades[j][i][nt] /= np.std(
+                            obs_cascade_timestep
+                        )
                 else:
                     # use the deterministic AR(p) model computed above if
                     # perturbations are disabled
@@ -3625,6 +3653,20 @@ def forecast(
     turns out to be a warranted functionality.
     """
 
+    # Check the input precip and ar_order to be consistent
+    # zero-precip/constant field in previous time steps has to be removed
+    # (constant field causes autoregression to fail)
+    precip, ar_order = check_previous_radar_obs(
+        precip,
+        ar_order,
+        {
+            "precip_thr": precip_thr,
+            "norain_thr": norain_thr,
+            # Set default window function to "tukey" to match the default in check_inputs() below
+            "win_fun": "tukey" if noise_kwargs is None else noise_kwargs["win_fun"],
+        },
+    )
+
     blending_config = StepsBlendingConfig(
         n_ens_members=n_ens_members,
         n_cascade_levels=n_cascade_levels,

pysteps/nowcasts/steps.py

@@ -28,7 +28,7 @@
 )
 from pysteps.postprocessing import probmatching
 from pysteps.timeseries import autoregression, correlation
-from pysteps.utils.check_norain import check_norain
+from pysteps.utils.check_norain import check_norain, check_previous_radar_obs
 
 try:
     import dask
@@ -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,20 @@ def forecast(
     :cite:`Seed2003`, :cite:`BPS2006`, :cite:`SPN2013`, :cite:`PCH2019b`
     """
 
+    # Check the input precip and ar_order to be consistent
+    # zero-precip/constant field in previous time steps has to be removed
+    # (constant field causes autoregression to fail)
+    precip, ar_order = check_previous_radar_obs(
+        precip,
+        ar_order,
+        {
+            "precip_thr": precip_thr,
+            "norain_thr": norain_thr,
+            # Set default window function to "tukey" to match the default in check_inputs() below
+            "win_fun": "tukey" if noise_kwargs is None else noise_kwargs["win_fun"],
+        },
+    )
+
     nowcaster_config = StepsNowcasterConfig(
         n_ens_members=n_ens_members,
         n_cascade_levels=n_cascade_levels,