Add option to smooth radar mask

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 -   repo: https://github.com/psf/black
-    rev: 24.3.0
+    rev: 24.4.2
     hooks:
     - id: black
       language_version: python3

pysteps/blending/steps.py

@@ -90,6 +90,7 @@
     conditional=False,
     probmatching_method="cdf",
     mask_method="incremental",
+    smooth_radar_mask_range=0,
     callback=None,
     return_output=True,
     seed=None,
@@ -210,6 +211,11 @@
       'obs' = apply precip_thr to the most recently observed precipitation intensity
       field, 'incremental' = iteratively buffer the mask with a certain rate
       (currently it is 1 km/min), None=no masking.
+    smooth_radar_mask_range: int, Default  is 0.
+      If 0 this generates a normal forecast. To create a smooth mask, this range
+      should be a positive value, representing a buffer band of a number of pixels
+      by which the mask is cropped and smoothed. The smooth radar mask removes
+      the hard edges between NWP and radar in the final blended product.
     probmatching_method: {'cdf','mean',None}, optional
       Method for matching the statistics of the forecast field with those of
       the most recently observed one. 'cdf'=map the forecast CDF to the observed
@@ -1451,10 +1457,52 @@
                         # forecast outside the radar domain. Therefore, fill these
                         # areas with the "..._mod_only" blended forecasts, consisting
                         # of the NWP and noise components.
+
                         nan_indices = np.isnan(R_f_new)
-                        R_f_new[nan_indices] = R_f_new_mod_only[nan_indices]
-                        nan_indices = np.isnan(R_pm_blended)
-                        R_pm_blended[nan_indices] = R_pm_blended_mod_only[nan_indices]
+                        if smooth_radar_mask_range != 0:
+                            # Compute the smooth dilated mask
+                            new_mask = blending.utils.compute_smooth_dilated_mask(
+                                nan_indices,
+                                max_padding_size_in_px=smooth_radar_mask_range,
+                                gaussian_kernel_size=9,
+                                inverted=False,
+                                non_linear_growth_kernel_sizes=False,
+                            )
+
+                            # Ensure mask values are between 0 and 1
+                            mask_model = np.clip(new_mask, 0, 1)
+                            mask_radar = np.clip(1 - new_mask, 0, 1)
+
+                            # Handle NaNs in R_f_new and R_f_new_mod_only by setting NaNs to 0 in the blending step
+                            R_f_new_mod_only_no_nan = np.nan_to_num(
+                                R_f_new_mod_only, nan=0
+                            )
+                            R_f_new_no_nan = np.nan_to_num(R_f_new, nan=0)
+
+                            # Perform the blending of radar and model inside the radar domain using a weighted combination
+                            R_f_new = np.nansum(
+                                [
+                                    mask_model * R_f_new_mod_only_no_nan,
+                                    mask_radar * R_f_new_no_nan,
+                                ],
+                                axis=0,
+                            )
+
+                            nan_indices = np.isnan(R_pm_blended)
+                            R_pm_blended = np.nansum(
+                                [
+                                    R_pm_blended * mask_radar,
+                                    R_pm_blended_mod_only * mask_model,
+                                ],
+                                axis=0,
+                            )
+                        else:
+                            R_f_new[nan_indices] = R_f_new_mod_only[nan_indices]
+                            nan_indices = np.isnan(R_pm_blended)
+                            R_pm_blended[nan_indices] = R_pm_blended_mod_only[
+                                nan_indices
+                            ]
+
                         # Finally, fill the remaining nan values, if present, with
                         # the minimum value in the forecast
                         nan_indices = np.isnan(R_f_new)

pysteps/blending/utils.py

@@ -16,6 +16,7 @@
     compute_store_nwp_motion
     load_NWP
     check_norain
+    compute_smooth_dilated_mask
 """
 
 import datetime
@@ -35,6 +36,13 @@
 except ImportError:
     NETCDF4_IMPORTED = False
 
+try:
+    import cv2
+
+    CV2_IMPORTED = True
+except ImportError:
+    CV2_IMPORTED = False
+
 
 def stack_cascades(R_d, donorm=True):
     """Stack the given cascades into a larger array.
@@ -557,3 +565,86 @@
         f"Rain fraction is: {str(rain_pixels.size / precip_arr.size)}, while minimum fraction is {str(norain_thr)}"
     )
     return norain
+
+
+def compute_smooth_dilated_mask(
+    original_mask,
+    max_padding_size_in_px=0,
+    gaussian_kernel_size=9,
+    inverted=False,
+    non_linear_growth_kernel_sizes=False,
+):
+    """
+    Compute a smooth dilated mask using Gaussian blur and dilation with varying kernel sizes.
+
+    Parameters
+    ----------
+    original_mask : array_like
+        Two-dimensional boolean array containing the input mask.
+    max_padding_size_in_px : int
+        The maximum size of the padding in pixels. Default is 100.
+    gaussian_kernel_size : int, optional
+        Size of the Gaussian kernel to use for blurring. Default is 9.
+    inverted : bool, optional
+        If True, invert the original mask before processing. Default is False.
+    non_linear_growth_kernel_sizes : bool, optional
+        If True, use non-linear growth for kernel sizes. Default is False.
+
+    Returns
+    -------
+    initial_mask : array_like
+        The initial binary mask after Gaussian blur and thresholding.
+    new_mask : array_like
+        The smooth dilated mask normalized to the range [0,1].
+    """
+    if not CV2_IMPORTED:
+        raise MissingOptionalDependency(
+            "CV2 package is required to transform the mask into a smoot mask."
+            " Please install it using `pip install opencv-python`."
+        )
+
+    if max_padding_size_in_px < 0:
+        raise ValueError("max_padding_size_in_px must be greater than or equal to 0.")
+
+    # Convert the original mask to uint8 numpy array and invert if needed
+    array_2d = np.array(original_mask, dtype=np.uint8)
+    if inverted:
+        array_2d = np.bitwise_not(array_2d)
+
+    # Rescale the 2D array values to 0-255 (black or white)
+    rescaled_array = array_2d * 255
+
+    # Apply Gaussian blur to the rescaled array
+    blurred_image = cv2.GaussianBlur(
+        rescaled_array, (gaussian_kernel_size, gaussian_kernel_size), 0
+    )
+
+    # Apply binary threshold to negate the blurring effect
+    _, binary_image = cv2.threshold(blurred_image, 128, 255, cv2.THRESH_BINARY)
+
+    # Define kernel sizes
+    if non_linear_growth_kernel_sizes:
+        lin_space = np.linspace(0, np.sqrt(max_padding_size_in_px), 10)
+        non_lin_space = np.power(lin_space, 2)
+        kernel_sizes = list(set(non_lin_space.astype(np.uint8)))
+    else:
+        kernel_sizes = np.linspace(0, max_padding_size_in_px, 10, dtype=np.uint8)
+
+    # Process each kernel size
+    final_mask = np.zeros_like(binary_image, dtype=np.float64)
+    for kernel_size in kernel_sizes:
+        if kernel_size == 0:
+            dilated_image = binary_image
+        else:
+            kernel = cv2.getStructuringElement(
+                cv2.MORPH_ELLIPSE, (kernel_size, kernel_size)
+            )
+            dilated_image = cv2.dilate(binary_image, kernel)
+
+        # Convert the dilated image to a binary array
+        _, binary_array = cv2.threshold(dilated_image, 128, 1, cv2.THRESH_BINARY)
+        final_mask += binary_array
+
+    final_mask = final_mask / final_mask.max()
+
+    return final_mask

pysteps/nowcasts/utils.py

@@ -22,7 +22,6 @@
 
 from pysteps import extrapolation
 
-
 try:
     import dask
 

pysteps/tests/test_blending_steps.py

@@ -8,31 +8,39 @@
 
 
 steps_arg_values = [
-    (1, 3, 4, 8, None, None, False, "spn", True, 4, False, False),
-    (1, 3, 4, 8, "obs", None, False, "spn", True, 4, False, False),
-    (1, 3, 4, 8, "incremental", None, False, "spn", True, 4, False, False),
-    (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False),
-    (1, 3, 4, 8, None, "cdf", False, "spn", True, 4, False, False),
-    (1, 3, 4, 8, "incremental", "cdf", False, "spn", True, 4, False, False),
-    (1, 3, 4, 6, "incremental", "cdf", False, "bps", True, 4, False, False),
-    (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False),
-    (1, 3, 4, 9, "incremental", "cdf", False, "spn", True, 4, False, False),
-    (2, 3, 10, 8, "incremental", "cdf", False, "spn", True, 10, False, False),
-    (5, 3, 5, 8, "incremental", "cdf", False, "spn", True, 5, False, False),
-    (1, 10, 1, 8, "incremental", "cdf", False, "spn", True, 1, False, False),
-    (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False),
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False),
+    (1, 3, 4, 8, None, None, False, "spn", True, 4, False, False, 0),
+    (1, 3, 4, 8, "obs", None, False, "spn", True, 4, False, False, 0),
+    (1, 3, 4, 8, "incremental", None, False, "spn", True, 4, False, False, 0),
+    (1, 3, 4, 8, None, "mean", False, "spn", True, 4, False, False, 0),
+    (1, 3, 4, 8, None, "cdf", False, "spn", True, 4, False, False, 0),
+    (1, 3, 4, 8, "incremental", "cdf", False, "spn", True, 4, False, False, 0),
+    (1, 3, 4, 6, "incremental", "cdf", False, "bps", True, 4, False, False, 0),
+    (1, 3, 4, 6, "incremental", "cdf", False, "bps", False, 4, False, False, 0),
+    (1, 3, 4, 9, "incremental", "cdf", False, "spn", True, 4, False, False, 0),
+    (2, 3, 10, 8, "incremental", "cdf", False, "spn", True, 10, False, False, 0),
+    (5, 3, 5, 8, "incremental", "cdf", False, "spn", True, 5, False, False, 0),
+    (1, 10, 1, 8, "incremental", "cdf", False, "spn", True, 1, False, False, 0),
+    (2, 3, 2, 8, "incremental", "cdf", True, "spn", True, 2, False, False, 0),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 0),
     #    Test the case where the radar image contains no rain.
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, True, False, 0),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 0),
     #   Test the case where the NWP fields contain no rain.
-    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, True),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 0),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, True, 0),
     # 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),
-    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, True),
-    (5, 3, 5, 6, "obs", "mean", True, "spn", True, 5, True, True),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, True, True, 0),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, True, 0),
+    (5, 3, 5, 6, "obs", "mean", True, "spn", True, 5, True, True, 0),
+    # Test for smooth radar mask
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, False, 80),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, False, False, 80),
+    (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, False, False, 80),
+    (1, 3, 6, 8, None, None, False, "spn", True, 6, False, True, 80),
+    (5, 3, 5, 6, "incremental", "cdf", False, "spn", False, 5, True, False, 80),
+    (5, 3, 5, 6, "obs", "mean", False, "spn", False, 5, True, True, 80),
 ]
+
 steps_arg_names = (
     "n_models",
     "n_timesteps",
@@ -46,6 +54,7 @@
     "expected_n_ens_members",
     "zero_radar",
     "zero_nwp",
+    "smooth_radar_mask_range",
 )
 
 
@@ -63,6 +72,7 @@ def test_steps_blending(
     expected_n_ens_members,
     zero_radar,
     zero_nwp,
+    smooth_radar_mask_range,
 ):
     pytest.importorskip("cv2")