import directly to xarray and rewrite utils_transformation to work with it

Author: mats-knmi

ci/ci_test_env.yml

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

pysteps/converters.py

@@ -0,0 +1,202 @@
+# -*- coding: utf-8 -*-
+"""
+pysteps.converters
+==================
+
+Module with data converter functions.
+
+.. autosummary::
+    :toctree: ../generated/
+
+    convert_to_xarray_dataset
+"""
+
+import numpy as np
+import pyproj
+import xarray as xr
+
+from pysteps.utils.conversion import cf_parameters_from_unit
+
+# TODO(converters): Write methods for converting Proj.4 projection definitions
+# into CF grid mapping attributes. Currently this has been implemented for
+# the stereographic projection.
+# The conversions implemented here are take from:
+# https://github.com/cf-convention/cf-convention.github.io/blob/master/wkt-proj-4.md
+
+
+def _convert_proj4_to_grid_mapping(proj4str):
+    tokens = proj4str.split("+")
+
+    d = {}
+    for t in tokens[1:]:
+        t = t.split("=")
+        if len(t) > 1:
+            d[t[0]] = t[1].strip()
+
+    params = {}
+    # TODO(exporters): implement more projection types here
+    if d["proj"] == "stere":
+        grid_mapping_var_name = "polar_stereographic"
+        grid_mapping_name = "polar_stereographic"
+        v = d["lon_0"] if d["lon_0"][-1] not in ["E", "W"] else d["lon_0"][:-1]
+        params["straight_vertical_longitude_from_pole"] = float(v)
+        v = d["lat_0"] if d["lat_0"][-1] not in ["N", "S"] else d["lat_0"][:-1]
+        params["latitude_of_projection_origin"] = float(v)
+        if "lat_ts" in list(d.keys()):
+            params["standard_parallel"] = float(d["lat_ts"])
+        elif "k_0" in list(d.keys()):
+            params["scale_factor_at_projection_origin"] = float(d["k_0"])
+        params["false_easting"] = float(d["x_0"])
+        params["false_northing"] = float(d["y_0"])
+    elif d["proj"] == "aea":  # Albers Conical Equal Area
+        grid_mapping_var_name = "proj"
+        grid_mapping_name = "albers_conical_equal_area"
+        params["false_easting"] = float(d["x_0"]) if "x_0" in d else float(0)
+        params["false_northing"] = float(d["y_0"]) if "y_0" in d else float(0)
+        v = d["lon_0"] if "lon_0" in d else float(0)
+        params["longitude_of_central_meridian"] = float(v)
+        v = d["lat_0"] if "lat_0" in d else float(0)
+        params["latitude_of_projection_origin"] = float(v)
+        v1 = d["lat_1"] if "lat_1" in d else float(0)
+        v2 = d["lat_2"] if "lat_2" in d else float(0)
+        params["standard_parallel"] = (float(v1), float(v2))
+    else:
+        print("unknown projection", d["proj"])
+        return None, None, None
+
+    return grid_mapping_var_name, grid_mapping_name, params
+
+
+def convert_to_xarray_dataset(
+    precip: np.ndarray,
+    quality: np.ndarray | None,
+    metadata: dict[str, str | float | None],
+) -> xr.Dataset:
+    """
+    Read a precip, quality, metadata tuple as returned by the importers
+    (:py:mod:`pysteps.io.importers`) and return an xarray dataset containing
+    this data.
+
+    Parameters
+    ----------
+    precip: array
+        2D array containing imported precipitation data.
+    quality: array, None
+        2D array containing the quality values of the imported precipitation
+        data, can be None.
+    metadata: dict
+        Metadata dictionary containing the attributes described in the
+        documentation of :py:mod:`pysteps.io.importers`.
+
+    Returns
+    -------
+    out: Dataset
+        A CF compliant xarray dataset, which contains all data and metadata.
+
+    """
+    var_name, attrs = cf_parameters_from_unit(metadata["unit"])
+    h, w = precip.shape
+    x_r = np.linspace(metadata["x1"], metadata["x2"], w + 1)[:-1]
+    x_r += 0.5 * (x_r[1] - x_r[0])
+    y_r = np.linspace(metadata["y1"], metadata["y2"], h + 1)[:-1]
+    y_r += 0.5 * (y_r[1] - y_r[0])
+
+    # flip yr vector if yorigin is upper
+    if metadata["yorigin"] == "upper":
+        y_r = np.flip(y_r)
+
+    x_2d, y_2d = np.meshgrid(x_r, y_r)
+    pr = pyproj.Proj(metadata["projection"])
+    lon, lat = pr(x_2d.flatten(), y_2d.flatten(), inverse=True)
+
+    (
+        grid_mapping_var_name,
+        grid_mapping_name,
+        grid_mapping_params,
+    ) = _convert_proj4_to_grid_mapping(metadata["projection"])
+
+    data_vars = {
+        var_name: (
+            ["y", "x"],
+            precip,
+            {
+                "units": attrs["units"],
+                "standard_name": attrs["standard_name"],
+                "long_name": attrs["long_name"],
+                "grid_mapping": "projection",
+                "transform": metadata["transform"],
+                "accutime": metadata["accutime"],
+                "threshold": metadata["threshold"],
+                "zerovalue": metadata["zerovalue"],
+                "zr_a": metadata["zr_a"],
+                "zr_b": metadata["zr_b"],
+            },
+        )
+    }
+    if quality is not None:
+        data_vars["quality"] = (
+            ["y", "x"],
+            quality,
+            {
+                "units": "1",
+                "standard_name": "quality_flag",
+                "grid_mapping": "projection",
+            },
+        )
+    coords = {
+        "y": (
+            ["y"],
+            y_r,
+            {
+                "axis": "Y",
+                "long_name": "y-coordinate in Cartesian system",
+                "standard_name": "projection_y_coordinate",
+                "units": metadata["cartesian_unit"],
+            },
+        ),
+        "x": (
+            ["x"],
+            x_r,
+            {
+                "axis": "X",
+                "long_name": "x-coordinate in Cartesian system",
+                "standard_name": "projection_x_coordinate",
+                "units": metadata["cartesian_unit"],
+            },
+        ),
+        "lon": (
+            ["y", "x"],
+            lon.reshape(precip.shape),
+            {
+                "long_name": "longitude coordinate",
+                "standard_name": "longitude",
+                # TODO(converters): Don't hard-code the unit.
+                "units": "degrees_east",
+            },
+        ),
+        "lat": (
+            ["y", "x"],
+            lat.reshape(precip.shape),
+            {
+                "long_name": "latitude coordinate",
+                "standard_name": "latitude",
+                # TODO(converters): Don't hard-code the unit.
+                "units": "degrees_north",
+            },
+        ),
+    }
+    if grid_mapping_var_name is not None:
+        coords[grid_mapping_name] = (
+            (
+                [],
+                None,
+                {"grid_mapping_name": grid_mapping_name, **grid_mapping_params},
+            ),
+        )
+    attrs = {
+        "Conventions": "CF-1.7",
+        "institution": metadata["institution"],
+        "projection": metadata["projection"],
+        "precip_var": var_name,
+    }
+    return xr.Dataset(data_vars=data_vars, coords=coords, attrs=attrs)

pysteps/decorators.py

@@ -22,6 +22,8 @@
 
 import numpy as np
 
+from pysteps.converters import convert_to_xarray_dataset
+
 
 def _add_extra_kwrds_to_docstrings(target_func, extra_kwargs_doc_text):
     """
@@ -66,7 +68,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 +90,7 @@ def _import_with_postprocessing(*args, **kwargs):
                     mask = ~np.isfinite(precip)
                     precip[mask] = _fillna
 
-            return (precip.astype(_dtype),) + tuple(other_args)
+            return convert_to_xarray_dataset(precip.astype(_dtype), quality, metadata)
 
         extra_kwargs_doc = """
             Other Parameters

pysteps/io/importers.py

@@ -89,12 +89,10 @@
 from functools import partial
 
 import numpy as np
-
 from matplotlib.pyplot import imread
 
 from pysteps.decorators import postprocess_import
-from pysteps.exceptions import DataModelError
-from pysteps.exceptions import MissingOptionalDependency
+from pysteps.exceptions import DataModelError, MissingOptionalDependency
 from pysteps.utils import aggregate_fields
 
 try:

pysteps/io/readers.py

@@ -12,13 +12,14 @@
 """
 
 import numpy as np
+import xarray as xr
 
 
-def read_timeseries(inputfns, importer, **kwargs):
+def read_timeseries(inputfns, importer, **kwargs) -> xr.Dataset | None:
     """
     Read a time series of input files using the methods implemented in the
-    :py:mod:`pysteps.io.importers` module and stack them into a 3d array of
-    shape (num_timesteps, height, width).
+    :py:mod:`pysteps.io.importers` module and stack them into a 3d xarray
+    dataset of shape (num_timesteps, height, width).
 
     Parameters
     ----------
@@ -32,50 +33,50 @@ def read_timeseries(inputfns, importer, **kwargs):
 
     Returns
     -------
-    out: tuple
-        A three-element tuple containing the read data and quality rasters and
+    out: Dataset
+        A dataset containing the read data and quality rasters and
         associated metadata. If an input file name is None, the corresponding
         precipitation and quality fields are filled with nan values. If all
         input file names are None or if the length of the file name list is
-        zero, a three-element tuple containing None values is returned.
+        zero, None is returned.
 
     """
 
     # check for missing data
-    precip_ref = None
+    dataset_ref = None
     if all(ifn is None for ifn in inputfns):
-        return None, None, None
+        return None
     else:
         if len(inputfns[0]) == 0:
-            return None, None, None
+            return None
         for ifn in inputfns[0]:
             if ifn is not None:
-                precip_ref, quality_ref, metadata = importer(ifn, **kwargs)
+                dataset_ref = importer(ifn, **kwargs)
                 break
 
-    if precip_ref is None:
-        return None, None, None
+    if dataset_ref is None:
+        return None
 
-    precip = []
-    quality = []
-    timestamps = []
+    startdate = min(inputfns[1])
+
+    datasets = []
     for i, ifn in enumerate(inputfns[0]):
         if ifn is not None:
-            precip_, quality_, _ = importer(ifn, **kwargs)
-            precip.append(precip_)
-            quality.append(quality_)
-            timestamps.append(inputfns[1][i])
+            dataset_ = importer(ifn, **kwargs)
         else:
-            precip.append(precip_ref * np.nan)
-            if quality_ref is not None:
-                quality.append(quality_ref * np.nan)
-            else:
-                quality.append(None)
-            timestamps.append(inputfns[1][i])
-
-    # Replace this with stack?
-    precip = np.concatenate([precip_[None, :, :] for precip_ in precip])
-    # TODO: Q should be organized as R, but this is not trivial as Q_ can be also None or a scalar
-    metadata["timestamps"] = np.array(timestamps)
+            dataset_ = dataset_ref * np.nan
+        dataset_ = dataset_.expand_dims(dim="time", axis=0)
+        dataset_ = dataset_.assign_coords(
+            time=(
+                "time",
+                [inputfns[1][i]],
+                {
+                    "long_name": "forecast time",
+                    "units": f"seconds since {startdate:%Y-%m-%d %H:%M:%S}",
+                },
+            )
+        )
+        datasets.append(dataset_)
 
-    return precip, quality, metadata
+    dataset = xr.concat(datasets, dim="time")
+    return dataset

pysteps/nowcasts/anvil.py

@@ -19,12 +19,13 @@
 """
 
 import time
+
 import numpy as np
 from scipy.ndimage import gaussian_filter
-from pysteps import cascade, extrapolation
+
+from pysteps import cascade, extrapolation, utils
 from pysteps.nowcasts.utils import nowcast_main_loop
 from pysteps.timeseries import autoregression
-from pysteps import utils
 
 try:
     import dask