Refactor steps in blending code

.gitignore

@@ -91,3 +91,7 @@ venv.bak/
 
 # Mac OS Stuff
 .DS_Store
+
+# Running local tests
+/tmp
+/pysteps/tests/tmp/

PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.2
 Name: pysteps
-Version: 1.13.0
+Version: 1.14.0
 Summary: Python framework for short-term ensemble prediction systems
 Home-page: http://pypi.python.org/pypi/pysteps/
 License: LICENSE

pysteps/motion/__init__.py

@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
 """
-    Implementations of optical flow methods."""
+Implementations of optical flow methods."""
 
 from .interface import get_method

pysteps/nowcasts/steps.py

@@ -14,12 +14,12 @@
 import numpy as np
 from scipy.ndimage import generate_binary_structure, iterate_structure
 import time
+from copy import deepcopy
 
 from pysteps import cascade
 from pysteps import extrapolation
 from pysteps import noise
 from pysteps import utils
-from pysteps.decorators import deprecate_args
 from pysteps.nowcasts import utils as nowcast_utils
 from pysteps.postprocessing import probmatching
 from pysteps.timeseries import autoregression, correlation
@@ -36,7 +36,7 @@
     DASK_IMPORTED = False
 
 
-@dataclass
+@dataclass(frozen=True)
 class StepsNowcasterConfig:
     """
     Parameters
@@ -248,6 +248,10 @@ class StepsNowcasterParams:
     xy_coordinates: np.ndarray | None = None
     velocity_perturbation_parallel: list[float] | None = None
     velocity_perturbation_perpendicular: list[float] | None = None
+    filter_kwargs: dict | None = None
+    noise_kwargs: dict | None = None
+    velocity_perturbation_kwargs: dict | None = None
+    mask_kwargs: dict | None = None
 
 
 @dataclass
@@ -269,6 +273,7 @@ class StepsNowcasterState:
     )
     velocity_perturbations: list[Callable] | None = field(default_factory=list)
     fft_objects: list[Any] | None = field(default_factory=list)
+    extrapolation_kwargs: dict[str, Any] | None = field(default_factory=dict)
 
 
 class StepsNowcaster:
@@ -341,9 +346,9 @@ def compute_forecast(self):
         if self.__config.measure_time:
             self.__start_time_init = time.time()
 
-        self.__initialize_nowcast_components()
         # Slice the precipitation field to only use the last ar_order + 1 fields
         self.__precip = self.__precip[-(self.__config.ar_order + 1) :, :, :].copy()
+        self.__initialize_nowcast_components()
 
         self.__perform_extrapolation()
         self.__apply_noise_and_ar_model()
@@ -352,15 +357,19 @@ def compute_forecast(self):
         self.__initialize_fft_objects()
         # Measure and print initialization time
         if self.__config.measure_time:
-            self.__measure_time("Initialization", self.__start_time_init)
+            self.__init_time = self.__measure_time(
+                "Initialization", self.__start_time_init
+            )
 
         # Run the main nowcast loop
         self.__nowcast_main()
 
+        # Unstack nowcast output if return_output is True
         if self.__config.measure_time:
-            self.__state.precip_forecast, self.__mainloop_time = (
-                self.__state.precip_forecast
-            )
+            (
+                self.__state.precip_forecast,
+                self.__mainloop_time,
+            ) = self.__state.precip_forecast
 
         # Stack and return the forecast output
         if self.__config.return_output:
@@ -386,14 +395,14 @@ def __nowcast_main(self):
         Main nowcast loop that iterates through the ensemble members and time steps
         to generate forecasts.
         """
-        # Isolate the last time slice of precipitation
+        # Isolate the last time slice of observed precipitation
         precip = self.__precip[
             -1, :, :
         ]  # Extract the last available precipitation field
 
         # Prepare state and params dictionaries, these need to be formatted a specific way for the nowcast_main_loop
-        state = self.__initialize_state()
-        params = self.__initialize_params(precip)
+        state = self.__return_state_dict()
+        params = self.__return_params_dict(precip)
 
         print("Starting nowcast computation.")
 
@@ -405,7 +414,7 @@ def __nowcast_main(self):
             self.__time_steps,
             self.__config.extrapolation_method,
             self.__update_state,  # Reference to the update function
-            extrap_kwargs=self.__config.extrapolation_kwargs,
+            extrap_kwargs=self.__state.extrapolation_kwargs,
             velocity_pert_gen=self.__state.velocity_perturbations,
             params=params,
             ensemble=True,
@@ -480,15 +489,33 @@ def __check_inputs(self):
 
         # Handle None values for various kwargs
         if self.__config.extrapolation_kwargs is None:
-            self.__config.extrapolation_kwargs = {}
+            self.__state.extrapolation_kwargs = dict()
+        else:
+            self.__state.extrapolation_kwargs = deepcopy(
+                self.__config.extrapolation_kwargs
+            )
+
         if self.__config.filter_kwargs is None:
-            self.__config.filter_kwargs = {}
+            self.__params.filter_kwargs = dict()
+        else:
+            self.__params.filter_kwargs = deepcopy(self.__config.filter_kwargs)
+
         if self.__config.noise_kwargs is None:
-            self.__config.noise_kwargs = {}
+            self.__params.noise_kwargs = dict()
+        else:
+            self.__params.noise_kwargs = deepcopy(self.__config.noise_kwargs)
+
         if self.__config.velocity_perturbation_kwargs is None:
-            self.__config.velocity_perturbation_kwargs = {}
+            self.__params.velocity_perturbation_kwargs = dict()
+        else:
+            self.__params.velocity_perturbation_kwargs = deepcopy(
+                self.__config.velocity_perturbation_kwargs
+            )
+
         if self.__config.mask_kwargs is None:
-            self.__config.mask_kwargs = {}
+            self.__params.mask_kwargs = dict()
+        else:
+            self.__params.mask_kwargs = deepcopy(self.__config.mask_kwargs)
 
         print("Inputs validated and initialized successfully.")
 
@@ -545,12 +572,12 @@ def __print_forecast_info(self):
 
         if self.__config.velocity_perturbation_method == "bps":
             self.__params.velocity_perturbation_parallel = (
-                self.__config.velocity_perturbation_kwargs.get(
+                self.__params.velocity_perturbation_kwargs.get(
                     "p_par", noise.motion.get_default_params_bps_par()
                 )
             )
             self.__params.velocity_perturbation_perpendicular = (
-                self.__config.velocity_perturbation_kwargs.get(
+                self.__params.velocity_perturbation_kwargs.get(
                     "p_perp", noise.motion.get_default_params_bps_perp()
                 )
             )
@@ -585,13 +612,14 @@ def __initialize_nowcast_components(self):
         self.__params.bandpass_filter = filter_method(
             (M, N),
             self.__config.n_cascade_levels,
-            **(self.__config.filter_kwargs or {}),
+            **(self.__params.filter_kwargs or {}),
         )
 
         # Get the decomposition method (e.g., FFT)
-        self.__params.decomposition_method, self.__params.recomposition_method = (
-            cascade.get_method(self.__config.decomposition_method)
-        )
+        (
+            self.__params.decomposition_method,
+            self.__params.recomposition_method,
+        ) = cascade.get_method(self.__config.decomposition_method)
 
         # Get the extrapolation method (e.g., semilagrangian)
         self.__params.extrapolation_method = extrapolation.get_method(
@@ -625,7 +653,7 @@ def __perform_extrapolation(self):
         else:
             self.__state.mask_threshold = None
 
-        extrap_kwargs = self.__config.extrapolation_kwargs.copy()
+        extrap_kwargs = self.__state.extrapolation_kwargs.copy()
         extrap_kwargs["xy_coords"] = self.__params.xy_coordinates
         extrap_kwargs["allow_nonfinite_values"] = (
             True if np.any(~np.isfinite(self.__precip)) else False
@@ -687,7 +715,7 @@ def __apply_noise_and_ar_model(self):
             self.__params.perturbation_generator = init_noise(
                 self.__precip,
                 fft_method=self.__params.fft,
-                **self.__config.noise_kwargs,
+                **self.__params.noise_kwargs,
             )
 
             # Handle noise standard deviation adjustments if necessary
@@ -715,7 +743,7 @@ def __apply_noise_and_ar_model(self):
 
                 # Measure and print time taken
                 if self.__config.measure_time:
-                    self.__measure_time(
+                    __ = self.__measure_time(
                         "Noise adjustment coefficient computation", starttime
                     )
                 else:
@@ -827,21 +855,16 @@ def __apply_noise_and_ar_model(self):
         if self.__config.noise_method is not None:
             self.__state.random_generator_precip = []
             self.__state.random_generator_motion = []
-
+            seed = self.__config.seed
             for _ in range(self.__config.n_ens_members):
                 # Create random state for precipitation noise generator
-                rs = np.random.RandomState(self.__config.seed)
+                rs = np.random.RandomState(seed)
                 self.__state.random_generator_precip.append(rs)
-                self.__config.seed = rs.randint(
-                    0, high=int(1e9)
-                )  # Update seed after generating
-
+                seed = rs.randint(0, high=int(1e9))
                 # Create random state for motion perturbations generator
-                rs = np.random.RandomState(self.__config.seed)
+                rs = np.random.RandomState(seed)
                 self.__state.random_generator_motion.append(rs)
-                self.__config.seed = rs.randint(
-                    0, high=int(1e9)
-                )  # Update seed after generating
+                seed = rs.randint(0, high=int(1e9))
         else:
             self.__state.random_generator_precip = None
             self.__state.random_generator_motion = None
@@ -861,10 +884,10 @@ def __initialize_velocity_perturbations(self):
             for j in range(self.__config.n_ens_members):
                 kwargs = {
                     "randstate": self.__state.random_generator_motion[j],
-                    "p_par": self.__config.velocity_perturbation_kwargs.get(
+                    "p_par": self.__params.velocity_perturbation_kwargs.get(
                         "p_par", self.__params.velocity_perturbation_parallel
                     ),
-                    "p_perp": self.__config.velocity_perturbation_kwargs.get(
+                    "p_perp": self.__params.velocity_perturbation_kwargs.get(
                         "p_perp", self.__params.velocity_perturbation_perpendicular
                     ),
                 }
@@ -916,8 +939,8 @@ def __initialize_precipitation_mask(self):
 
             elif self.__config.mask_method == "incremental":
                 # Get mask parameters
-                self.__params.mask_rim = self.__config.mask_kwargs.get("mask_rim", 10)
-                mask_f = self.__config.mask_kwargs.get("mask_f", 1.0)
+                self.__params.mask_rim = self.__params.mask_kwargs.get("mask_rim", 10)
+                mask_f = self.__params.mask_kwargs.get("mask_f", 1.0)
                 # Initialize the structuring element
                 self.__params.structuring_element = generate_binary_structure(2, 1)
                 # Expand the structuring element based on mask factor and timestep
@@ -957,7 +980,7 @@ def __initialize_fft_objects(self):
             self.__state.fft_objs.append(fft_obj)
         print("FFT objects initialized successfully.")
 
-    def __initialize_state(self):
+    def __return_state_dict(self):
         """
         Initialize the state dictionary used during the nowcast iteration.
         """
@@ -971,7 +994,7 @@ def __initialize_state(self):
             "randgen_prec": self.__state.random_generator_precip,
         }
 
-    def __initialize_params(self, precip):
+    def __return_params_dict(self, precip):
         """
         Initialize the params dictionary used during the nowcast iteration.
         """
@@ -1196,6 +1219,8 @@ def __measure_time(self, label, start_time):
         if self.__config.measure_time:
             elapsed_time = time.time() - start_time
             print(f"{label} took {elapsed_time:.2f} seconds.")
+            return elapsed_time
+        return None
 
     def reset_states_and_params(self):
         """
@@ -1214,7 +1239,6 @@ def reset_states_and_params(self):
 
 
 # Wrapper function to preserve backward compatibility
-@deprecate_args({"R": "precip", "V": "velocity", "R_thr": "precip_thr"}, "1.8.0")
 def forecast(
     precip,
     velocity,

pysteps/postprocessing/probmatching.py

@@ -289,7 +289,7 @@ def resample_distributions(
         cascade). It must be of the same shape as `second_array`. Input must not contain NaNs.
     second_array: array_like
         One of the two arrays from which the distribution should be sampled (e.g., the NWP (model)
-        cascade). It must be of the same shape as `first_array`.. Input must not contain NaNs.
+        cascade). It must be of the same shape as `first_array`. Input must not contain NaNs.
     probability_first_array: float
         The weight that `first_array` should get (a value between 0 and 1). This determines the
         likelihood of selecting elements from `first_array` over `second_array`.

pysteps/tests/test_blending_steps.py

@@ -28,6 +28,9 @@
     (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),
     #    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),

pysteps/tests/test_motion_lk.py

@@ -1,7 +1,6 @@
 # coding: utf-8
 
-"""
-"""
+""" """
 
 import pytest
 import numpy as np