diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1e915e89e..6421a7ef1 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/pysteps/blending/steps.py b/pysteps/blending/steps.py index 21a0ffc05..c4a1f6b21 100644 --- a/pysteps/blending/steps.py +++ b/pysteps/blending/steps.py @@ -90,6 +90,7 @@ def forecast( conditional=False, probmatching_method="cdf", mask_method="incremental", + smooth_radar_mask_range=0, callback=None, return_output=True, seed=None, @@ -210,6 +211,13 @@ def forecast( '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. + 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 + not present anymore or is not reliable. If set to 0 (grid cells), this generates a normal forecast without smoothing. 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. Typically, a value between 50 and 100 km can be used. 80 km generally gives good results. 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 +1459,49 @@ def worker(j): # 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, + ) + + # 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) diff --git a/pysteps/blending/utils.py b/pysteps/blending/utils.py index 6349f193f..a7b56494b 100644 --- a/pysteps/blending/utils.py +++ b/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,92 @@ def check_norain(precip_arr, precip_thr=None, norain_thr=0.0): 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, this should be an uneven number. This option ensures + that the nan-fields are large enough to start the smoothing. Without it, the method will also be applied + to local nan-values in the radar domain. Default is 9, which is generally a recommended number to work + with. + inverted : bool, optional + Typically, the smoothed mask works from the outside of the radar domain inward, using the + max_padding_size_in_px. If set to True, it works from the edge of the radar domain outward + (generally not recommended). Default is False. + non_linear_growth_kernel_sizes : bool, optional + If True, use non-linear growth for kernel sizes. Default is False. + + Returns + ------- + final_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.") + + # Check if gaussian_kernel_size is an uneven number + assert gaussian_kernel_size % 2 + + # 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 diff --git a/pysteps/nowcasts/utils.py b/pysteps/nowcasts/utils.py index 7d7b8dc0d..fd111e28d 100644 --- a/pysteps/nowcasts/utils.py +++ b/pysteps/nowcasts/utils.py @@ -22,7 +22,6 @@ from pysteps import extrapolation - try: import dask diff --git a/pysteps/tests/test_blending_steps.py b/pysteps/tests/test_blending_steps.py index 3f4a0580b..ea20eaf25 100644 --- a/pysteps/tests/test_blending_steps.py +++ b/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") @@ -254,6 +264,7 @@ def test_steps_blending( conditional=False, probmatching_method=probmatching_method, mask_method=mask_method, + smooth_radar_mask_range=smooth_radar_mask_range, callback=None, return_output=True, seed=None, diff --git a/pysteps/tests/test_blending_utils.py b/pysteps/tests/test_blending_utils.py index 63fc1ffa5..22d0a0045 100644 --- a/pysteps/tests/test_blending_utils.py +++ b/pysteps/tests/test_blending_utils.py @@ -16,6 +16,7 @@ compute_store_nwp_motion, load_NWP, check_norain, + compute_smooth_dilated_mask, ) pytest.importorskip("netCDF4") @@ -140,12 +141,27 @@ ) ] +smoothing_arg_names = ( + "precip_nwp", + "max_padding_size_in_px", + "gaussian_kernel_size", + "inverted", + "non_linear_growth_kernel_sizes", +) + +smoothing_arg_values = [ + (precip_nwp, 80, 9, False, False), + (precip_nwp, 10, 9, False, False), + (precip_nwp, 80, 5, False, False), + (precip_nwp, 80, 9, True, False), + (precip_nwp, 80, 9, False, True), +] + ### # The test ### @pytest.mark.parametrize(utils_arg_names, utils_arg_values) - # The test function to be used def test_blending_utils( precip_nwp, @@ -404,3 +420,28 @@ def test_blending_utils( # should always give norain if the threshold is set to 100% assert check_norain(precip_arr, norain_thr=1.0) + + +# Finally, also test the compute_smooth_dilated mask functionality +@pytest.mark.parametrize(smoothing_arg_names, smoothing_arg_values) +def test_blending_smoothing_utils( + precip_nwp, + max_padding_size_in_px, + gaussian_kernel_size, + inverted, + non_linear_growth_kernel_sizes, +): + # First add some nans to indicate a mask + precip_nwp[:, 0:100, 0:100] = np.nan + nan_indices = np.isnan(precip_nwp[0]) + new_mask = compute_smooth_dilated_mask( + nan_indices, + max_padding_size_in_px=max_padding_size_in_px, + gaussian_kernel_size=gaussian_kernel_size, + inverted=inverted, + non_linear_growth_kernel_sizes=non_linear_growth_kernel_sizes, + ) + assert new_mask.shape == nan_indices.shape + + if max_padding_size_in_px > 0 and inverted == False: + assert np.sum((new_mask > 0) & (new_mask < 1)) > 0