Allow for semi dynamic incremental mask in blending setup (#459)

RubenImhoff · ladc · web-flow · commit 769880ce9ccb · 2025-03-28T21:55:35.000+01:00
* feat: allow for semi dynamic incremental mask in blending setup

* fix: remove remnant code

* Simplify max_rim range computation

* fix: add tests and update docstrings

---------

Co-authored-by: Lesley De Cruz &lt;lesley.decruz@meteo.be&gt;
diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py
@@ -147,10 +147,10 @@ class StepsBlendingConfig:
       field, 'incremental' = iteratively buffer the mask with a certain rate
       (currently it is 1 km/min), None=no masking.
     resample_distribution: bool, optional
-        Method to resample the distribution from the extrapolation and NWP cascade as input
-        for the probability matching. Not resampling these distributions may lead to losing
-        some extremes when the weight of both the extrapolation and NWP cascade is similar.
-        Defaults to True.
+      Method to resample the distribution from the extrapolation and NWP cascade as input
+      for the probability matching. Not resampling these distributions may lead to losing
+      some extremes when the weight of both the extrapolation and NWP cascade is similar.
+      Defaults to True.
     smooth_radar_mask_range: int, Default is 0.
       Method to smooth the transition between the radar-NWP-noise blend and the NWP-noise
       blend near the edge of the radar domain (radar mask), where the radar data is either
@@ -249,9 +249,9 @@ class StepsBlendingConfig:
       (the number of NWP models) and 'window_length' (the minimum number of
       days the clim file should have, otherwise the default is used).
     mask_kwargs: dict
-      Optional dictionary containing mask keyword arguments 'mask_f' and
-      'mask_rim', the factor defining the the mask increment and the rim size,
-      respectively.
+      Optional dictionary containing mask keyword arguments 'mask_f',
+      'mask_rim' and 'max_mask_rim', the factor defining the the mask
+      increment and the (maximum) rim size, respectively.
       The mask increment is defined as mask_f*timestep/kmperpixel.
     measure_time: bool
       If set to True, measure, print and return the computation time.
@@ -639,7 +639,10 @@ def worker(j):
                         self.__recompose_cascade_to_rainfall_field(j, worker_state)
                         final_blended_forecast_single_member = (
                             self.__post_process_output(
-                                j, final_blended_forecast_single_member, worker_state
+                                j,
+                                t_sub,
+                                final_blended_forecast_single_member,
+                                worker_state,
                             )
                         )
                     final_blended_forecast_all_members_one_timestep[j] = (
@@ -1480,6 +1483,9 @@ def __prepare_forecast_loop(self):
         if self.__config.mask_method == "incremental":
             # get mask parameters
             self.__params.mask_rim = self.__params.mask_kwargs.get("mask_rim", 10)
+            self.__params.max_mask_rim = self.__params.mask_kwargs.get(
+                "max_mask_rim", 10
+            )
             mask_f = self.__params.mask_kwargs.get("mask_f", 1.0)
             # initialize the structuring element
             struct = generate_binary_structure(2, 1)
@@ -2511,7 +2517,7 @@ def __recompose_cascade_to_rainfall_field(self, j, worker_state):
             )
 
     def __post_process_output(
-        self, j, final_blended_forecast_single_member, worker_state
+        self, j, t_sub, final_blended_forecast_single_member, worker_state
     ):
         """
         Apply post-processing steps to refine the final blended forecast. This
@@ -2674,16 +2680,16 @@ def __post_process_output(
                 worker_state.final_blended_forecast_recomposed.min()
             )
             if self.__config.mask_method == "incremental":
-                # The incremental mask is slightly different from
-                # the implementation in the non-blended steps.py, as
-                # it is not based on the last forecast, but instead
-                # on R_pm_blended. Therefore, the buffer does not
-                # increase over time.
-                # Get the mask for this forecast
+                # The incremental mask is slightly different from the implementation in
+                # nowcasts.steps.py, as it is not computed in the Lagrangian space. Instead,
+                # we use precip_forecast_probability_matched and let the mask_rim increase with
+                # the time step until mask_rim_max. This ensures that for the first t time
+                # steps, the buffer mask keeps increasing.
                 precip_field_mask = (
                     precip_forecast_probability_matching_blended
                     >= self.__config.precip_threshold
                 )
+
                 # Buffer the mask
                 # Convert the precipitation field mask into an 8-bit unsigned integer mask
                 obs_mask_uint8 = precip_field_mask.astype("uint8")
@@ -2698,7 +2704,10 @@ def __post_process_output(
                 accumulated_mask = dilated_mask.astype(float)
 
                 # Iteratively dilate the mask and accumulate the results to create a grayscale rim
-                for _ in range(self.__params.mask_rim):
+                mask_rim_temp = min(
+                    self.__params.mask_rim + t_sub - 1, self.__params.max_mask_rim
+                )
+                for _ in range(mask_rim_temp):
                     dilated_mask = binary_dilation(dilated_mask, struct_element)
                     accumulated_mask += dilated_mask
 
@@ -3095,9 +3104,9 @@ def forecast(
       (the number of NWP models) and 'window_length' (the minimum number of
       days the clim file should have, otherwise the default is used).
     mask_kwargs: dict
-      Optional dictionary containing mask keyword arguments 'mask_f' and
-      'mask_rim', the factor defining the the mask increment and the rim size,
-      respectively.
+      Optional dictionary containing mask keyword arguments 'mask_f',
+      'mask_rim' and 'max_mask_rim', the factor defining the the mask
+      increment and the (maximum) rim size, respectively.
       The mask increment is defined as mask_f*timestep/kmperpixel.
     measure_time: bool
       If set to True, measure, print and return the computation time.
diff --git a/pysteps/tests/test_blending_steps.py b/pysteps/tests/test_blending_steps.py
@@ -10,47 +10,50 @@
 
 # fmt:off
 steps_arg_values = [
-    (1, 3, 4, 8, None, None, False, "spn", True, 4, False, False, 0, False, None),
-    (1, 3, 4, 8, "obs", None, False, "spn", True, 4, False, False, 0, False, None),
-    (1, 3, 4, 8, "incremental", None, False, "spn", True, 4, False, False, 0, False, None),
-    (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, False, None),
-    (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, True, None),
-    (1, 3, 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False, None),
-    (1, [1, 2, 3], 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False, None),
-    (1, 3, 4, 8, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False, None),
-    (1, 3, 4, 6, "incremental", "cdf", False, "bps", True, 4, False, False, 0, False, None),
-    (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, False, None),
-    (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, True, None),
-    (1, 3, 4, 9, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False, None),
-    (2, 3, 10, 8, "incremental", "cdf", False, "spn", True, 10, False, False, 0, False, None),
-    (5, 3, 5, 8, "incremental", "cdf", False, "spn", True, 5, False, False, 0, False, None),
-    (1, 10, 1, 8, "incremental", "cdf", False, "spn", True, 1, False, False, 0, False, None),
-    (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False, None),
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, None),
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, "bps"),
-    # TODO: make next test work! This is currently not working on the main branch
-    # (2, 3, 4, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False),
-    # (2, 3, 4, 8, "incremental", "cdf", False, "spn", True, 2, False, False, 0, False),
+    (1, 3, 4, 8, None, None, False, "spn", True, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 8, "obs", None, False, "spn", True, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 8, "incremental", None, False, "spn", True, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0, True, None, None),
+    (1, 3, 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False, None, None),
+    (1, [1, 2, 3], 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 8, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 6, "incremental", "cdf", False, "bps", True, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, False, None, None),
+    (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0, True, None, None),
+    (1, 3, 4, 9, "incremental", "cdf", False, "spn", True, 4, False, False, 0, False, None, None),
+    (2, 3, 10, 8, "incremental", "cdf", False, "spn", True, 10, False, False, 0, False, None, None),
+    (5, 3, 5, 8, "incremental", "cdf", False, "spn", True, 5, False, False, 0, False, None, None),
+    (1, 10, 1, 8, "incremental", "cdf", False, "spn", True, 1, False, False, 0, False, None, None),
+    (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0, False, None, None),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, None, None),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0, False, "bps", None),
     #    Test the case where the radar image contains no rain.
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0, False, None),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, False, None),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, True, None),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0, False, None, None),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, False, None, None),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0, True, None, None),
     #   Test the case where the NWP fields contain no rain.
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 0, False, None),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, True, 0, True, None),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 0, False, None, None),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, True, 0, True, None, None),
     # Test the case where both the radar image and the NWP fields contain no rain.
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, True, True, 0, False, None),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, True, 0, False, None),
-    (5, 3, 5, 6, "obs", "mean", True, "spn", True, 5, True, True, 0, False, None),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, True, True, 0, False, None, None),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, True, 0, False, None, None),
+    (5, 3, 5, 6, "obs", "mean", True, "spn", True, 5, True, True, 0, False, None, None),
     # Test for smooth radar mask
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 80, False, None),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, False, 80, False, None),
-    (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, False, False, 80, False, None),
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 80, False, None),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 80, True, None),
-    (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None),
-    (5, [1, 2, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None),
-    (5, [1, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 80, False, None, None),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, False, 80, False, None, None),
+    (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, False, False, 80, False, None, None),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 80, False, None, None),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 80, True, None, None),
+    (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None, None),
+    (5, [1, 2, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None, None),
+    (5, [1, 3], 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80, False, None, None),
+    # Test the usage of a max_mask_rim in the mask_kwargs
+    (1, 3, 6, 8, None, None, False, "bps", True, 6, False, False, 80, False, None, 40),
+    (5, 3, 5, 6, "obs", "mean", False, "bps", False, 5, False, False, 80, False, None, 40),
+    (5, 3, 5, 6, "incremental", "cdf", False, "bps", False, 5, False, False, 80, False, None, 25),
+    (5, 3, 5, 6, "incremental", "cdf", False, "bps", False, 5, False, False, 80, False, None, 40),
+    (5, 3, 5, 6, "incremental", "cdf", False, "bps", False, 5, False, False, 80, False, None, 60),
 ]
 # fmt:on
 
@@ -70,6 +73,7 @@
     "smooth_radar_mask_range",
     "resample_distribution",
     "vel_pert_method",
+    "max_mask_rim",
 )
 
 
@@ -90,6 +94,7 @@ def test_steps_blending(
     smooth_radar_mask_range,
     resample_distribution,
     vel_pert_method,
+    max_mask_rim,
 ):
     pytest.importorskip("cv2")
 
@@ -162,13 +167,18 @@ def test_steps_blending(
     metadata["zr_a"] = 200.0
     metadata["zr_b"] = 1.6
 
-    # Also set the outdir_path and clim_kwargs
+    # Also set the outdir_path, clim_kwargs and mask_kwargs
     outdir_path_skill = "./tmp/"
     if n_models == 1:
         clim_kwargs = None
     else:
         clim_kwargs = dict({"n_models": n_models, "window_length": 30})
 
+    if max_mask_rim is not None:
+        mask_kwargs = dict({"mask_rim": 10, "max_mask_rim": max_mask_rim})
+    else:
+        mask_kwargs = None
+
     ###
     # First threshold the data and convert it to dBR
     ###
@@ -288,6 +298,7 @@ def test_steps_blending(
         conditional=False,
         probmatching_method=probmatching_method,
         mask_method=mask_method,
+        resample_distribution=resample_distribution,
         smooth_radar_mask_range=smooth_radar_mask_range,
         callback=None,
         return_output=True,
@@ -301,7 +312,7 @@ def test_steps_blending(
         noise_kwargs=None,
         vel_pert_kwargs=None,
         clim_kwargs=clim_kwargs,
-        mask_kwargs=None,
+        mask_kwargs=mask_kwargs,
         measure_time=False,
     )
 
