Draft: xarray/main

ci/ci_test_env.yml

@@ -18,6 +18,7 @@ dependencies:
   - pillow
   - pyproj
   - scipy
+  - xarray
   # Optional dependencies
   - dask
   - pyfftw

environment.yml

@@ -13,3 +13,4 @@ dependencies:
   - pillow
   - pyproj
   - scipy
+  - xarray

environment_dev.yml

@@ -21,7 +21,6 @@ dependencies:
   - dask
   - pyfftw
   - h5py
-  - PyWavelets
   - pygrib
   - black
   - pytest-cov
@@ -31,3 +30,6 @@ dependencies:
   - scikit-image
   - pandas
   - rasterio
+  - xarray
+  - geotiff
+  - cookiecutter

pysteps/datasets.py

@@ -434,10 +434,8 @@ def load_dataset(case="fmi", frames=14):
 
     Returns
     -------
-    rainrate: array-like
-        Precipitation data in mm/h. Dimensions: [time, lat, lon]
-    metadata: dict
-        The metadata observations attributes.
+    rainrate: Dataset
+        A dataset containing the precipitation data in mm/h and quality rasters and associated metadata. Dimensions: [time, lat, lon]
     timestep: number
         Time interval between composites in minutes.
     """
@@ -476,11 +474,11 @@ def load_dataset(case="fmi", frames=14):
     # Read the radar composites
     importer = io.get_method(data_source["importer"], "importer")
     importer_kwargs = data_source["importer_kwargs"]
-    reflectivity, _, metadata = io.read_timeseries(
-        file_names, importer, **importer_kwargs
+    reflectivity = io.read_timeseries(
+        file_names, importer, timestep=data_source["timestep"], **importer_kwargs
     )
 
     # Convert to rain rate
-    precip, metadata = conversion.to_rainrate(reflectivity, metadata)
+    precip = conversion.to_rainrate(reflectivity)
 
-    return precip, metadata, data_source["timestep"]
+    return precip, data_source["timestep"]

pysteps/decorators.py

@@ -22,6 +22,8 @@
 
 import numpy as np
 
+from pysteps.xarray_helpers import convert_input_to_xarray_dataset
+
 
 def _add_extra_kwrds_to_docstrings(target_func, extra_kwargs_doc_text):
     """
@@ -30,6 +32,8 @@ def _add_extra_kwrds_to_docstrings(target_func, extra_kwargs_doc_text):
     """
     # Clean up indentation from docstrings for the
     # docstrings to be merged correctly.
+    if target_func.__doc__ is None:
+        return target_func
     extra_kwargs_doc = inspect.cleandoc(extra_kwargs_doc_text)
     target_func.__doc__ = inspect.cleandoc(target_func.__doc__)
 
@@ -66,7 +70,7 @@ def postprocess_import(fillna=np.nan, dtype="double"):
     def _postprocess_import(importer):
         @wraps(importer)
         def _import_with_postprocessing(*args, **kwargs):
-            precip, *other_args = importer(*args, **kwargs)
+            precip, quality, metadata = importer(*args, **kwargs)
 
             _dtype = kwargs.get("dtype", dtype)
 
@@ -88,7 +92,9 @@ def _import_with_postprocessing(*args, **kwargs):
                     mask = ~np.isfinite(precip)
                     precip[mask] = _fillna
 
-            return (precip.astype(_dtype),) + tuple(other_args)
+            return convert_input_to_xarray_dataset(
+                precip.astype(_dtype), quality, metadata
+            )
 
         extra_kwargs_doc = """
             Other Parameters
@@ -124,7 +130,9 @@ def new_function(*args, **kwargs):
             target motion_method_func function.
             """
 
-            input_images = args[0]
+            dataset = args[0]
+            precip_var = dataset.attrs["precip_var"]
+            input_images = dataset[precip_var].values
             if input_images.ndim != 3:
                 raise ValueError(
                     "input_images dimension mismatch.\n"

pysteps/downscaling/rainfarm.py

@@ -18,11 +18,12 @@
 """
 
 import warnings
-
+import xarray as xr
 import numpy as np
 from scipy.signal import convolve
 from pysteps.utils.spectral import rapsd
 from pysteps.utils.dimension import aggregate_fields
+from pysteps.xarray_helpers import compute_lat_lon
 
 
 def _gaussianize(precip):
@@ -210,7 +211,7 @@ def _balanced_spatial_average(array, kernel):
 
 
 def downscale(
-    precip,
+    precip_dataset: xr.Dataset,
     ds_factor,
     alpha=None,
     threshold=None,
@@ -224,8 +225,8 @@ def downscale(
 
     Parameters
     ----------
-    precip: array_like
-        Array of shape (m, n) containing the input field.
+    precip: xarray.Dataset
+        Xarray dataset containing the input field.
         The input is expected to contain rain rate values.
         All values are required to be finite.
     alpha: float, optional
@@ -266,6 +267,8 @@ def downscale(
     """
 
     # Validate inputs
+    precip_var = precip_dataset.attrs["precip_var"]
+    precip = precip_dataset[precip_var].values.squeeze()
     if not np.isfinite(precip).all():
         raise ValueError("All values in 'precip' must be finite.")
     if not isinstance(ds_factor, int) or ds_factor <= 0:
@@ -297,8 +300,50 @@ def downscale(
     noise_field /= noise_field.std()
     noise_field = np.exp(noise_field)
 
+    y_new = np.arange(
+        precip_dataset.y.values[0]
+        - precip_dataset.y.attrs["stepsize"]
+        + precip_dataset.y.attrs["stepsize"] / ds_factor,
+        precip_dataset.y.values[-1] + precip_dataset.y.attrs["stepsize"] / ds_factor,
+        precip_dataset.y.attrs["stepsize"] / ds_factor,
+    )
+    x_new = np.arange(
+        precip_dataset.x.values[0]
+        - precip_dataset.y.attrs["stepsize"]
+        + precip_dataset.x.attrs["stepsize"] / ds_factor,
+        precip_dataset.x.values[-1] + precip_dataset.x.attrs["stepsize"] / ds_factor,
+        precip_dataset.x.attrs["stepsize"] / ds_factor,
+    )
+    lat, lon = compute_lat_lon(x_new, y_new, precip_dataset.attrs["projection"])
+    noise_dataset = xr.Dataset(
+        data_vars={precip_var: (["time", "y", "x"], [noise_field])},
+        coords={
+            "time": (["time"], precip_dataset.time.values, precip_dataset.time.attrs),
+            "y": (
+                ["y"],
+                y_new,
+                {
+                    **precip_dataset.y.attrs,
+                    "stepsize": precip_dataset.y.attrs["stepsize"] / ds_factor,
+                },
+            ),
+            "x": (
+                ["x"],
+                x_new,
+                {
+                    **precip_dataset.x.attrs,
+                    "stepsize": precip_dataset.x.attrs["stepsize"] / ds_factor,
+                },
+            ),
+            "lon": (["y", "x"], lon, precip_dataset.lon.attrs),
+            "lat": (["y", "x"], lat, precip_dataset.lat.attrs),
+        },
+        attrs=precip_dataset.attrs,
+    )
+
     # Aggregate the noise field to low resolution
-    noise_lowres = aggregate_fields(noise_field, ds_factor, axis=(0, 1))
+    noise_lowres_dataset = aggregate_fields(noise_dataset, ds_factor, dim=("y", "x"))
+    noise_lowres = noise_lowres_dataset[precip_var].values.squeeze()
 
     # Expand input and noise fields to high resolution
     precip_expanded = np.kron(precip, np.ones((ds_factor, ds_factor)))
@@ -322,8 +367,11 @@ def downscale(
     if threshold is not None:
         precip_highres[precip_highres < threshold] = 0
 
+    precip_highres_dataset = noise_dataset.copy(deep=True)
+    precip_highres_dataset[precip_var][:] = [precip_highres]
+
     # Return the downscaled field and optionally the spectral slope alpha
     if return_alpha:
-        return precip_highres, alpha
+        return precip_highres_dataset, alpha
 
-    return precip_highres
+    return precip_highres_dataset

pysteps/io/exporters.py

@@ -511,12 +511,12 @@ def initialize_forecast_exporter_netcdf(
 
     if metadata["unit"] == "mm/h":
         var_name = "precip_intensity"
-        var_standard_name = None
+        var_standard_name = "instantaneous_precipitation_rate"
         var_long_name = "instantaneous precipitation rate"
         var_unit = "mm h-1"
     elif metadata["unit"] == "mm":
         var_name = "precip_accum"
-        var_standard_name = None
+        var_standard_name = "accumulated_precipitation"
         var_long_name = "accumulated precipitation"
         var_unit = "mm"
     elif metadata["unit"] == "dBZ":