fix: Advection of domain mask for smooth blending at domain edges in reduced-space EnKF (#533)

feat: add smooth_radar_mask functionality from steps blending to pca_ens_kalman_filter blending


---------

Co-authored-by: Ruben Imhoff 
Co-authored-by: mats-knm

Author: m-rempel

pysteps/blending/pca_ens_kalman_filter.py

@@ -176,6 +176,8 @@ class EnKFCombinationConfig:
         :py:mod:`pysteps.blending.ens_kalman_filter_methods`.
     measure_time: bool
         If set to True, measure, print and return the computation time.
+    verbose_output: bool
+        If set to True, return additionally the background ensemble of the EnKF for further statistics.
     callback: function, optional
       Optional function that is called after computation of each time step of
       the nowcast. The function takes one argument: a three-dimensional array
@@ -214,6 +216,7 @@ class EnKFCombinationConfig:
     noise_kwargs: dict[str, Any] = field(default_factory=dict)
     combination_kwargs: dict[str, Any] = field(default_factory=dict)
     measure_time: bool = False
+    verbose_output: bool = False
     callback: Any | None = None
     return_output: bool = True
     n_noise_fields: int = 30
@@ -609,6 +612,7 @@ def __init__(
         self.nwc_prediction_btf = self.nwc_prediction.copy()
 
         self.final_combined_forecast = []
+        self.background_ensemble = {}
 
         return
 
@@ -785,7 +789,10 @@ def __advect(self):
             displacement_previous=self.__previous_displacement,
             **self.__forecast_state.params.extrapolation_kwargs,
         )
-        if self.__forecast_state.config.smooth_radar_mask_range > 0:
+        if (
+            self.__forecast_state.config.smooth_radar_mask_range > 0
+            and self.__ens_member == 0
+        ):
             self.__forecast_state.params.domain_mask = (
                 self.__forecast_state.params.extrapolation_method(
                     self.__forecast_state.params.domain_mask,
@@ -1082,6 +1089,8 @@ def compute_forecast(self):
                     self.__fc_init,
                     self.__mainloop_time,
                 )
+            if self.__config.verbose_output:
+                return self.FS.final_combined_forecast, self.FS.background_ensemble
             return self.FS.final_combined_forecast
 
         # Else, return None
@@ -1406,7 +1415,7 @@ def __integrated_nowcast_main_loop(self):
                 print(f"Full NWP weight is reached for lead time + {fc_leadtime} min")
                 if is_correction_timestep:
                     t_corr = np.where(
-                        self.__correction_leadtimes == self.__forecast_leadtimes[t - 1]
+                        self.__correction_leadtimes == self.__forecast_leadtimes[t]
                     )[0][0]
                 self.FS.nwc_prediction = self.__nwp_precip[:, t_corr]
 
@@ -1476,6 +1485,11 @@ def __forecast_loop(self, t, is_correction_timestep, is_nowcasting_timestep):
                 self.__correction_leadtimes == self.__forecast_leadtimes[t - 1]
             )[0][0]
 
+            if self.__config.verbose_output:
+                self.FS.background_ensemble[self.__correction_leadtimes[t_corr]] = (
+                    self.FS.nwc_prediction.copy()
+                )
+
             self.FS.nwc_prediction, self.FS.fc_resampled = (
                 self.KalmanFilterModel.correct_step(
                     self.FS.nwc_prediction,
@@ -1571,6 +1585,7 @@ def forecast(
     noise_kwargs=None,
     combination_kwargs=None,
     measure_time=False,
+    verbose_output=False,
 ):
     """
     Generate a combined nowcast ensemble by using the reduced-space ensemble Kalman
@@ -1700,7 +1715,9 @@ def forecast(
         ensemble Kalman filter method. See the documentation of
         :py:mod:`pysteps.blending.ens_kalman_filter_methods`.
     measure_time: bool
-      If set to True, measure, print and return the computation time.
+        If set to True, measure, print and return the computation time.
+    verbose_output: bool
+        If set to True, return additionally the background ensemble of the EnKF for further statistics.
 
     Returns
     -------
@@ -1751,6 +1768,7 @@ def forecast(
         noise_kwargs=noise_kwargs,
         combination_kwargs=combination_kwargs,
         measure_time=measure_time,
+        verbose_output=verbose_output,
         callback=callback,
         return_output=return_output,
         n_noise_fields=30,

pysteps/tests/test_blending_pca_ens_kalman_filter.py

@@ -10,39 +10,41 @@
 # fmt: off
 pca_enkf_arg_values = [
     # Standard setting
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Smooth radar mask
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,20),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,20,False),
     # Coarser NWP temporal resolution
-    (20,30,0,-60,False,False,5,15,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-60,False,False,5,15,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Coarser Obs temporal resolution
-    (20,30,0,-60,False,False,10,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-60,False,False,10,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Larger shift of the NWP init
-    (20,30,0,-30,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-30,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Zero rain case in observation
-    (20,30,0,-60,True,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-60,True,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Zero rain case in NWP
-    (20,30,0,-60,False,True,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-60,False,True,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Zero rain in both
-    (20,30,0,-60,True,True,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-60,True,True,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Accumulated sampling probability
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",True,False,0),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",True,False,0,False),
     # Use full NWP weight
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,True,0),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,True,0,False),
     # Both
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",True,True,0),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",True,True,0,False),
     # Explained variance as sampling probability source
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"explained_var",False,False,0),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"explained_var",False,False,0,False),
     # No combination
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",False,None,1.0,"ensemble",False,False,0),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",False,None,1.0,"ensemble",False,False,0,False),
     # Standard deviation adjustment
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,"auto",1.0,"ensemble",False,False,0),
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,"auto",1.0,"ensemble",False,False,0,False),
     # Other number of ensemble members
-    (10,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (10,30,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Other forecast length
-    (20,35,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0),
+    (20,35,0,-60,False,False,5,5,0.05,0.01,"ssft","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
     # Other noise method
-    (20,30,0,-60,False,False,5,5,0.05,0.01,"nonparametric","masked_enkf",True,None,1.0,"ensemble",False,False,0),]
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"nonparametric","masked_enkf",True,None,1.0,"ensemble",False,False,0,False),
+    # Verbose output
+    (20,30,0,-60,False,False,5,5,0.05,0.01,"nonparametric","masked_enkf",True,None,1.0,"ensemble",False,False,0,True),]
 # fmt: on
 
 pca_enkf_arg_names = (
@@ -65,6 +67,7 @@
     "use_accum_sampling_prob",
     "ensure_full_nwp_weight",
     "smooth_radar_mask_range",
+    "verbose_output",
 )
 
 
@@ -89,6 +92,7 @@ def test_pca_enkf_combination(
     use_accum_sampling_prob,
     ensure_full_nwp_weight,
     smooth_radar_mask_range,
+    verbose_output,
 ):
     pytest.importorskip("sklearn")
 
@@ -248,8 +252,13 @@ def test_pca_enkf_combination(
         noise_kwargs=None,
         combination_kwargs=combination_kwargs,
         measure_time=False,
+        verbose_output=verbose_output,
     )
 
+    if verbose_output:
+        assert len(combined_forecast) == 2, "Wrong amount of output data"
+        combined_forecast = combined_forecast[0]
+
     assert combined_forecast.ndim == 4, "Wrong amount of dimensions in forecast output"
     assert (
         combined_forecast.shape[0] == n_ens_members