wrap all nowcasts in xarray

pysteps/io/importers.py

@@ -87,7 +87,7 @@
 +---------------+-------------------------------------------------------------------------------------------+
 |   time        | forecast time in seconds since forecast start time                                        |
 +---------------+-------------------------------------------------------------------------------------------+
-|   member      | ensemble member number (integer)                                                          |
+|   ens_number  | ensemble member number (integer)                                                          |
 +---------------+-------------------------------------------------------------------------------------------+
 
 

pysteps/nowcasts/lagrangian_probability.py

@@ -12,22 +12,20 @@
 """
 
 import numpy as np
+import xarray as xr
 from scipy.signal import convolve
 
 from pysteps.nowcasts import extrapolation
-from pysteps.xarray_decorators import xarray_nowcast
 
 
-@xarray_nowcast
 def forecast(
-    precip,
-    velocity,
+    dataset: xr.Dataset,
     timesteps,
     threshold,
     extrap_method="semilagrangian",
     extrap_kwargs=None,
     slope=5,
-):
+) -> xr.Dataset:
     """
     Generate a probability nowcast by a local lagrangian approach. The ouput is
     the probability of exceeding a given intensity threshold, i.e.
@@ -73,13 +71,11 @@ def forecast(
         timesteps = np.arange(1, timesteps + 1)
     elif not isinstance(timesteps, list):
         raise ValueError(f"invalid value for argument 'timesteps': {timesteps}")
-    precip_forecast = extrapolation.forecast(
-        precip,
-        velocity,
-        timesteps,
-        extrap_method,
-        extrap_kwargs,
+    dataset_forecast = extrapolation.forecast(
+        dataset, timesteps, extrap_method, extrap_kwargs
     )
+    precip_var = dataset_forecast.attrs["precip_var"]
+    precip_forecast = dataset_forecast[precip_var].values
 
     # Ignore missing values
     nanmask = np.isnan(precip_forecast)
@@ -106,7 +102,8 @@ def forecast(
         precip_forecast[i, ...] /= kernel_sum
     precip_forecast = np.clip(precip_forecast, 0, 1)
     precip_forecast[nanmask] = np.nan
-    return precip_forecast
+    dataset_forecast[precip_var].data[:] = precip_forecast
+    return dataset_forecast
 
 
 def _get_kernel(size):

pysteps/tests/test_nowcasts_anvil.py

@@ -31,29 +31,27 @@ def test_anvil_rainrate(
 ):
     """Tests ANVIL nowcast using rain rate precipitation fields."""
     # inputs
-    precip_input = get_precipitation_fields(
+    dataset_input = get_precipitation_fields(
         num_prev_files=4,
         num_next_files=0,
         return_raw=False,
         metadata=False,
         upscale=2000,
     )
-    precip_input = precip_input.filled()
 
-    precip_obs = get_precipitation_fields(
+    dataset_obs = get_precipitation_fields(
         num_prev_files=0, num_next_files=3, return_raw=False, upscale=2000
-    )[1:, :, :]
-    precip_obs = precip_obs.filled()
+    ).isel(time=slice(1, None, None))
+    precip_var = dataset_input.attrs["precip_var"]
 
     pytest.importorskip("cv2")
     oflow_method = motion.get_method("LK")
-    retrieved_motion = oflow_method(precip_input)
+    dataset_w_motion = oflow_method(dataset_input)
 
     nowcast_method = nowcasts.get_method("anvil")
 
     output = nowcast_method(
-        precip_input[-(ar_order + 2) :],
-        retrieved_motion,
+        dataset_w_motion.isel(time=slice(-(ar_order + 2), None, None)),
         timesteps=timesteps,
         rainrate=None,  # no R(VIL) conversion is done
         n_cascade_levels=n_cascade_levels,
@@ -63,17 +61,18 @@ def test_anvil_rainrate(
         measure_time=measure_time,
     )
     if measure_time:
-        precip_forecast, __, __ = output
+        dataset_forecast, __, __ = output
     else:
-        precip_forecast = output
+        dataset_forecast = output
+    precip_forecast = dataset_forecast[precip_var].values
 
     assert precip_forecast.ndim == 3
     assert precip_forecast.shape[0] == (
         timesteps if isinstance(timesteps, int) else len(timesteps)
     )
 
     result = verification.det_cat_fct(
-        precip_forecast[-1], precip_obs[-1], thr=0.1, scores="CSI"
+        precip_forecast[-1], dataset_obs[precip_var].values[-1], thr=0.1, scores="CSI"
     )["CSI"]
     assert result > min_csi, f"CSI={result:.2f}, required > {min_csi:.2f}"
 

pysteps/tests/test_nowcasts_lagrangian_probability.py

@@ -1,23 +1,33 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 import pytest
+import xarray as xr
 
+from pysteps.motion.lucaskanade import dense_lucaskanade
 from pysteps.nowcasts.lagrangian_probability import forecast
 from pysteps.tests.helpers import get_precipitation_fields
-from pysteps.motion.lucaskanade import dense_lucaskanade
 
 
 def test_numerical_example():
     """"""
     precip = np.zeros((20, 20))
     precip[5:10, 5:10] = 1
     velocity = np.zeros((2, *precip.shape))
+    dataset_input = xr.Dataset(
+        data_vars={
+            "precip_intensity": (["y", "x"], precip),
+            "velocity_x": (["y", "x"], velocity[0]),
+            "velocity_y": (["y", "x"], velocity[1]),
+        },
+        attrs={"precip_var": "precip_intensity"},
+    )
     timesteps = 4
     thr = 0.5
     slope = 1  # pixels / timestep
 
     # compute probability forecast
-    fct = forecast(precip, velocity, timesteps, thr, slope=slope)
+    dataset_forecast = forecast(dataset_input, timesteps, thr, slope=slope)
+    fct = dataset_forecast["precip_intensity"].values
 
     assert fct.ndim == 3
     assert fct.shape[0] == timesteps
@@ -26,7 +36,8 @@ def test_numerical_example():
     assert fct.min() >= 0.0
 
     # slope = 0 should return a binary field
-    fct = forecast(precip, velocity, timesteps, thr, slope=0)
+    dataset_forecast = forecast(dataset_input, timesteps, thr, slope=0)
+    fct = dataset_forecast["precip_intensity"].values
     ref = (np.repeat(precip[None, ...], timesteps, axis=0) >= thr).astype(float)
     assert np.allclose(fct, fct.astype(bool))
     assert np.allclose(fct, ref)
@@ -37,12 +48,21 @@ def test_numerical_example_with_float_slope_and_float_list_timesteps():
     precip = np.zeros((20, 20))
     precip[5:10, 5:10] = 1
     velocity = np.zeros((2, *precip.shape))
+    dataset_input = xr.Dataset(
+        data_vars={
+            "precip_intensity": (["y", "x"], precip),
+            "velocity_x": (["y", "x"], velocity[0]),
+            "velocity_y": (["y", "x"], velocity[1]),
+        },
+        attrs={"precip_var": "precip_intensity"},
+    )
     timesteps = [1.0, 2.0, 5.0, 12.0]
     thr = 0.5
     slope = 1.0  # pixels / timestep
 
     # compute probability forecast
-    fct = forecast(precip, velocity, timesteps, thr, slope=slope)
+    dataset_forecast = forecast(dataset_input, timesteps, thr, slope=slope)
+    fct = dataset_forecast["precip_intensity"].values
 
     assert fct.ndim == 3
     assert fct.shape[0] == len(timesteps)
@@ -56,16 +76,18 @@ def test_real_case():
     pytest.importorskip("cv2")
 
     # inputs
-    precip, metadata = get_precipitation_fields(
+    dataset_input = get_precipitation_fields(
         num_prev_files=2,
         num_next_files=0,
         return_raw=False,
         metadata=True,
         upscale=2000,
     )
+    precip_var = dataset_input.attrs["precip_var"]
+    metadata = dataset_input[precip_var].attrs
 
     # motion
-    motion = dense_lucaskanade(precip)
+    dataset_w_motion = dense_lucaskanade(dataset_input)
 
     # parameters
     timesteps = [1, 2, 3]
@@ -74,13 +96,18 @@ def test_real_case():
 
     # compute probability forecast
     extrap_kwargs = dict(allow_nonfinite_values=True)
-    fct = forecast(
-        precip[-1], motion, timesteps, thr, slope=slope, extrap_kwargs=extrap_kwargs
+    dataset_forecast = forecast(
+        dataset_w_motion.isel(time=-1).drop_vars(["time"]),
+        timesteps,
+        thr,
+        slope=slope,
+        extrap_kwargs=extrap_kwargs,
     )
+    fct = dataset_forecast["precip_intensity"].values
 
     assert fct.ndim == 3
     assert fct.shape[0] == len(timesteps)
-    assert fct.shape[1:] == precip.shape[1:]
+    assert fct.shape[1:] == dataset_input[precip_var].values.shape[1:]
     assert np.nanmax(fct) <= 1.0
     assert np.nanmin(fct) >= 0.0
 
@@ -89,11 +116,19 @@ def test_wrong_inputs():
     # dummy inputs
     precip = np.zeros((3, 3))
     velocity = np.zeros((2, *precip.shape))
+    dataset_input = xr.Dataset(
+        data_vars={
+            "precip_intensity": (["y", "x"], precip),
+            "velocity_x": (["y", "x"], velocity[0]),
+            "velocity_y": (["y", "x"], velocity[1]),
+        },
+        attrs={"precip_var": "precip_intensity"},
+    )
 
     # timesteps must be > 0
     with pytest.raises(ValueError):
-        forecast(precip, velocity, 0, 1)
+        forecast(dataset_input, 0, 1)
 
     # timesteps must be a sorted list
     with pytest.raises(ValueError):
-        forecast(precip, velocity, [2, 1], 1)
+        forecast(dataset_input, [2, 1], 1)