Xarray/fix even more tests (#513)

* fix test_importer_decorator

* fixed some more tests

* remove test.nc

Author: mats-knmi

pysteps/decorators.py

@@ -32,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__)
 

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":

pysteps/io/importers.py

@@ -352,7 +352,6 @@ def _get_threshold_value(precip):
         return np.nan
 
 
-@postprocess_import(dtype="float32")
 def import_mrms_grib(filename, extent=None, window_size=4, **kwargs):
     """
     Importer for NSSL's Multi-Radar/Multi-Sensor System
@@ -407,7 +406,14 @@ def import_mrms_grib(filename, extent=None, window_size=4, **kwargs):
         image dimensions.
         Default: window_size=4.
 
-    {extra_kwargs_doc}
+    Other Parameters
+    ----------------
+    dtype: str
+        Data-type to which the array is cast.
+        Valid values:  "float32", "float64", "single", and "double".
+    fillna: float or np.nan
+        Value used to represent the missing data ("No Coverage").
+        By default, np.nan is used.
 
     Returns
     -------
@@ -420,7 +426,32 @@ def import_mrms_grib(filename, extent=None, window_size=4, **kwargs):
     metadata: dict
         Associated metadata (pixel sizes, map projections, etc.).
     """
+    dataset = _import_mrms_grib(filename, extent, window_size, **kwargs)
+    # Create a function with default arguments for aggregate_fields
+    block_reduce = partial(aggregate_fields, method="mean", trim=True)
+
+    if window_size != (1, 1):
+        # Downscale data
+        precip_var = dataset.attrs["precip_var"]
+        # block_reduce does not handle nan values
+        if "fillna" in kwargs:
+            no_data_mask = dataset[precip_var].values == kwargs["fillna"]
+        else:
+            no_data_mask = np.isnan(dataset[precip_var].values)
+        dataset[precip_var].data[no_data_mask] = 0
+        dataset["no_data_mask"] = (("y", "x"), no_data_mask)
+        dataset = block_reduce(dataset, window_size, dim=("y", "x"))
+
+        # Consider that if a single invalid observation is located in the block,
+        # then mark that value as invalid.
+        no_data_mask = dataset.no_data_mask.values == 1.0
+        dataset = dataset.drop_vars("no_data_mask")
+
+    return dataset
+
 
+@postprocess_import(dtype="float32")
+def _import_mrms_grib(filename, extent=None, window_size=4, **kwargs):
     del kwargs
 
     if not PYGRIB_IMPORTED:
@@ -466,32 +497,6 @@ def import_mrms_grib(filename, extent=None, window_size=4, **kwargs):
     precip = grib_msg.values
     no_data_mask = precip == -3  # Missing values
 
-    # Create a function with default arguments for aggregate_fields
-    block_reduce = partial(aggregate_fields, method="mean", trim=True)
-
-    if window_size != (1, 1):
-        # Downscale data
-        lats = block_reduce(lats, window_size[0])
-        lons = block_reduce(lons, window_size[1])
-
-        # Update the limits
-        ul_lat, lr_lat = (
-            lats[0],
-            lats[-1],
-        )  # Lat from North to south!
-        ul_lon, lr_lon = lons[0], lons[-1]
-
-        precip[no_data_mask] = 0  # block_reduce does not handle nan values
-        precip = block_reduce(precip, window_size, axis=(0, 1))
-
-        # Consider that if a single invalid observation is located in the block,
-        # then mark that value as invalid.
-        no_data_mask = block_reduce(
-            no_data_mask.astype("int"),
-            window_size,
-            axis=(0, 1),
-        ).astype(bool)
-
     lons, lats = np.meshgrid(lons, lats)
     precip[no_data_mask] = np.nan
 

pysteps/tests/test_downscaling_rainfarm.py

@@ -8,13 +8,12 @@
 
 
 @pytest.fixture(scope="module")
-def data():
-    precip, metadata = get_precipitation_fields(
+def dataset():
+    precip_dataset = get_precipitation_fields(
         num_prev_files=0, num_next_files=0, return_raw=False, metadata=True
     )
-    precip = precip.filled()
-    precip, metadata = square_domain(precip, metadata, "crop")
-    return precip, metadata
+    precip_dataset = square_domain(precip_dataset, "crop")
+    return precip_dataset
 
 
 rainfarm_arg_names = (
@@ -35,7 +34,7 @@ def data():
 
 @pytest.mark.parametrize(rainfarm_arg_names, rainfarm_arg_values)
 def test_rainfarm_shape(
-    data,
+    dataset,
     alpha,
     ds_factor,
     threshold,
@@ -44,13 +43,13 @@ def test_rainfarm_shape(
     kernel_type,
 ):
     """Test that the output of rainfarm is consistent with the downscaling factor."""
-    precip, metadata = data
-    window = metadata["xpixelsize"] * ds_factor
-    precip_lr, __ = aggregate_fields_space(precip, metadata, window)
+    precip_var = dataset.attrs["precip_var"]
+    window = dataset.x.attrs["stepsize"] * ds_factor
+    precip_lr_dataset = aggregate_fields_space(dataset, window)
 
     rainfarm = downscaling.get_method("rainfarm")
-    precip_hr = rainfarm(
-        precip_lr,
+    precip_hr_dataset = rainfarm(
+        precip_lr_dataset,
         alpha=alpha,
         ds_factor=ds_factor,
         threshold=threshold,
@@ -59,9 +58,15 @@ def test_rainfarm_shape(
         kernel_type=kernel_type,
     )
 
-    assert precip_hr.ndim == precip.ndim
-    assert precip_hr.shape[0] == precip.shape[0]
-    assert precip_hr.shape[1] == precip.shape[1]
+    assert precip_hr_dataset[precip_var].values.ndim == dataset[precip_var].values.ndim
+    assert (
+        precip_hr_dataset[precip_var].values.shape[0]
+        == dataset[precip_var].values.shape[0]
+    )
+    assert (
+        precip_hr_dataset[precip_var].values.shape[1]
+        == dataset[precip_var].values.shape[1]
+    )
 
 
 rainfarm_arg_values = [
@@ -74,7 +79,7 @@ def test_rainfarm_shape(
 
 @pytest.mark.parametrize(rainfarm_arg_names, rainfarm_arg_values)
 def test_rainfarm_aggregate(
-    data,
+    dataset,
     alpha,
     ds_factor,
     threshold,
@@ -83,22 +88,24 @@ def test_rainfarm_aggregate(
     kernel_type,
 ):
     """Test that the output of rainfarm is equal to original when aggregated."""
-    precip, metadata = data
-    window = metadata["xpixelsize"] * ds_factor
-    precip_lr, __ = aggregate_fields_space(precip, metadata, window)
+    precip_var = dataset.attrs["precip_var"]
+    window = dataset.x.attrs["stepsize"] * ds_factor
+    precip_lr_dataset = aggregate_fields_space(dataset, window)
 
     rainfarm = downscaling.get_method("rainfarm")
-    precip_hr = rainfarm(
-        precip_lr,
+    precip_hr_dataset = rainfarm(
+        precip_lr_dataset,
         alpha=alpha,
         ds_factor=ds_factor,
         threshold=threshold,
         return_alpha=return_alpha,
         spectral_fusion=spectral_fusion,
         kernel_type=kernel_type,
     )
-    precip_low = aggregate_fields(precip_hr, ds_factor, axis=(0, 1))
+    precip_low_dataset = aggregate_fields(precip_hr_dataset, ds_factor, dim=("y", "x"))
+    precip_lr = precip_lr_dataset[precip_var].values
     precip_lr[precip_lr < threshold] = 0.0
+    precip_low = precip_low_dataset[precip_var].values
 
     np.testing.assert_array_almost_equal(precip_lr, precip_low)
 
@@ -108,7 +115,7 @@ def test_rainfarm_aggregate(
 
 @pytest.mark.parametrize(rainfarm_arg_names, rainfarm_arg_values)
 def test_rainfarm_alpha(
-    data,
+    dataset,
     alpha,
     ds_factor,
     threshold,
@@ -117,13 +124,12 @@ def test_rainfarm_alpha(
     kernel_type,
 ):
     """Test that rainfarm computes and returns alpha."""
-    precip, metadata = data
-    window = metadata["xpixelsize"] * ds_factor
-    precip_lr, __ = aggregate_fields_space(precip, metadata, window)
+    window = dataset.x.attrs["stepsize"] * ds_factor
+    precip_lr_dataset = aggregate_fields_space(dataset, window)
 
     rainfarm = downscaling.get_method("rainfarm")
-    precip_hr = rainfarm(
-        precip_lr,
+    precip_hr_dataset = rainfarm(
+        precip_lr_dataset,
         alpha=alpha,
         ds_factor=ds_factor,
         threshold=threshold,
@@ -132,8 +138,8 @@ def test_rainfarm_alpha(
         kernel_type=kernel_type,
     )
 
-    assert len(precip_hr) == 2
+    assert len(precip_hr_dataset) == 2
     if alpha is None:
-        assert not precip_hr[1] == alpha
+        assert not precip_hr_dataset[1] == alpha
     else:
-        assert precip_hr[1] == alpha
+        assert precip_hr_dataset[1] == alpha

pysteps/tests/test_feature.py

@@ -14,14 +14,16 @@ def test_feature(method, max_num_features):
     if method == "shitomasi":
         pytest.importorskip("cv2")
 
-    input_field = get_precipitation_fields(
+    input_dataset = get_precipitation_fields(
         num_prev_files=0,
         num_next_files=0,
         return_raw=True,
         metadata=False,
         upscale=None,
         source="mch",
     )
+    precip_var = input_dataset.attrs["precip_var"]
+    input_field = input_dataset[precip_var].values
 
     detector = feature.get_method(method)