pysteps.io with xarray (#219)

* Add xarray dependency

* MCH importer returns an xarray Dataset

* Remove plot lines

* Remove import

* Adapt readers to xarray format

* Rewrite as more general decorator

* Add missing metadata

* Adapt io tests

* Mrms bounding box (#222)

* Fix bounding box coordinates

* Add missing metadata

* Import xarray DataArray

Ignore quality field

* Black

* Do not hardcode metadata

* Address review comments by ruben

* Add a legacy option to the io functions

A "legacy" options is added to revert back the importers and readers behavior to version 1. This is a temporary solution to allow the examples, and other functions, to run as usual (v1.*).

Hopefully, this is will allow a easier transition into version 2 during the development process and will allow testing functions that were not updated to v2.

* Fix compatibility problems with tests

Many of the tests were updated to use the legacy data structures (v1). The tests that still contains issues, were tagged with a TODO message and they are skipped.

This will allow incremental changes to be tested in the new v2.

IMPORTANT: once the v2 branch is stable, we may remove the legacy compatibility from the code base and the tests.

* Update dependencies

* Ignore plugins test

Co-authored-by: Andres Perez Hortal <16256571+aperezhortal@users.noreply.github.com>

.github/workflows/test_pysteps.yml

@@ -8,9 +8,13 @@ env:
 on:
   # Triggers the workflow on push or pull request events but only for the master branch
   push:
-    branches: [ master ]
+    branches:
+      - master
+      - pysteps-v2
   pull_request:
-    branches: [ master ]
+      branches:
+        - master
+        - pysteps-v2
 
 jobs:
   unit_tests:

environment.yml

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

environment_dev.yml

@@ -30,3 +30,4 @@ dependencies:
   - cartopy>=0.18
   - scikit-image
   - pandas
+  - xarray

examples/LK_buffer_mask.py

@@ -55,7 +55,7 @@
 
 # Read the radar composites
 importer = io.get_method(importer_name, "importer")
-R, quality, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+R, quality, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 del quality  # Not used
 
@@ -210,7 +210,7 @@
 )
 
 # Read and convert the radar composites
-R_o, _, metadata_o = io.read_timeseries(fns, importer, **importer_kwargs)
+R_o, _, metadata_o = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 R_o, metadata_o = conversion.to_rainrate(R_o, metadata_o)
 
 # Compute Spearman correlation

examples/advection_correction.py

@@ -60,7 +60,7 @@
 
 # Read the radar composites
 importer = io.get_method(importer_name, "importer")
-R, __, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+R, __, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 # Convert to mm/h
 R, metadata = conversion.to_rainrate(R, metadata)

examples/anvil_nowcast.py

@@ -48,7 +48,7 @@
 # Read the input time series
 importer = io.get_method(importer_name, "importer")
 rainrate_field, quality, metadata = io.read_timeseries(
-    filenames, importer, **importer_kwargs
+    filenames, importer, legacy=True, **importer_kwargs
 )
 
 # Convert to rain rate (mm/h)
@@ -118,8 +118,8 @@
     date, root_path, path_fmt, fn_pattern, fn_ext, timestep=5, num_next_files=3
 )
 
-refobs_field, quality, metadata = io.read_timeseries(
-    filenames, importer, **importer_kwargs
+refobs_field, _, metadata = io.read_timeseries(
+    filenames, importer, legacy=True, **importer_kwargs
 )
 
 refobs_field, metadata = utils.to_rainrate(refobs_field[-1], metadata)

examples/data_transformations.py

@@ -58,7 +58,7 @@
 
 # Read the radar composites
 importer = io.get_method(importer_name, "importer")
-Z, _, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+Z, _, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 # Keep only positive rainfall values
 Z = Z[Z > metadata["zerovalue"]].flatten()

examples/optical_flow_methods_convergence.py

@@ -62,7 +62,7 @@
 # Read the reference radar composite
 importer = io.get_method(importer_name, "importer")
 reference_field, quality, metadata = io.read_timeseries(
-    fns, importer, **importer_kwargs
+    fns, importer, legacy=True, **importer_kwargs
 )
 
 del quality  # Not used

examples/plot_cascade_decomposition.py

@@ -30,7 +30,7 @@
 filename = os.path.join(
     root_path, "20160928", "201609281600_fmi.radar.composite.lowest_FIN_SUOMI1.pgm.gz"
 )
-R, _, metadata = io.import_fmi_pgm(filename, gzipped=True)
+R, _, metadata = io.import_fmi_pgm(filename, gzipped=True, legacy=True)
 
 # Convert to rain rate
 R, metadata = conversion.to_rainrate(R, metadata)

examples/plot_ensemble_verification.py

@@ -56,7 +56,7 @@
 
 # Read the data from the archive
 importer = io.get_method(importer_name, "importer")
-R, _, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+R, _, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 # Convert to rain rate
 R, metadata = conversion.to_rainrate(R, metadata)
@@ -142,7 +142,7 @@
 )
 
 # Read the observations
-R_o, _, metadata_o = io.read_timeseries(fns, importer, **importer_kwargs)
+R_o, _, metadata_o = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 # Convert to mm/h
 R_o, metadata_o = conversion.to_rainrate(R_o, metadata_o)

examples/plot_extrapolation_nowcast.py

@@ -48,7 +48,7 @@
 
 # Read the radar composites
 importer = io.get_method(importer_name, "importer")
-Z, _, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+Z, _, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 # Convert to rain rate
 R, metadata = conversion.to_rainrate(Z, metadata)
@@ -113,7 +113,7 @@
     num_next_files=n_leadtimes,
 )
 # Read the radar composites
-R_o, _, metadata_o = io.read_timeseries(fns, importer, **importer_kwargs)
+R_o, _, metadata_o = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 R_o, metadata_o = conversion.to_rainrate(R_o, metadata_o, 223.0, 1.53)
 
 # Compute fractions skill score (FSS) for all lead times, a set of scales and 1 mm/h

examples/plot_noise_generators.py

@@ -34,7 +34,9 @@
 # Import the example radar composite
 root_path = rcparams.data_sources["mch"]["root_path"]
 filename = os.path.join(root_path, "20160711", "AQC161932100V_00005.801.gif")
-R, _, metadata = io.import_mch_gif(filename, product="AQC", unit="mm", accutime=5.0)
+R, _, metadata = io.import_mch_gif(
+    filename, product="AQC", unit="mm", accutime=5.0, legacy=True
+)
 
 # Convert to mm/h
 R, metadata = conversion.to_rainrate(R, metadata)

examples/plot_optical_flow.py

@@ -47,7 +47,7 @@
 
 # Read the radar composites
 importer = io.get_method(importer_name, "importer")
-R, quality, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+R, quality, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 del quality  # Not used
 

examples/plot_steps_nowcast.py

@@ -51,7 +51,7 @@
 
 # Read the data from the archive
 importer = io.get_method(importer_name, "importer")
-R, _, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+R, _, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 # Convert to rain rate
 R, metadata = conversion.to_rainrate(R, metadata)

examples/probability_forecast.py

@@ -62,7 +62,7 @@
 date = datetime.strptime("201607112100", "%Y%m%d%H%M")
 fns = io.find_by_date(date, root, fmt, pattern, ext, timestep, num_prev_files=2)
 importer = io.get_method(importer_name, "importer")
-precip, __, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+precip, __, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 precip, metadata = utils.to_rainrate(precip, metadata)
 # precip[np.isnan(precip)] = 0
 
@@ -106,7 +106,7 @@
 fns = io.find_by_date(
     date, root, fmt, pattern, ext, timestep, num_next_files=nleadtimes
 )
-obs, __, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+obs, __, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 obs, metadata = utils.to_rainrate(obs, metadata)
 obs[np.isnan(obs)] = 0
 

examples/rainfarm_downscale.py

@@ -33,7 +33,7 @@
 root_path = rcparams.data_sources["mch"]["root_path"]
 filename = os.path.join(root_path, "20160711", "AQC161932100V_00005.801.gif")
 precip, _, metadata = io.import_mch_gif(
-    filename, product="AQC", unit="mm", accutime=5.0
+    filename, product="AQC", unit="mm", accutime=5.0, legacy=True
 )
 
 # Convert to mm/h

examples/thunderstorm_detection_and_tracking.py

@@ -63,7 +63,7 @@
     date, root_path, path_fmt, fn_pattern, fn_ext, timestep, num_next_files=20
 )
 importer = io.get_method(importer_name, "importer")
-R, _, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+R, _, metadata = io.read_timeseries(fns, importer, legacy=True, **importer_kwargs)
 
 # Convert to reflectivity (it is possible to give the a- and b- parameters of the
 # Marshall-Palmer relationship here: zr_a = and zr_b =).

pysteps/decorators.py

@@ -16,6 +16,7 @@
 """
 import inspect
 import uuid
+import xarray as xr
 from collections import defaultdict
 from functools import wraps
 
@@ -67,7 +68,11 @@ def _postprocess_import(importer):
         @wraps(importer)
         def _import_with_postprocessing(*args, **kwargs):
 
-            precip, *other_args = importer(*args, **kwargs)
+            data_array = importer(*args, **kwargs)
+
+            if not isinstance(data_array, xr.DataArray):
+                array, _, metadata = data_array
+                data_array = _to_xarray(array, metadata)
 
             _dtype = kwargs.get("dtype", dtype)
 
@@ -78,18 +83,15 @@ def _import_with_postprocessing(*args, **kwargs):
                     "The accepted values are: " + str(accepted_precisions)
                 )
 
-            if isinstance(precip, np.ma.MaskedArray):
-                invalid_mask = np.ma.getmaskarray(precip)
-                precip.data[invalid_mask] = fillna
-            else:
-                # If plain numpy arrays are used, the importers should indicate
-                # the invalid values with np.nan.
-                _fillna = kwargs.get("fillna", fillna)
-                if _fillna is not np.nan:
-                    mask = ~np.isfinite(precip)
-                    precip[mask] = _fillna
+            _fillna = kwargs.get("fillna", fillna)
+            if _fillna is not np.nan:
+                data_array = data_array.fillna(_fillna)
+
+            data_array = data_array.astype(_dtype)
 
-            return (precip.astype(_dtype),) + tuple(other_args)
+            if kwargs.get("legacy", False):
+                return _xarray2legacy(data_array)
+            return data_array
 
         extra_kwargs_doc = """
             Other Parameters
@@ -283,3 +285,84 @@ def _func_with_cache(*args, **kwargs):
         return _func_with_cache
 
     return _memoize
+
+
+def _to_xarray(array, metadata):
+    """Convert to xarray DataArray."""
+    x1, x2 = metadata["x1"], metadata["x2"]
+    y1, y2 = metadata["y1"], metadata["y2"]
+    xsize, ysize = metadata["xpixelsize"], metadata["ypixelsize"]
+    # x_coords = np.arange(x1, x2, xsize) + xsize / 2
+    # y_coords = np.arange(y1, y2, ysize) + ysize / 2
+    x_coords = np.arange(x1, x1 + xsize * array.shape[1], xsize) + xsize / 2
+    y_coords = np.arange(y1, y1 + ysize * array.shape[0], ysize) + ysize / 2
+
+    data_array = xr.DataArray(
+        data=array,
+        dims=("y", "x"),
+        coords=dict(
+            x=("x", x_coords),
+            y=("y", y_coords),
+        ),
+    )
+
+    data_array.attrs.update(
+        {
+            "unit": metadata["unit"],
+            "accutime": metadata["accutime"],
+            "transform": metadata["transform"],
+            "zerovalue": metadata["zerovalue"],
+            "threshold": metadata["threshold"],
+            "zr_a": metadata.get("zr_a", None),
+            "zr_b": metadata.get("zr_b", None),
+            "institution": metadata["institution"],
+            "projection": metadata["projection"],
+            #
+            # TODO: Remove before final 2.0 version
+            "yorigin": metadata["yorigin"],
+            "xpixelsize": metadata["xpixelsize"],
+            "ypixelsize": metadata["ypixelsize"],
+        }
+    )
+
+    data_array.x.attrs.update(
+        {
+            "standard_name": "projection_x_coordinate",
+            "units": metadata["cartesian_unit"],
+            #
+            # TODO: Remove before final 2.0 version
+            "x1": x1,
+            "x2": x2,
+        }
+    )
+
+    data_array.y.attrs.update(
+        {
+            "standard_name": "projection_y_coordinate",
+            "units": metadata["cartesian_unit"],
+            #
+            # TODO: Remove before final 2.0 version
+            "y1": y1,
+            "y2": y2,
+        }
+    )
+
+    return data_array
+
+
+# TODO: Remove before final 2.0 version
+def _xarray2legacy(data_array):
+    """
+    Convert the new DataArrays to the legacy format used in pysteps v1.*
+    """
+    _array = data_array.values
+
+    metadata = data_array.x.attrs.copy()
+    metadata.update(**data_array.y.attrs)
+    metadata.update(**data_array.attrs)
+
+    if "t" in data_array.coords:
+        print(data_array["t"])
+        metadata["timestamps"] = data_array["t"]
+
+    return _array, None, metadata

pysteps/io/importers.py

@@ -90,6 +90,7 @@
 from functools import partial
 
 import numpy as np
+import xarray as xr
 from matplotlib.pyplot import imread
 
 from pysteps.decorators import postprocess_import
@@ -392,22 +393,27 @@ def import_mrms_grib(filename, extent=None, window_size=4, **kwargs):
     pr = pyproj.Proj(proj_params)
     proj_def = " ".join([f"+{key}={value} " for key, value in proj_params.items()])
 
+    xsize = grib_msg["iDirectionIncrementInDegrees"] * window_size[0]
+    ysize = grib_msg["jDirectionIncrementInDegrees"] * window_size[1]
+
     x1, y1 = pr(ul_lon, lr_lat)
     x2, y2 = pr(lr_lon, ul_lat)
 
     metadata = dict(
-        xpixelsize=grib_msg["iDirectionIncrementInDegrees"] * window_size[0],
-        ypixelsize=grib_msg["jDirectionIncrementInDegrees"] * window_size[1],
+        institution="NOAA National Severe Storms Laboratory",
+        xpixelsize=xsize,
+        ypixelsize=ysize,
         unit="mm/h",
+        accutime=2.0,
         transform=None,
         zerovalue=0,
         projection=proj_def.strip(),
         yorigin="upper",
         threshold=_get_threshold_value(precip),
-        x1=x1,
-        x2=x2,
-        y1=y1,
-        y2=y2,
+        x1=x1 - xsize / 2,
+        x2=x2 + xsize / 2,
+        y1=y1 - ysize / 2,
+        y2=y2 + ysize / 2,
         cartesian_unit="degrees",
     )