diff --git a/pysteps/blending/pca_ens_kalman_filter.py b/pysteps/blending/pca_ens_kalman_filter.py index b68fd026..dcab310c 100644 --- a/pysteps/blending/pca_ens_kalman_filter.py +++ b/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, diff --git a/pysteps/tests/test_blending_pca_ens_kalman_filter.py b/pysteps/tests/test_blending_pca_ens_kalman_filter.py index 74216b0c..3ed29bb1 100644 --- a/pysteps/tests/test_blending_pca_ens_kalman_filter.py +++ b/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