Use max rain pixels for timestep and member to decompose for precip cascades in case of zero radar (#461)

* Use max rain pixels for timestep and member to decompose for precip cascades in case of zero radar

* store precip threshold in state

---------

Co-authored-by: Daniele Nerini <daniele.nerini@gmail.com>

Author: mats-knmi

pysteps/blending/steps.py

@@ -325,6 +325,7 @@ class StepsBlendingParams:
     precip_models_provided_is_cascade: bool = False
     xy_coordinates: np.ndarray | None = None
     precip_zerovalue: float | None = None
+    precip_threshold: float | None = None
     mask_threshold: np.ndarray | None = None
     zero_precip_radar: bool = False
     zero_precip_model_fields: bool = False
@@ -806,6 +807,8 @@ def __check_inputs(self):
         else:
             self.__params.mask_kwargs = deepcopy(self.__config.mask_kwargs)
 
+        self.__params.precip_threshold = self.__config.precip_threshold
+
         if np.any(~np.isfinite(self.__velocity)):
             raise ValueError("velocity contains non-finite values")
 
@@ -815,12 +818,12 @@ def __check_inputs(self):
                 % self.__config.mask_method
             )
 
-        if self.__config.conditional and self.__config.precip_threshold is None:
+        if self.__config.conditional and self.__params.precip_threshold is None:
             raise ValueError("conditional=True but precip_thr is not set")
 
         if (
             self.__config.mask_method is not None
-            and self.__config.precip_threshold is None
+            and self.__params.precip_threshold is None
         ):
             raise ValueError("mask_method!=None but precip_thr=None")
 
@@ -931,7 +934,7 @@ def __print_forecast_info(self):
             ) = (None, None)
 
         if self.__config.conditional or self.__config.mask_method is not None:
-            print(f"precip. intensity threshold: {self.__config.precip_threshold}")
+            print(f"precip. intensity threshold: {self.__params.precip_threshold}")
         print(f"no-rain fraction threshold for radar: {self.__config.norain_threshold}")
         print("")
 
@@ -994,7 +997,7 @@ def __prepare_radar_and_NWP_fields(self):
             # TODO: is this logical_and correct here? Now only those places where precip is in all images is saved?
             self.__params.mask_threshold = np.logical_and.reduce(
                 [
-                    self.__precip[i, :, :] >= self.__config.precip_threshold
+                    self.__precip[i, :, :] >= self.__params.precip_threshold
                     for i in range(self.__precip.shape[0])
                 ]
             )
@@ -1098,15 +1101,15 @@ def transform_to_lagrangian(precip, i):
         # Check for zero input fields in the radar and NWP data.
         self.__params.zero_precip_radar = check_norain(
             self.__precip,
-            self.__config.precip_threshold,
+            self.__params.precip_threshold,
             self.__config.norain_threshold,
             self.__params.noise_kwargs["win_fun"],
         )
         # The norain fraction threshold used for nwp is the default value of 0.0,
         # since nwp does not suffer from clutter.
         self.__params.zero_precip_model_fields = check_norain(
             self.__precip_models,
-            self.__config.precip_threshold,
+            self.__params.precip_threshold,
             self.__config.norain_threshold,
             self.__params.noise_kwargs["win_fun"],
         )
@@ -1173,43 +1176,6 @@ def __prepare_nowcast_for_zero_radar(self):
         updates the cascade with NWP data, uses the NWP velocity field, and
         initializes the noise model based on the time step with the most rain.
         """
-        # If zero_precip_radar, make sure that precip_cascade does not contain
-        # only nans or infs. If so, fill it with the zero value.
-
-        # Look for a timestep and member with rain so that we have a sensible decomposition
-        done = False
-        for t in self.__timesteps:
-            if done:
-                break
-            for j in range(self.__precip_models.shape[0]):
-                if not check_norain(
-                    self.__precip_models[j, t],
-                    self.__config.precip_threshold,
-                    self.__config.norain_threshold,
-                    self.__params.noise_kwargs["win_fun"],
-                ):
-                    if self.__state.precip_models_cascades is not None:
-                        self.__state.precip_cascades[
-                            ~np.isfinite(self.__state.precip_cascades)
-                        ] = np.nanmin(
-                            self.__state.precip_models_cascades[j, t]["cascade_levels"]
-                        )
-                        continue
-                    precip_models_cascade_timestep = self.__params.decomposition_method(
-                        self.__precip_models[j, t, :, :],
-                        bp_filter=self.__params.bandpass_filter,
-                        fft_method=self.__params.fft,
-                        output_domain=self.__config.domain,
-                        normalize=True,
-                        compute_stats=True,
-                        compact_output=True,
-                    )["cascade_levels"]
-                    self.__state.precip_cascades[
-                        ~np.isfinite(self.__state.precip_cascades)
-                    ] = np.nanmin(precip_models_cascade_timestep)
-                    done = True
-                    break
-
         # If zero_precip_radar is True, only use the velocity field of the NWP
         # forecast. I.e., velocity (radar) equals velocity_model at the first time
         # step.
@@ -1227,16 +1193,16 @@ def __prepare_nowcast_for_zero_radar(self):
         # might be zero as well). Else, initialize the noise with the radar
         # rainfall data
         # Initialize noise based on the NWP field time step where the fraction of rainy cells is highest
-        if self.__config.precip_threshold is None:
-            self.__config.precip_threshold = np.nanmin(self.__precip_models)
+        if self.__params.precip_threshold is None:
+            self.__params.precip_threshold = np.nanmin(self.__precip_models)
 
         max_rain_pixels = -1
         max_rain_pixels_j = -1
         max_rain_pixels_t = -1
         for j in range(self.__precip_models.shape[0]):
             for t in self.__timesteps:
                 rain_pixels = self.__precip_models[j][t][
-                    self.__precip_models[j][t] > self.__config.precip_threshold
+                    self.__precip_models[j][t] > self.__params.precip_threshold
                 ].size
                 if rain_pixels > max_rain_pixels:
                     max_rain_pixels = rain_pixels
@@ -1249,6 +1215,30 @@ def __prepare_nowcast_for_zero_radar(self):
             np.float64, copy=False
         )
 
+        # If zero_precip_radar, make sure that precip_cascade does not contain
+        # only nans or infs. If so, fill it with the zero value.
+        if self.__state.precip_models_cascades is not None:
+            self.__state.precip_cascades[~np.isfinite(self.__state.precip_cascades)] = (
+                np.nanmin(
+                    self.__state.precip_models_cascades[
+                        max_rain_pixels_j, max_rain_pixels_t
+                    ]["cascade_levels"]
+                )
+            )
+        else:
+            precip_models_cascade_timestep = self.__params.decomposition_method(
+                self.__precip_models[max_rain_pixels_j, max_rain_pixels_t, :, :],
+                bp_filter=self.__params.bandpass_filter,
+                fft_method=self.__params.fft,
+                output_domain=self.__config.domain,
+                normalize=True,
+                compute_stats=True,
+                compact_output=True,
+            )["cascade_levels"]
+            self.__state.precip_cascades[~np.isfinite(self.__state.precip_cascades)] = (
+                np.nanmin(precip_models_cascade_timestep)
+            )
+
         # Make sure precip_noise_input is three-dimensional
         if len(self.__state.precip_noise_input.shape) != 3:
             self.__state.precip_noise_input = self.__state.precip_noise_input[
@@ -1281,7 +1271,7 @@ def __initialize_noise(self):
                 precip_forecast_min = np.min(self.__state.precip_noise_input)
                 self.__params.noise_std_coeffs = noise.utils.compute_noise_stddev_adjs(
                     self.__state.precip_noise_input[-1, :, :],
-                    self.__config.precip_threshold,
+                    self.__params.precip_threshold,
                     precip_forecast_min,
                     self.__params.bandpass_filter,
                     self.__params.decomposition_method,
@@ -2687,7 +2677,7 @@ def __post_process_output(
                 # steps, the buffer mask keeps increasing.
                 precip_field_mask = (
                     precip_forecast_probability_matching_blended
-                    >= self.__config.precip_threshold
+                    >= self.__params.precip_threshold
                 )
 
                 # Buffer the mask
@@ -2731,7 +2721,7 @@ def __post_process_output(
                 # rainfall field
                 precip_field_mask_temp = (
                     precip_forecast_probability_matching_blended
-                    >= self.__config.precip_threshold
+                    >= self.__params.precip_threshold
                 )
 
             # Set to min value outside of mask
@@ -2789,12 +2779,12 @@ def __post_process_output(
             mean_probabiltity_matching_forecast = np.mean(
                 precip_forecast_probability_matching_resampled[
                     precip_forecast_probability_matching_resampled
-                    >= self.__config.precip_threshold
+                    >= self.__params.precip_threshold
                 ]
             )
             no_rain_mask = (
                 worker_state.final_blended_forecast_recomposed
-                >= self.__config.precip_threshold
+                >= self.__params.precip_threshold
             )
             mean_precip_forecast = np.mean(
                 worker_state.final_blended_forecast_recomposed[no_rain_mask]