Xarray/avijit dev

.gitignore

@@ -91,3 +91,11 @@ venv.bak/
 
 # Mac OS Stuff
 .DS_Store
+
+# Running local tests
+/tmp
+/pysteps/tests/tmp/
+
+/pysteps_data/
+/examples/data_download.py
+

ci/ci_test_env.yml

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

examples/advection_correction.py

@@ -1,29 +1,33 @@
 """
-Advection correction
-====================
+Advection correction (xarray version)
+=====================================
 
 This tutorial shows how to use the optical flow routines of pysteps to implement
 the advection correction procedure described in Anagnostou and Krajewski (1999).
 
 Advection correction is a temporal interpolation procedure that is often used
 when estimating rainfall accumulations to correct for the shift of rainfall patterns
-between consecutive radar rainfall maps. This shift becomes particularly 
+between consecutive radar rainfall maps. This shift becomes particularly
 significant for long radar scanning cycles and in presence of fast moving
 precipitation features.
 
 .. note:: The code for the advection correction using pysteps was originally
           written by `Daniel Wolfensberger <https://github.com/wolfidan>`_.
 
+Ported to xarray: the workflow now uses an xr.Dataset with a DataArray
+'precip_intensity' on dims ('time','y','x'), with projection info in attrs.
 """
 
 from datetime import datetime
-import matplotlib.pyplot as plt
 import numpy as np
+import matplotlib.pyplot as plt
+import xarray as xr
+
+from scipy.ndimage import map_coordinates
 
 from pysteps import io, motion, rcparams
 from pysteps.utils import conversion, dimension
 from pysteps.visualization import plot_precip_field
-from scipy.ndimage import map_coordinates
 
 ################################################################################
 # Read the radar input images
@@ -51,26 +55,42 @@
 fn_ext = data_source["fn_ext"]
 importer_name = data_source["importer"]
 importer_kwargs = data_source["importer_kwargs"]
-timestep = data_source["timestep"]
 
 # Find the input files from the archive
 fns = io.archive.find_by_date(
     date, root_path, path_fmt, fn_pattern, fn_ext, timestep=5, num_next_files=35
 )
 
-# Read the radar composites
+# Read as xarray Dataset (expects variable 'precip_intensity' and coords x,y,time)
 importer = io.get_method(importer_name, "importer")
-R, __, metadata = io.read_timeseries(fns, importer, **importer_kwargs)
+precip_ds = io.read_timeseries(fns, importer, **importer_kwargs)
 
 # Convert to mm/h
-R, metadata = conversion.to_rainrate(R, metadata)
+precip_ds = conversion.to_rainrate(precip_ds)
 
-# Upscale to 2 km (simply to reduce the memory demand)
-R, metadata = dimension.aggregate_fields_space(R, metadata, 2000)
+# Upscale to 2 km (to reduce memory)
+precip_ds = dimension.aggregate_fields_space(precip_ds, 2000)
 
 # Keep only one frame every 10 minutes (i.e., every 2 timesteps)
 # (to highlight the need for advection correction)
-R = R[::2]
+precip_ds = precip_ds.isel(time=slice(None, None, 2))
+
+# Ensure precip_var is present
+precip_ds.attrs.setdefault("precip_var", "precip_intensity")
+precip_var = precip_ds.attrs["precip_var"]
+
+# Convenience handle to the intensity DataArray (time, y, x)
+R = precip_ds[precip_var]
+
+# Build geodata for plotting from xarray metadata/coordinates
+geodata = {
+    "projection": precip_ds.attrs.get("projection", None),
+    "x1": float(R.x.values[0]),
+    "x2": float(R.x.values[-1]),
+    "y1": float(R.y.values[0]),
+    "y2": float(R.y.values[-1]),
+    "yorigin": "lower",
+}
 
 ################################################################################
 # Advection correction
@@ -84,43 +104,85 @@
 # going to use the Lucas-Kanade optical flow routine available in pysteps.
 
 
-def advection_correction(R, T=5, t=1):
+def advection_correction_xr(pair_da: xr.DataArray, T: int = 10, t: int = 1) -> xr.DataArray:
     """
-    R = np.array([qpe_previous, qpe_current])
-    T = time between two observations (5 min)
-    t = interpolation timestep (1 min)
+    Advection correction for a pair of successive frames using optical flow motion.
+
+    Parameters
+    ----------
+    pair_da : xr.DataArray
+        'precip_intensity' for two consecutive times with dims ('time','y','x').
+        pair_da.sizes['time'] must be 2.
+    T : int
+        Minutes between the two observations (here 10 after subsampling).
+    t : int
+        Interpolation timestep in minutes (1 minute).
+
+    Returns
+    -------
+    xr.DataArray
+        Time-averaged field from the temporally interpolated sequence between
+        the two frames (same dims as a single frame: ('y','x')).
     """
+
+
+    assert pair_da.sizes["time"] == 2, "pair_da must have exactly two time steps"
 
-    # Evaluate advection
-    oflow_method = motion.get_method("LK")
-    fd_kwargs = {"buffer_mask": 10}  # avoid edge effects
-    V = oflow_method(np.log(R), fd_kwargs=fd_kwargs)
+    # Optical flow on log-intensity (avoid log(0)) 
+    eps = 1e-6
+    oflow = motion.get_method("LK")  
+    fd_kwargs = {"buffer_mask": 10}  
 
-    # Perform temporal interpolation
-    Rd = np.zeros((R[0].shape))
-    x, y = np.meshgrid(
-        np.arange(R[0].shape[1], dtype=float), np.arange(R[0].shape[0], dtype=float)
-    )
+    pair_np = pair_da.values  # shape (2, y, x) — kept for interpolation step
+
+    # Build a two-frame Dataset that LK expects and set precip_var attr
+    pair_ds = pair_da.to_dataset(name=precip_var)
+    # carry over global attrs to satisfy decorators
+    pair_ds.attrs.update(precip_ds.attrs)
+    pair_ds.attrs.setdefault("precip_var", precip_var)
+
+    # log-transform 
+    pair_ds[precip_var] = np.log(pair_ds[precip_var] + eps)
+
+    # Run optical flow motion algorithm (returns velocities as variables on the same grid)
+    V_ds = oflow(pair_ds)
+    u = V_ds["velocity_x"].values  # x-component
+    v = V_ds["velocity_y"].values  # y-component
+
+    # Temporal interpolation loop
+    ny, nx = pair_np.shape[-2], pair_np.shape[-1]
+    y, x = np.meshgrid(np.arange(ny, dtype=float), np.arange(nx, dtype=float), indexing="ij")
+
+    Rd = np.zeros((ny, nx), dtype=float)
     for i in range(t, T + t, t):
-        pos1 = (y - i / T * V[1], x - i / T * V[0])
-        R1 = map_coordinates(R[0], pos1, order=1)
+        # Backward sample from the first frame
+        pos1 = (y - i / T * v, x - i / T * u)
+        R1 = map_coordinates(pair_np[0], pos1, order=1, mode="nearest")
 
-        pos2 = (y + (T - i) / T * V[1], x + (T - i) / T * V[0])
-        R2 = map_coordinates(R[1], pos2, order=1)
+        # Forward sample from the second frame
+        pos2 = (y + (T - i) / T * v, x + (T - i) / T * u)
+        R2 = map_coordinates(pair_np[1], pos2, order=1, mode="nearest")
 
         Rd += (T - i) * R1 + i * R2
 
-    return t / T**2 * Rd
+    Rd = (t / (T ** 2)) * Rd  # time-weighted mean over the interval
 
+    # Wrap back into a DataArray with original coords
+    return xr.DataArray(
+        Rd, dims=("y", "x"), coords={"y": pair_da.y, "x": pair_da.x}, attrs=pair_da.attrs
+    )
 
 ###############################################################################
 # Finally, we apply the advection correction to the whole sequence of radar
 # images and produce the rainfall accumulation map.
 
-R_ac = R[0].copy()
-for i in range(R.shape[0] - 1):
-    R_ac += advection_correction(R[i : (i + 2)], T=10, t=1)
-R_ac /= R.shape[0]
+R_mean = R.mean(dim="time")
+R_ac = R.isel(time=0).copy(deep=True)
+for k in range(R.sizes["time"] - 1):
+    pair = R.isel(time=slice(k, k + 2))
+    R_ac = R_ac + advection_correction_xr(pair, T=10, t=1)
+
+R_ac = R_ac / R.sizes["time"]
 
 ###############################################################################
 # Results
@@ -134,10 +196,10 @@ def advection_correction(R, T=5, t=1):
 # rainfall accumulation map.
 
 plt.figure(figsize=(9, 4))
-plt.subplot(121)
-plot_precip_field(R.mean(axis=0), title="3-h rainfall accumulation")
-plt.subplot(122)
-plot_precip_field(R_ac, title="Same with advection correction")
+plt.subplot(1, 2, 1)
+plot_precip_field(R_mean, geodata=geodata, title="3-h rainfall accumulation")
+plt.subplot(1, 2, 2)
+plot_precip_field(R_ac, geodata=geodata, title="Same with advection correction")
 plt.tight_layout()
 plt.show()
 
@@ -149,3 +211,5 @@ def advection_correction(R, T=5, t=1):
 # Estimation. Part I: Algorithm Formulation." Journal of Atmospheric and
 # Oceanic Technology 16: 189–97.
 # https://doi.org/10.1175/1520-0426(1999)016<0189:RTRREP>2.0.CO;2
+
+# sphinx_gallery_thumbnail_number = 2

examples/my_first_nowcast.ipynb

@@ -805,7 +805,7 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "pysteps_dev",
    "language": "python",
    "name": "python3"
   },
@@ -819,7 +819,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.2"
+   "version": "3.13.5"
   },
   "pycharm": {
    "stem_cell": {
@@ -833,4 +833,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}