fix tests again

examples/ens_kalman_filter_blended_forecast.py

@@ -87,11 +87,9 @@
     + nwp_data_source["fn_ext"],
 )
 nwp_importer = io.get_method("dwd_nwp", "importer")
-kwargs = nwp_data_source["importer_kwargs"]
-# Resolve grid_file_path relative to PYSTEPS_DATA_PATH
-kwargs["grid_file_path"] = os.path.join(
-    os.environ["PYSTEPS_DATA_PATH"], kwargs["grid_file_path"]
-)
+# grid_file_path in importer_kwargs is already resolvable from the repo root,
+# like every other data_sources path in pystepsrc.
+kwargs = dict(nwp_data_source["importer_kwargs"])
 nwp_precip, _, nwp_metadata = nwp_importer(filename, **kwargs)
 # We lower the number of ens members to 10 to reduce the memory needs in the
 # example here. However, it is advised to have a minimum of 20 members for the
@@ -250,14 +248,18 @@
     n_cascade_levels=6,  # No. of cascade levels
     precip_thr=log_thr_prec,  # Precip threshold
     norain_thr=0.0005,  # Minimum of 0.5% precip needed, otherwise 'zero rainfall'
-    num_workers=4,  # No. of parallel threads
+    # No. of parallel threads. Kept at 1 so that the forecast (and the plots
+    # in this example) are exactly reproducible: multi-threaded FFT/dask
+    # reductions are not guaranteed to be bit-for-bit deterministic.
+    num_workers=1,
     noise_stddev_adj="auto",  # Standard deviation adjustment
     noise_method="ssft",  # SSFT as noise method
     enable_combination=True,  # Enable combination
     noise_kwargs={"win_size": (512, 512), "win_fun": "hann", "overlap": 0.5},
     extrap_kwargs={"interp_order": 3, "map_coordinates_mode": "nearest"},
     combination_kwargs=combination_kwargs,
     filter_kwargs={"include_mean": True},
+    seed=42,  # Fixed seed for reproducible ensemble members
 )
 
 # Transform the data back into mm/h

examples/linda_nowcasts.py

@@ -137,6 +137,7 @@
     n_ens_members=40,
     num_workers=8,
     measure_time=True,
+    seed=42,  # Fixed seed for reproducible ensemble members
 )[0]
 
 # Compute 40-member STEPS nowcast for comparison
@@ -151,6 +152,7 @@
     kmperpixel=2.0,
     timestep=datasource_params["timestep"],
     vel_pert_method=None,
+    seed=42,  # Fixed seed for reproducible ensemble members
 )
 
 # Convert reflectivity nowcast to rain rate

examples/rainfarm_downscale.py

@@ -114,6 +114,9 @@ def upscale_field(precip, metadata, scale_factor):
 #
 # We can now use RainFARM to downscale the precipitation field.
 
+# Fixed seed for reproducible stochastic noise fields
+np.random.seed(42)
+
 # Basic downscaling
 precip_hr = rainfarm.downscale(precip_lr, ds_factor=scale_factor)
 

examples/steps_blended_forecast.py

@@ -169,6 +169,7 @@
     kmperpixel=radar_metadata["xpixelsize"] / 1000.0,
     noise_stddev_adj="auto",
     vel_pert_method=None,
+    seed=42,  # Fixed seed for reproducible ensemble members
 )
 
 # Transform the data back into mm/h
@@ -375,6 +376,7 @@
     probmatching_method="cdf",
     mask_method="incremental",
     weights_method="bps",
+    seed=42,  # Fixed seed for reproducible ensemble members
 )
 
 # Transform the data back into mm/h

pysteps/blending/ens_kalman_filter_methods.py

@@ -433,6 +433,16 @@ def __init__(self, config, params):
         self.__degradation_timestep = 0.2
         self.__inflation_factor_obs_tmp = 1.0
 
+        # Dedicated, per-member random generators for the iterative probability
+        # matching resampling below, so that "seed" fully controls the
+        # reproducibility of the forecast instead of relying on the global,
+        # unseeded numpy random state.
+        seed = self._config.seed
+        self.__randgen_probmatch = [
+            np.random.RandomState(None if seed is None else seed + j)
+            for j in range(self._config.n_ens_members)
+        ]
+
         print("Initialize masked ensemble Kalman filter")
         print("========================================")
         print("")
@@ -603,6 +613,7 @@ def worker(j):
                     first_array=background_ensemble[j],
                     second_array=observation_ensemble[j],
                     probability_first_array=1 - self.__sampling_probability,
+                    randgen=self.__randgen_probmatch[j],
                 ).reshape(self.__params.len_y, self.__params.len_x)
 
             dask_worker_collection = []

pysteps/blending/pca_ens_kalman_filter.py

@@ -55,6 +55,7 @@
 from copy import deepcopy
 
 import numpy as np
+import xarray as xr
 from scipy.ndimage import (
     binary_dilation,
     gaussian_filter,
@@ -66,6 +67,7 @@
 from pysteps.postprocessing import probmatching
 from pysteps.timeseries import autoregression, correlation
 from pysteps.utils.check_norain import check_norain
+from pysteps.xarray_helpers import convert_output_to_xarray_dataset
 
 try:
     import dask
@@ -666,6 +668,15 @@ def __init__(
 
         self.__ens_member = ens_member
 
+        # Dedicated, per-member random generator for selecting noise fields from
+        # the pool in __iterate, so that the "seed" config option fully controls
+        # the reproducibility of the forecast instead of relying on the global,
+        # unseeded numpy random state.
+        seed = self.__forecast_state.config.seed
+        self.__randgen_select = np.random.RandomState(
+            None if seed is None else seed + ens_member
+        )
+
     # Bundle single steps of the forecast.
     def run_forecast_step(self, nwp, is_correction_timestep=False):
         # Decompose precipitation field.
@@ -813,7 +824,9 @@ def __iterate(self):
         # precipitation mask and the standard deviation coefficients.
         epsilon = (
             self.__forecast_state.noise_field_pool[
-                np.random.randint(self.__forecast_state.config.n_noise_fields)
+                self.__randgen_select.randint(
+                    self.__forecast_state.config.n_noise_fields
+                )
             ]
             * self.__forecast_state.precip_mask[self.__ens_member][None, :, :]
             * self.__forecast_state.params.noise_std_coeffs[:, None, None]
@@ -923,18 +936,32 @@ def fill_backtransform(self, nwp):
 class EnKFCombinationNowcaster:
     def __init__(
         self,
-        obs_precip: np.ndarray,
-        obs_timestamps: np.ndarray,
-        nwp_precip: np.ndarray,
-        nwp_timestamps: np.ndarray,
-        obs_velocity: np.ndarray,
+        radar_dataset: xr.Dataset,
+        model_dataset: xr.Dataset,
         fc_period: int,
         fc_init: datetime.datetime,
         enkf_combination_config: EnKFCombinationConfig,
     ):
         """
         Initialize EnKFCombinationNowcaster with inputs and configurations.
         """
+        # Extract the raw arrays that the rest of this class operates on from
+        # the input datasets.
+        radar_precip_var = radar_dataset.attrs["precip_var"]
+        model_precip_var = model_dataset.attrs["precip_var"]
+
+        obs_precip = radar_dataset[radar_precip_var].values
+        obs_timestamps = np.array(
+            radar_dataset["time"].values.astype("datetime64[us]").tolist()
+        )
+        nwp_precip = model_dataset[model_precip_var].values
+        nwp_timestamps = np.array(
+            model_dataset["time"].values.astype("datetime64[us]").tolist()
+        )
+        obs_velocity = np.array(
+            [radar_dataset["velocity_x"].values, radar_dataset["velocity_y"].values]
+        )
+
         # Store inputs
         self.__obs_precip = obs_precip
         self.__nwp_precip = nwp_precip
@@ -952,47 +979,26 @@ def __init__(
         self.__obs_timestamps = obs_timestamps
         self.__nwp_timestamps = nwp_timestamps
 
+        # Store the input radar dataset to use its coordinates/metadata when
+        # building the output dataset.
+        self.__input_radar_dataset = radar_dataset
+
     def compute_forecast(self):
         """
         Generate a combined nowcast ensemble by using the reduced-space ensemble Kalman
         filter method.
 
-        Parameters
-        ----------
-        obs_precip: np.ndarray
-            Array of shape (ar_order+1,m,n) containing the observed input precipitation
-            fields ordered by timestamp from oldest to newst. The time steps between
-            the inputs are assumed to be regular.
-        obs_timestamps: np.ndarray
-            Array of shape (ar_order+1) containing the corresponding time stamps of
-            observed input precipitation fields as datetime objects.
-        nwp_precip: np.ndarray
-            Array of shape (n_ens,n_times,m,n) containing the (NWP) ensemble model
-            forecast.
-        nwp_timestamps: np.ndarray
-            Array of shape (n_times) containing the corresponding time stamps of the
-            (NWP) ensemble model forecast as datetime objects.
-        obs_velocity: np.ndarray
-            Array of shape (2,m,n) containing the x- and y-components of the advection
-            field. The velocities are based on the observed input precipitation fields
-            and are assumed to represent one time step between the inputs. All values
-            are required to be finite.
-        fc_period: int
-            Forecast range in minutes.
-        fc_init: datetime object
-            Issuetime of the combined forecast to compute.
-        enkf_combination_config: EnKFCombinationConfig
-            Provides a set of configuration parameters for the nowcast ensemble
-            generation.
+        Uses the ``radar_dataset``, ``model_dataset``, ``fc_period``, ``fc_init``
+        and ``enkf_combination_config`` provided to the constructor.
 
         Returns
         -------
-        out: np.ndarray
-          If return_output is True, a four-dimensional array of shape
-          (n_ens_members,num_timesteps,m,n) containing a time series of forecast
-          precipitation fields for each ensemble member. Otherwise, a None value
-          is returned. The time series starts from t0. The timestep is taken from the
-          input precipitation fields precip.
+        out: xarray.Dataset
+          If return_output is True, a dataset as described in the documentation of
+          :py:mod:`pysteps.io.importers` is returned containing a time series of
+          forecast precipitation fields for each ensemble member. Otherwise, a
+          None value is returned. The time series starts from ``fc_init``, and the
+          timestep is taken from the input radar dataset.
 
         See also
         --------
@@ -1083,15 +1089,21 @@ def compute_forecast(self):
                 self.FS.final_combined_forecast
             ).swapaxes(0, 1)
 
+            result_dataset = convert_output_to_xarray_dataset(
+                self.__input_radar_dataset,
+                list(range(len(self.__forecast_leadtimes))),
+                self.FS.final_combined_forecast,
+            )
+
             if self.__config.measure_time:
                 return (
-                    self.FS.final_combined_forecast,
+                    result_dataset,
                     self.__fc_init,
                     self.__mainloop_time,
                 )
             if self.__config.verbose_output:
-                return self.FS.final_combined_forecast, self.FS.background_ensemble
-            return self.FS.final_combined_forecast
+                return result_dataset, self.FS.background_ensemble
+            return result_dataset
 
         # Else, return None
         return None
@@ -1553,11 +1565,8 @@ def __measure_time(self, label, start_time):
 
 
 def forecast(
-    obs_precip,
-    obs_timestamps,
-    nwp_precip,
-    nwp_timestamps,
-    velocity,
+    radar_dataset,
+    model_dataset,
     forecast_horizon,
     issuetime,
     n_ens_members,
@@ -1593,24 +1602,19 @@ def forecast(
 
     Parameters
     ----------
-    obs_precip: np.ndarray
-        Array of shape (ar_order+1,m,n) containing the observed input precipitation
-        fields ordered by timestamp from oldest to newst. The time steps between
-        the inputs are assumed to be regular.
-    obs_timestamps: np.ndarray
-        Array of shape (ar_order+1) containing the corresponding time stamps of
-        observed input precipitation fields as datetime objects.
-    nwp_precip: np.ndarray
-        Array of shape (n_ens,n_times,m,n) containing the (NWP) ensemble model
-        forecast.
-    nwp_timestamps: np.ndarray
-        Array of shape (n_times) containing the corresponding time stamps of the
-        (NWP) ensemble model forecast as datetime objects.
-    velocity: np.ndarray
-        Array of shape (2,m,n) containing the x- and y-components of the advection
-        field. The velocities are based on the observed input precipitation fields
-        and are assumed to represent one time step between the inputs. All values
-        are required to be finite.
+    radar_dataset: xarray.Dataset
+        Dataset containing the observed input precipitation fields as described
+        in the documentation of :py:mod:`pysteps.io.importers`. It has to contain
+        the ``velocity_x`` and ``velocity_y`` data variables (computed from the
+        observed precipitation fields, representing one time step between the
+        inputs), as well as any precipitation data variable. The time dimension
+        of the dataset has to be size ``ar_order + 1``, ordered by timestamp from
+        oldest to newest, and evenly spaced. All velocity values are required to
+        be finite.
+    model_dataset: xarray.Dataset
+        Dataset containing the (NWP) ensemble model forecast, as described in the
+        documentation of :py:mod:`pysteps.io.importers`. It has to contain an
+        ``ens_number`` dimension in addition to ``time``, ``y`` and ``x``.
     forecast_horizon: int
         The length of the forecast horizon (the length of the forecast) in minutes.
     issuetime: datetime object
@@ -1721,12 +1725,13 @@ def forecast(
 
     Returns
     -------
-    out: np.ndarray
-        If return_output is True, a four-dimensional array of shape
-        (n_ens_members,num_timesteps,m,n) containing a time series of forecast
-        precipitation fields for each ensemble member. Otherwise, a None value
-        is returned. The time series starts from t0. The timestep is taken from the
-        input precipitation fields precip.
+    out: xarray.Dataset
+        If return_output is True, a dataset as described in the documentation of
+        :py:mod:`pysteps.io.importers` is returned containing a time series of
+        forecast precipitation fields for each ensemble member. Otherwise, a None
+        value is returned. The time series starts from ``issuetime``, and the
+        timestep is taken from the metadata of the ``radar_dataset`` time
+        coordinate.
 
     See also
     --------
@@ -1775,11 +1780,8 @@ def forecast(
     )
 
     combination_nowcaster = EnKFCombinationNowcaster(
-        obs_precip=obs_precip,
-        obs_timestamps=obs_timestamps,
-        nwp_precip=nwp_precip,
-        nwp_timestamps=nwp_timestamps,
-        obs_velocity=velocity,
+        radar_dataset=radar_dataset,
+        model_dataset=model_dataset,
         fc_period=forecast_horizon,
         fc_init=issuetime,
         enkf_combination_config=combination_config,

pysteps/blending/steps.py

@@ -3793,6 +3793,8 @@ def forecast(
     # Check the input precip and ar_order to be consistent
     # zero-precip/constant field in previous time steps has to be removed
     # (constant field causes autoregression to fail)
+    radar_precip_var = radar_dataset.attrs["precip_var"]
+    precip = radar_dataset[radar_precip_var].values
     precip, ar_order = check_previous_radar_obs(
         precip,
         ar_order,
@@ -3803,6 +3805,10 @@ def forecast(
             "win_fun": "tukey" if noise_kwargs is None else noise_kwargs["win_fun"],
         },
     )
+    radar_dataset = radar_dataset.isel(time=slice(-precip.shape[0], None)).copy(
+        deep=True
+    )
+    radar_dataset[radar_precip_var].values = precip
 
     blending_config = StepsBlendingConfig(
         n_ens_members=n_ens_members,