Support for mixed precip/noprecip input and AR order 1

pysteps/blending/steps.py

@@ -1196,12 +1196,22 @@ def transform_to_lagrangian(precip, i):
             self.__precip_models = np.stack(temp_precip_models)
 
         # Check for zero input fields in the radar, nowcast and NWP data.
+        # decide based on latest input file only
         self.__params.zero_precip_radar = check_norain(
-            self.__precip,
+            self.__precip[-1],
             self.__params.precip_threshold,
             self.__config.norain_threshold,
             self.__params.noise_kwargs["win_fun"],
         )
+        # Set all inputs to 0 (also previous times)
+        # not strictly necessary but to be consistent with checking only
+        # latest observation rain field to set zero_precip_radar
+        if self.__params.zero_precip_radar:
+            self.__precip = np.where(
+                np.isfinite(self.__precip),
+                np.ones(self.__precip.shape) * np.nanmin(self.__precip),
+                self.__precip,
+            )
 
         # The norain fraction threshold used for nwp is the default value of 0.0,
         # since nwp does not suffer from clutter.
@@ -1481,6 +1491,8 @@ def __estimate_ar_parameters_radar(self):
 
             # Finally, transpose GAMMA to ensure that the shape is the same as np.empty((n_cascade_levels, ar_order))
             GAMMA = GAMMA.transpose()
+            if len(GAMMA.shape) == 1:
+                GAMMA = GAMMA.reshape((GAMMA.size, 1))
             assert GAMMA.shape == (
                 self.__config.n_cascade_levels,
                 self.__config.ar_order,
@@ -3625,6 +3637,11 @@ def forecast(
     turns out to be a warranted functionality.
     """
 
+    # Check the input precip and ar_order to be consistent
+    # zero-precip in previous time steps has to be removed
+    # (constant field causes autoregression to fail)
+    precip, ar_order = check_previous_radar_obs(precip, ar_order)
+
     blending_config = StepsBlendingConfig(
         n_ens_members=n_ens_members,
         n_cascade_levels=n_cascade_levels,
@@ -3686,6 +3703,54 @@ def forecast(
     return forecast_steps_nowcast
 
 
+# TODO: Where does this piece of code best fit: in utils or inside the class?
+def check_previous_radar_obs(precip, ar_order):
+    """Check all radar time steps and remove zero precipitation time steps
+
+    Parameters
+    ----------
+    precip : array-like
+      Array of shape (ar_order+1,m,n) containing the input precipitation fields
+      ordered by timestamp from oldest to newest. The time steps between the
+      inputs are assumed to be regular.
+    ar_order : int
+      The order of the autoregressive model to use. Must be >= 1.
+
+    Returns
+    -------
+    precip : numpy array
+      Array of shape (ar_order+1,m,n) containing the modified array with
+      input precipitation fields ordered by timestamp from oldest to newest.
+      The time steps between the inputs are assumed to be regular.
+    ar_order : int
+      The order of the autoregressive model to use. Must be >= 1.
+      Adapted to match with precip.shape equal (ar_order+1,m,n).
+    """
+    # Quick check radar input - at least 2 time steps
+    assert precip.shape[0] >= 2
+
+    # Check all time steps for zero-precip (constant field, minimum==maximum)
+    zero_precip = [np.nanmin(obs) == np.nanmax(obs) for obs in precip]
+    if zero_precip[-1] or ~np.any(zero_precip):
+        # Unchanged if no rain in latest time step -> will be processed as zero_precip_radar=True
+        # or Unchanged if all time steps contain rain
+        return precip, ar_order
+    elif zero_precip[-2]:
+        # try to use a previous time step
+        if not np.all(zero_precip[:-2]):
+            # find latest non-zero precip
+            # ATTENTION: This changes the time between precip[-2] and precip[-1] from initial 5min to a longer period
+            prev = np.arange(len(zero_precip[:-2]))[~np.array(zero_precip[:-2])][-1]
+            return precip[[prev, -1]], 1
+        raise ValueError(
+            "Precipitation in latest but no previous time step. Not possible to calculate autoregression."
+        )
+    else:
+        # Keep the latest time steps that do all contain precip
+        precip = precip[np.max(np.arange(len(zero_precip))[zero_precip]) + 1 :].copy()
+        return precip, min(ar_order, precip.shape[0] - 1)
+
+
 # TODO: Where does this piece of code best fit: in utils or inside the class?
 def calculate_ratios(correlations):
     """Calculate explained variance ratios from correlation.

pysteps/tests/test_blending_steps.py

@@ -77,6 +77,35 @@
 
 # fmt:on
 
+
+def run_and_assert_forecast(
+    precip, forecast_kwargs, expected_n_ens_members, n_timesteps, converter, metadata
+):
+    """Run a blended nowcast and assert the output has the expected shape."""
+    precip_forecast = blending.steps.forecast(precip=precip, **forecast_kwargs)
+
+    assert precip_forecast.ndim == 4, "Wrong amount of dimensions in forecast output"
+    assert (
+        precip_forecast.shape[0] == expected_n_ens_members
+    ), "Wrong amount of output ensemble members in forecast output"
+    assert (
+        precip_forecast.shape[1] == n_timesteps
+    ), "Wrong amount of output time steps in forecast output"
+
+    # Transform the data back into mm/h
+    precip_forecast, _ = converter(precip_forecast, metadata)
+
+    assert (
+        precip_forecast.ndim == 4
+    ), "Wrong amount of dimensions in converted forecast output"
+    assert (
+        precip_forecast.shape[0] == expected_n_ens_members
+    ), "Wrong amount of output ensemble members in converted forecast output"
+    assert (
+        precip_forecast.shape[1] == n_timesteps
+    ), "Wrong amount of output time steps in converted forecast output"
+
+
 steps_arg_names = (
     "n_models",
     "timesteps",
@@ -357,10 +386,9 @@ def test_steps_blending(
     assert nwp_velocity.ndim == 5, "nwp_velocity must be a five-dimensional array"
 
     ###
-    # The blending
+    # Shared forecast kwargs
     ###
-    precip_forecast = blending.steps.forecast(
-        precip=radar_precip,
+    forecast_kwargs = dict(
         precip_models=nwp_precip_decomp,
         velocity=radar_velocity,
         velocity_models=nwp_velocity,
@@ -402,27 +430,125 @@ def test_steps_blending(
         measure_time=False,
     )
 
-    assert precip_forecast.ndim == 4, "Wrong amount of dimensions in forecast output"
+    ###
+    # The blending
+    ###
+    # Test with full radar data
+    run_and_assert_forecast(
+        radar_precip,
+        forecast_kwargs,
+        expected_n_ens_members,
+        n_timesteps,
+        converter,
+        metadata,
+    )
 
-    assert (
-        precip_forecast.shape[0] == expected_n_ens_members
-    ), "Wrong amount of output ensemble members in forecast output"
 
-    assert (
-        precip_forecast.shape[1] == n_timesteps
-    ), "Wrong amount of output time steps in forecast output"
+@pytest.mark.parametrize("ar_order", [1, 2])
+def test_steps_blending_partial_zero_radar(ar_order):
+    """Test that a forecast succeeds when only the latest radar frame has
+    precipitation (initiating cell corner case that produces NaN autocorrelations
+    for the earlier, all-zero cascade levels)."""
+    pytest.importorskip("cv2")
 
-    # Transform the data back into mm/h
-    precip_forecast, _ = converter(precip_forecast, metadata)
+    n_timesteps = 3
+    metadata = dict(
+        unit="mm",
+        transformation="dB",
+        accutime=5.0,
+        transform="dB",
+        zerovalue=0.0,
+        threshold=0.01,
+        zr_a=200.0,
+        zr_b=1.6,
+    )
 
-    assert (
-        precip_forecast.ndim == 4
-    ), "Wrong amount of dimensions in converted forecast output"
+    # Build minimal NWP data (1 model, 4 time steps)
+    nwp_precip = np.zeros((1, n_timesteps + 1, 200, 200))
+    for i in range(nwp_precip.shape[1]):
+        nwp_precip[0, i, 30:185, 32 + i] = 1.0
+        nwp_precip[0, i, 30:185, 33 + i] = 5.0
+        nwp_precip[0, i, 30:185, 34 + i] = 5.0
+        nwp_precip[0, i, 30:185, 35 + i] = 4.5
 
-    assert (
-        precip_forecast.shape[0] == expected_n_ens_members
-    ), "Wrong amount of output ensemble members in converted forecast output"
+    # Build radar data: only the latest (most recent) frame has precipitation
+    radar_precip = np.zeros((3, 200, 200))
+    radar_precip[2, 30:155, 32] = 1.0
+    radar_precip[2, 30:155, 33] = 5.0
+    radar_precip[2, 30:155, 34] = 5.0
+    radar_precip[2, 30:155, 35] = 4.5
 
-    assert (
-        precip_forecast.shape[1] == n_timesteps
-    ), "Wrong amount of output time steps in converted forecast output"
+    # Threshold, convert and transform
+    radar_precip[radar_precip < metadata["threshold"]] = 0.0
+    nwp_precip[nwp_precip < metadata["threshold"]] = 0.0
+    converter = pysteps.utils.get_method("mm/h")
+    radar_precip, _ = converter(radar_precip, metadata)
+    nwp_precip, metadata = converter(nwp_precip, metadata)
+    transformer = pysteps.utils.get_method(metadata["transformation"])
+    radar_precip, _ = transformer(radar_precip, metadata)
+    nwp_precip, metadata = transformer(nwp_precip, metadata)
+    radar_precip[~np.isfinite(radar_precip)] = metadata["zerovalue"]
+    nwp_precip[~np.isfinite(nwp_precip)] = metadata["zerovalue"]
+
+    # Decompose NWP
+    n_cascade_levels = 6
+    decomp_method, _ = cascade.get_method("fft")
+    bp_filter = cascade.get_method("gaussian")(radar_precip.shape[1:], n_cascade_levels)
+    nwp_precip_decomp = nwp_precip.copy()
+
+    # Velocity fields
+    oflow_method = pysteps.motion.get_method("lucaskanade")
+    radar_velocity = oflow_method(radar_precip)
+    _V_NWP = [
+        oflow_method(nwp_precip[0, t - 1 : t + 1, :])
+        for t in range(1, nwp_precip.shape[1])
+    ]
+    nwp_velocity = np.stack([np.insert(_V_NWP, 0, _V_NWP[0], axis=0)])
+
+    run_and_assert_forecast(
+        radar_precip,
+        dict(
+            precip_models=nwp_precip_decomp,
+            velocity=radar_velocity,
+            velocity_models=nwp_velocity,
+            timesteps=n_timesteps,
+            timestep=5.0,
+            issuetime=datetime.datetime.strptime("202112012355", "%Y%m%d%H%M"),
+            n_ens_members=2,
+            n_cascade_levels=n_cascade_levels,
+            blend_nwp_members=False,
+            precip_thr=metadata["threshold"],
+            kmperpixel=1.0,
+            extrap_method="semilagrangian",
+            decomp_method="fft",
+            bandpass_filter_method="gaussian",
+            noise_method="nonparametric",
+            noise_stddev_adj="auto",
+            ar_order=ar_order,
+            vel_pert_method=None,
+            weights_method="bps",
+            conditional=False,
+            probmatching_method=None,
+            mask_method="incremental",
+            resample_distribution=False,
+            smooth_radar_mask_range=0,
+            callback=None,
+            return_output=True,
+            seed=42,
+            num_workers=1,
+            fft_method="numpy",
+            domain="spatial",
+            outdir_path_skill="./tmp/",
+            extrap_kwargs=None,
+            filter_kwargs=None,
+            noise_kwargs=None,
+            vel_pert_kwargs=None,
+            clim_kwargs=None,
+            mask_kwargs=None,
+            measure_time=False,
+        ),
+        expected_n_ens_members=2,
+        n_timesteps=n_timesteps,
+        converter=converter,
+        metadata=metadata,
+    )