Merge remote-tracking branch 'origin/xarray/main' into xarray/blend-utils

Author: mats-knmi

environment_dev.yml

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

pysteps/blending/linear_blending.py

@@ -21,27 +21,28 @@
 """
 
 import numpy as np
+import xarray as xr
 from pysteps import nowcasts
 from pysteps.utils import conversion
 from scipy.stats import rankdata
 
+from pysteps.xarray_helpers import convert_output_to_xarray_dataset
+
 
 def forecast(
-    precip,
-    precip_metadata,
-    velocity,
+    radar_dataset: xr.Dataset,
     timesteps,
     timestep,
     nowcast_method,
-    precip_nwp=None,
-    precip_nwp_metadata=None,
+    model_dataset: xr.Dataset = None,
     start_blending=120,
     end_blending=240,
     fill_nwp=True,
     saliency=False,
     nowcast_kwargs=None,
 ):
     """Generate a forecast by linearly or saliency-based blending of nowcasts with NWP data
+    # XR: Update docstring
 
     Parameters
     ----------
@@ -105,31 +106,27 @@ def forecast(
     if nowcast_kwargs is None:
         nowcast_kwargs = dict()
 
-    # Ensure that only the most recent precip timestep is used
-    if len(precip.shape) == 3:
-        precip = precip[-1, :, :]
-
     # First calculate the number of needed timesteps (up to end_blending) for the nowcast
     # to ensure that the nowcast calculation time is limited.
     timesteps_nowcast = int(end_blending / timestep)
 
     nowcast_method_func = nowcasts.get_method(nowcast_method)
 
     # Check if NWP data is given as input
-    if precip_nwp is not None:
+    if model_dataset is not None:
         # Calculate the nowcast
-        precip_nowcast = nowcast_method_func(
-            precip,
-            velocity,
-            timesteps_nowcast,
-            **nowcast_kwargs,
+        nowcast_dataset = nowcast_method_func(
+            radar_dataset, timesteps_nowcast, **nowcast_kwargs
         )
 
         # Make sure that precip_nowcast and precip_nwp are in mm/h
-        precip_nowcast, _ = conversion.to_rainrate(
-            precip_nowcast, metadata=precip_metadata
-        )
-        precip_nwp, _ = conversion.to_rainrate(precip_nwp, metadata=precip_nwp_metadata)
+        nowcast_dataset = conversion.to_rainrate(nowcast_dataset)
+        nowcast_precip_var = nowcast_dataset.attrs["precip_var"]
+        precip_nowcast = nowcast_dataset[nowcast_precip_var].values
+
+        model_dataset = conversion.to_rainrate(model_dataset)
+        model_precip_var = model_dataset.attrs["precip_var"]
+        precip_nwp = model_dataset[model_precip_var].values
 
         if len(precip_nowcast.shape) == 4:
             n_ens_members_nowcast = precip_nowcast.shape[0]
@@ -261,22 +258,19 @@ def forecast(
 
     else:
         # Calculate the nowcast
-        precip_nowcast = nowcast_method_func(
-            precip,
-            velocity,
-            timesteps,
-            **nowcast_kwargs,
+        nowcast_dataset = nowcast_method_func(
+            radar_dataset, timesteps, **nowcast_kwargs
         )
 
         # Make sure that precip_nowcast and precip_nwp are in mm/h
-        precip_nowcast, _ = conversion.to_rainrate(
-            precip_nowcast, metadata=precip_metadata
-        )
+        nowcast_dataset = conversion.to_rainrate(nowcast_dataset)
+        nowcast_precip_var = nowcast_dataset.attrs["precip_var"]
+        precip_nowcast = nowcast_dataset[nowcast_precip_var].values
 
         # If no NWP data is given, the blended field is simply equal to the nowcast field
         precip_blended = precip_nowcast
 
-    return precip_blended
+    return convert_output_to_xarray_dataset(radar_dataset, timesteps, precip_blended)
 
 
 def _get_slice(n_dims, ref_dim, ref_id):

pysteps/extrapolation/eulerian_persistence.py

@@ -0,0 +1,61 @@
+import numpy as np
+import xarray as xr
+from pysteps.xarray_helpers import convert_output_to_xarray_dataset
+
+
+def extrapolate(precip_dataset: xr.Dataset, timesteps, outval=np.nan, **kwargs):
+    """
+    A dummy extrapolation method to apply Eulerian persistence to a
+    two-dimensional precipitation field. The method returns the a sequence
+    of the same initial field with no extrapolation applied (i.e. Eulerian
+    persistence).
+
+    Parameters
+    ----------
+    precip_dataset : xarray.Dataset
+        xarray dataset containing the input precipitation field. All
+        values are required to be finite.
+    timesteps : int or list of floats
+        Number of time steps or a list of time steps.
+    outval : float, optional
+        Not used by the method.
+
+    Other Parameters
+    ----------------
+    return_displacement : bool
+        If True, return the total advection velocity (displacement) between the
+        initial input field and the advected one integrated along
+        the trajectory. Default : False
+
+    Returns
+    -------
+    out : array or tuple
+        If return_displacement=False, return a sequence of the same initial field
+        of shape (num_timesteps,m,n). Otherwise, return a tuple containing the
+        replicated fields and a (2,m,n) array of zeros.
+
+    References
+    ----------
+    :cite:`GZ2002`
+
+    """
+    del outval  # Unused by _eulerian_persistence
+    precip_var = precip_dataset.attrs["precip_var"]
+    precip = precip_dataset[precip_var].values[-1]
+
+    if isinstance(timesteps, int):
+        num_timesteps = timesteps
+    else:
+        num_timesteps = len(timesteps)
+
+    return_displacement = kwargs.get("return_displacement", False)
+
+    extrapolated_precip = np.repeat(precip[np.newaxis, :, :], num_timesteps, axis=0)
+    extrapolated_precip_dataset = convert_output_to_xarray_dataset(
+        precip_dataset, timesteps, extrapolated_precip
+    )
+
+    if not return_displacement:
+        return extrapolated_precip_dataset
+    else:
+        return extrapolated_precip_dataset, np.zeros((2,) + extrapolated_precip.shape)

pysteps/extrapolation/interface.py

@@ -34,63 +34,7 @@
 """
 
 import numpy as np
-
-from pysteps.extrapolation import semilagrangian
-
-
-def eulerian_persistence(precip, velocity, timesteps, outval=np.nan, **kwargs):
-    """
-    A dummy extrapolation method to apply Eulerian persistence to a
-    two-dimensional precipitation field. The method returns the a sequence
-    of the same initial field with no extrapolation applied (i.e. Eulerian
-    persistence).
-
-    Parameters
-    ----------
-    precip : array-like
-        Array of shape (m,n) containing the input precipitation field. All
-        values are required to be finite.
-    velocity : array-like
-        Not used by the method.
-    timesteps : int or list of floats
-        Number of time steps or a list of time steps.
-    outval : float, optional
-        Not used by the method.
-
-    Other Parameters
-    ----------------
-    return_displacement : bool
-        If True, return the total advection velocity (displacement) between the
-        initial input field and the advected one integrated along
-        the trajectory. Default : False
-
-    Returns
-    -------
-    out : array or tuple
-        If return_displacement=False, return a sequence of the same initial field
-        of shape (num_timesteps,m,n). Otherwise, return a tuple containing the
-        replicated fields and a (2,m,n) array of zeros.
-
-    References
-    ----------
-    :cite:`GZ2002`
-
-    """
-    del velocity, outval  # Unused by _eulerian_persistence
-
-    if isinstance(timesteps, int):
-        num_timesteps = timesteps
-    else:
-        num_timesteps = len(timesteps)
-
-    return_displacement = kwargs.get("return_displacement", False)
-
-    extrapolated_precip = np.repeat(precip[np.newaxis, :, :], num_timesteps, axis=0)
-
-    if not return_displacement:
-        return extrapolated_precip
-    else:
-        return extrapolated_precip, np.zeros((2,) + extrapolated_precip.shape)
+from pysteps.extrapolation import semilagrangian, eulerian_persistence
 
 
 def _do_nothing(precip, velocity, timesteps, outval=np.nan, **kwargs):
@@ -105,7 +49,7 @@ def _return_none(**kwargs):
 
 
 _extrapolation_methods = dict()
-_extrapolation_methods["eulerian"] = eulerian_persistence
+_extrapolation_methods["eulerian"] = eulerian_persistence.extrapolate
 _extrapolation_methods["semilagrangian"] = semilagrangian.extrapolate
 _extrapolation_methods[None] = _do_nothing
 _extrapolation_methods["none"] = _do_nothing

pysteps/io/importers.py

@@ -2021,7 +2021,6 @@ def _read_hdf5_cont(f, d):
                 }
 
         else:
-
             # Save h5py.Dataset by group name
             d[key] = np.array(value)
 

pysteps/nowcasts/interface.py

@@ -43,7 +43,7 @@
 
 _nowcast_methods = dict()
 _nowcast_methods["anvil"] = anvil.forecast
-_nowcast_methods["eulerian"] = eulerian_persistence
+_nowcast_methods["eulerian"] = eulerian_persistence.extrapolate
 _nowcast_methods["extrapolation"] = extrapolation.forecast
 _nowcast_methods["lagrangian"] = extrapolation.forecast
 _nowcast_methods["lagrangian_probability"] = lagrangian_probability.forecast

pysteps/nowcasts/steps.py

@@ -55,10 +55,10 @@ class StepsNowcasterConfig:
         Specifies the threshold value for minimum observable precipitation
         intensity. Required if mask_method is not None or conditional is True.
     norain_threshold: float
-      Specifies the threshold value for the fraction of rainy (see above) pixels
-      in the radar rainfall field below which we consider there to be no rain.
-      Depends on the amount of clutter typically present.
-      Standard set to 0.0
+        Specifies the threshold value for the fraction of rainy (see above) pixels
+        in the radar rainfall field below which we consider there to be no rain.
+        Depends on the amount of clutter typically present.
+        Standard set to 0.0
     kmperpixel: float, optional
         Spatial resolution of the input data (kilometers/pixel). Required if
         vel_pert_method is not None or mask_method is 'incremental'.

pysteps/pystepsrc

@@ -51,7 +51,7 @@
             "path_fmt": "%Y%m%d",
             "fn_pattern": "%Y%m%d%H%M_FINUTM",
             "fn_ext": "tif",
-            "importer": "geotiff",
+            "importer": "fmi_geotiff",
             "timestep": 5,
             "importer_kwargs": {}
         },

pysteps/tests/test_blending_linear_blending.py

@@ -1,10 +1,12 @@
 # -*- coding: utf-8 -*-
 
+from datetime import datetime
 import numpy as np
 import pytest
 from pysteps.blending.linear_blending import forecast, _get_ranked_salience, _get_ws
 from numpy.testing import assert_array_almost_equal
 from pysteps.utils import transformation
+from pysteps.xarray_helpers import convert_input_to_xarray_dataset
 
 # Test function arguments
 linear_arg_values = [
@@ -218,30 +220,68 @@ def test_linear_blending(
         for i in range(100, 200):
             r_input[:, i, :] = 11.0
     else:
-        r_input = np.zeros((200, 200))
+        r_input = np.zeros((1, 200, 200))
         for i in range(100, 200):
-            r_input[i, :] = 11.0
+            r_input[0, i, :] = 11.0
+
+    metadata = dict()
+    metadata["unit"] = "mm/h"
+    metadata["cartesian_unit"] = "km"
+    metadata["accutime"] = 5.0
+    metadata["zerovalue"] = 0.0
+    metadata["threshold"] = 0.01
+    metadata["zr_a"] = 200.0
+    metadata["zr_b"] = 1.6
+    metadata["x1"] = 0.0
+    metadata["x2"] = 200.0
+    metadata["y1"] = 0.0
+    metadata["y2"] = 200.0
+    metadata["yorigin"] = "lower"
+    metadata["institution"] = "test"
+    metadata["projection"] = (
+        "+proj=lcc +lon_0=4.55 +lat_1=50.8 +lat_2=50.8 +a=6371229 +es=0 +lat_0=50.8 +x_0=365950 +y_0=-365950.000000001"
+    )
+    radar_dataset = convert_input_to_xarray_dataset(
+        r_input,
+        None,
+        metadata,
+        datetime.fromisoformat("2021-07-04T11:50:00.000000000"),
+        300,
+    )
 
     # Transform from mm/h to dB
-    r_input, _ = transformation.dB_transform(
-        r_input, None, threshold=0.1, zerovalue=-15.0
+    radar_dataset = transformation.dB_transform(
+        radar_dataset, threshold=0.1, zerovalue=-15.0
     )
+    if V is not None:
+        radar_dataset["velocity_x"] = (["y", "x"], V[0])
+        radar_dataset["velocity_y"] = (["y", "x"], V[1])
+
+    if r_nwp is None:
+        model_dataset = None
+    else:
+        model_dataset = convert_input_to_xarray_dataset(
+            r_nwp,
+            None,
+            metadata,
+            datetime.fromisoformat("2021-07-04T11:50:00.000000000"),
+            300,
+        )
 
     # Calculate the blended field
-    r_blended = forecast(
-        r_input,
-        dict({"unit": "mm/h", "transform": "dB"}),
-        V,
+    blended_dataset = forecast(
+        radar_dataset,
         n_timesteps,
         timestep,
         nowcast_method,
-        r_nwp,
-        dict({"unit": "mm/h"}),
+        model_dataset,
         start_blending=start_blending,
         end_blending=end_blending,
         fill_nwp=fill_nwp,
         saliency=salient_blending,
     )
+    blended_precip_var = blended_dataset.attrs["precip_var"]
+    r_blended = blended_dataset[blended_precip_var].values
 
     # Assert that the blended field has the expected dimension
     if n_models > 1: