Repository: Ouranosinc/xclim

Issue:
Make run_length methods return the coordinate values
Right now, functions like `rl.first_run` return indexes (integers) along the given dimension. Retrieving the coordinate values (like the datetime objects) is not an easy task and usually gets done by unelegant hacky lines. There should be an option on the run_length methods to do this.

Moreover, it should be lazy (which might not be simple) and we could add options to return sub-indexers (like "time.dt.dayofyear"). And, there could be a way to manage NaNs in the input : right now functions like `first_run` return NaN if there are any.

And I think there is a potential need for a "last_run" function, which is also currently done by a hacky workaround.

Ça pourrait intéresser @sbiner et @Zeitsperre !

Discussion:

Patch:
diff --git a/HISTORY.rst b/HISTORY.rst
index ba05dad9..9915f137 100644
--- a/HISTORY.rst
+++ b/HISTORY.rst
@@ -4,6 +4,7 @@ History
 
 0.15.x
 ------
+* Add feature to retrieve coordinate values instead of index in `run_length.first_run`. Add `run_length.last_run`.
 * Fix bug in subset_gridpoint to work on lat/lon coords of any dimension when they are not a dimension of the data.
 
 0.14.x (2020-02-21)
diff --git a/setup.cfg b/setup.cfg
index 7d41e226..c0fac8ef 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.14.0
+current_version = 0.14.1-beta
 commit = False
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+).(?P<patch>\d+)(\-(?P<release>[a-z]+))?
diff --git a/setup.py b/setup.py
index e226e131..7310c16f 100644
--- a/setup.py
+++ b/setup.py
@@ -12,7 +12,7 @@ URL = "https://github.com/Ouranosinc/xclim"
 AUTHOR = "Travis Logan"
 AUTHOR_EMAIL = "logan.travis@ouranos.ca"
 REQUIRES_PYTHON = ">=3.6.0"
-VERSION = "0.14.0"
+VERSION = "0.14.1-beta"
 LICENSE = "Apache Software License 2.0"
 
 with open("README.rst") as readme_file:
diff --git a/xclim/__init__.py b/xclim/__init__.py
index 2f4f2ca5..66878af8 100644
--- a/xclim/__init__.py
+++ b/xclim/__init__.py
@@ -12,7 +12,7 @@ from xclim import seaIce
 
 __author__ = """Travis Logan"""
 __email__ = "logan.travis@ouranos.ca"
-__version__ = "0.14.0"
+__version__ = "0.14.1-beta"
 
 
 def build_module(name, objs, doc="", source=None, mode="ignore"):
diff --git a/xclim/indices/_threshold.py b/xclim/indices/_threshold.py
index 8a289dbe..e1e7b740 100644
--- a/xclim/indices/_threshold.py
+++ b/xclim/indices/_threshold.py
@@ -264,8 +264,9 @@ def freshet_start(
     """
     thresh = utils.convert_units_to(thresh, tas)
     over = tas > thresh
-    group = over.resample(time=freq)
-    return group.apply(rl.first_run_ufunc, window=window, index="dayofyear")
+    return over.resample(time=freq).map(
+        rl.first_run, dim="time", window=window, coord="dayofyear"
+    )
 
 
 @declare_units("C days", tas="[temperature]", thresh="[temperature]")
diff --git a/xclim/run_length.py b/xclim/run_length.py
index 83e09618..a4fe16c6 100644
--- a/xclim/run_length.py
+++ b/xclim/run_length.py
@@ -1,12 +1,14 @@
 # -*- coding: utf-8 -*-
 """Run length algorithms module"""
 import logging
+from functools import partial
 from typing import Optional
 from typing import Sequence
 from typing import Tuple
 from typing import Union
 from warnings import warn
 
+import dask.array as dsk
 import numpy as np
 import xarray as xr
 
@@ -186,47 +188,103 @@ def first_run(
     da: xr.DataArray,
     window: int,
     dim: str = "time",
+    coord: Optional[Union[str, bool]] = False,
     ufunc_1dim: Union[str, bool] = "auto",
 ):
-    """Return the index of the first item of a run of at least a given length.
+    """Return the index of the first item of the first run of at least a given length.
 
-        Parameters
-        ----------
-        da : xr.DataArray
-          Input N-dimensional DataArray (boolean)
-        window : int
-          Minimum duration of consecutive run to accumulate values.
-        dim : str
-          Dimension along which to calculate consecutive run (default: 'time').
-        ufunc_1dim : Union[str, bool]
-          Use the 1d 'ufunc' version of this function : default (auto) will attempt to select optimal
-          usage based on number of data points.  Using 1D_ufunc=True is typically more efficient
-          for dataarray with a small number of gridpoints.
+    Parameters
+    ----------
+    da : xr.DataArray
+      Input N-dimensional DataArray (boolean)
+    window : int
+      Minimum duration of consecutive run to accumulate values.
+    dim : str
+      Dimension along which to calculate consecutive run (default: 'time').
+    coord : Optional[str]
+      If not False, the function returns values along `dim` instead of indexes.
+      If `dim` has a datetime dtype, `coord` can also be a str of the name of the
+      DateTimeAccessor object to use (ex: 'dayofyear').
+    ufunc_1dim : Union[str, bool]
+      Use the 1d 'ufunc' version of this function : default (auto) will attempt to select optimal
+      usage based on number of data points.  Using 1D_ufunc=True is typically more efficient
+      for dataarray with a small number of gridpoints.
 
-        Returns
-        -------
-        out : N-dimensional xarray data array (int)
-          Index of first item in first valid run. Returns np.nan if there are no valid run.
-        """
+    Returns
+    -------
+    out : xr.DataArray
+      Index (or coordinate if `coord` is not False) of first item in first valid run. Returns np.nan if there are no valid run.
+    """
     if ufunc_1dim == "auto":
-        npts = get_npts(da)
-        ufunc_1dim = npts <= npts_opt
+        if isinstance(da.data, dsk.Array) and len(da.chunks[da.dims.index(dim)]) > 1:
+            ufunc_1dim = False
+        else:
+            npts = get_npts(da)
+            ufunc_1dim = npts <= npts_opt
+
+    da = da.fillna(0)  # We expect a boolean array, but there could be NaNs nonetheless
 
     if ufunc_1dim:
-        out = first_run_ufunc(da, window)
+        out = first_run_ufunc(x=da, window=window, dim=dim)
 
     else:
-        dims = list(da.dims)
-        if "time" not in dims:
-            da["time"] = da[dim]
-            da.swap_dims({dim: "time"})
         da = da.astype("int")
-        i = xr.DataArray(np.arange(da[dim].size), dims=dim).chunk({"time": 1})
-        ind = xr.broadcast(i, da)[0].chunk(da.chunks)
+        i = xr.DataArray(np.arange(da[dim].size), dims=dim)
+        ind = xr.broadcast(i, da)[0].transpose(*da.dims)
+        if isinstance(da.data, dsk.Array):
+            ind = ind.chunk(da.chunks)
         wind_sum = da.rolling(time=window).sum(allow_lazy=True, skipna=False)
         out = ind.where(wind_sum >= window).min(dim=dim) - (
             window - 1
         )  # remove window -1 as rolling result index is last element of the moving window
+
+    if coord:
+        crd = da[dim]
+        if isinstance(coord, str):
+            crd = getattr(crd.dt, coord)
+
+        return lazy_indexing(crd, out)
+
+    return out
+
+
+def last_run(
+    da: xr.DataArray,
+    window: int,
+    dim: str = "time",
+    coord: Optional[Union[str, bool]] = False,
+    ufunc_1dim: Union[str, bool] = "auto",
+):
+    """Return the index of the last item of the last run of at least a given length.
+
+    Parameters
+    ----------
+    da : xr.DataArray
+      Input N-dimensional DataArray (boolean)
+    window : int
+      Minimum duration of consecutive run to accumulate values.
+    dim : str
+      Dimension along which to calculate consecutive run (default: 'time').
+    coord : Optional[str]
+      If not False, the function returns values along `dim` instead of indexes.
+      If `dim` has a datetime dtype, `coord` can also be a str of the name of the
+      DateTimeAccessor object to use (ex: 'dayofyear').
+    ufunc_1dim : Union[str, bool]
+      Use the 1d 'ufunc' version of this function : default (auto) will attempt to select optimal
+      usage based on number of data points.  Using 1D_ufunc=True is typically more efficient
+      for dataarray with a small number of gridpoints.
+
+    Returns
+    -------
+    out : xr.DataArray
+      Index (or coordinate if `coord` is not False) of last item in last valid run. Returns np.nan if there are no valid run.
+    """
+    reversed_da = da.sortby(dim, ascending=False)
+    out = first_run(
+        reversed_da, window=window, dim=dim, coord=coord, ufunc_1dim=ufunc_1dim
+    )
+    if not coord:
+        return reversed_da[dim].size - out - 1
     return out
 
 
@@ -431,9 +489,7 @@ def longest_run_ufunc(x: Sequence[bool]) -> xr.apply_ufunc:
     )
 
 
-def first_run_ufunc(
-    x: xr.DataArray, window, index: Optional[str] = None
-) -> xr.apply_ufunc:
+def first_run_ufunc(x: xr.DataArray, window: int, dim: str = "time",) -> xr.apply_ufunc:
     """Dask-parallel version of first_run_1d, ie the first entry in array of consecutive true values.
 
     Parameters
@@ -452,7 +508,7 @@ def first_run_ufunc(
     ind = xr.apply_ufunc(
         first_run_1d,
         x,
-        input_core_dims=[["time"]],
+        input_core_dims=[[dim]],
         vectorize=True,
         dask="parallelized",
         output_dtypes=[np.float],
@@ -460,9 +516,36 @@ def first_run_ufunc(
         kwargs={"window": window},
     )
 
-    if index is not None and ind.notnull().all():
-        val = getattr(x.indexes["time"], index)
-        i = ind.data.astype(int)
-        ind.data = val[i]
-
     return ind
+
+
+def lazy_indexing(da: xr.DataArray, index: xr.DataArray):
+    """Get values of `da` at indices `index` in a NaN-aware and lazy manner.
+
+    Parameters
+    ----------
+    da : xr.DataArray
+      1D Input array
+    index : xr.DataArray
+      N-d integer indices
+
+    Returns
+    -------
+    xr.DataArray
+      Values of `da` at indices `index`
+    """
+
+    def _index_from_1d_array(array, indices):
+        return array[
+            indices,
+        ]
+
+    invalid = index.isnull()
+    index = index.fillna(0).astype(int)
+    func = partial(_index_from_1d_array, da)
+
+    if isinstance(index, dsk.Array):
+        out = dsk.map_blocks(func, index, dtype=index.dtype)
+    else:
+        out = func(index)
+    return out.where(~invalid)


Test patch:
diff --git a/tests/test_run_length.py b/tests/test_run_length.py
index 01b3d0c6..49963643 100644
--- a/tests/test_run_length.py
+++ b/tests/test_run_length.py
@@ -2,6 +2,7 @@ import os
 
 import numpy as np
 import pandas as pd
+import pytest
 import xarray as xr
 
 from xclim import run_length as rl
@@ -139,20 +140,59 @@ class TestLongestRun:
 class TestFirstRun:
     nc_pr = os.path.join(TESTS_DATA, "NRCANdaily", "nrcan_canada_daily_pr_1990.nc")
 
-    def test_simple(self):
+    def test_run_1d(self):
         a = np.zeros(100, bool)
         a[10:20] = 1
         i = rl.first_run_1d(a, 5)
         assert 10 == i
 
+    def test_real_data(self):
         # n-dim version versus ufunc
         da3d = xr.open_dataset(self.nc_pr).pr[:, 40:50, 50:68] != 0
-        lt_orig = da3d.resample(time="M").apply(rl.first_run_ufunc, window=5)
-        lt_Ndim = da3d.resample(time="M").apply(
+        lt_orig = da3d.resample(time="M").map(rl.first_run_ufunc, window=5)
+        lt_Ndim = da3d.resample(time="M").map(
             rl.first_run, window=5, dim="time", ufunc_1dim=False
         )
         np.testing.assert_array_equal(lt_orig, lt_Ndim)
 
+    @pytest.mark.parametrize(
+        "coord,expected",
+        [(False, 30), (True, np.datetime64("2000-01-31")), ("dayofyear", 31)],
+    )
+    @pytest.mark.parametrize(
+        "use_dask,use_1dim", [(True, False), (False, False), (False, True)]
+    )
+    def test_simple(self, tas_series, coord, expected, use_dask, use_1dim):
+        t = np.zeros(60)
+        t[30:40] = 2
+        tas = tas_series(t, start="2000-01-01")
+        runs = xr.concat((tas, tas), dim="dim0")
+
+        if use_dask:
+            runs = runs.chunk({"time": 10, "dim0": 1})
+
+        out = rl.first_run(runs, window=1, dim="time", coord=coord, ufunc_1dim=use_1dim)
+        np.testing.assert_array_equal(out.load(), expected)
+
+    @pytest.mark.parametrize(
+        "coord,expected",
+        [(False, 39), (True, np.datetime64("2000-02-09")), ("dayofyear", 40)],
+    )
+    @pytest.mark.parametrize(
+        "use_dask,use_1dim", [(True, False), (False, False), (False, True)]
+    )
+    def test_last_run(self, tas_series, coord, expected, use_dask, use_1dim):
+        t = np.zeros(60)
+        t[30:40] = 2
+        tas = tas_series(t, start="2000-01-01")
+        runs = xr.concat((tas, tas), dim="dim0")
+
+        if use_dask:
+            runs = runs.chunk({"time": 10, "dim0": 1})
+
+        out = rl.last_run(runs, window=1, dim="time", coord=coord, ufunc_1dim=use_1dim)
+        np.testing.assert_array_equal(out.load(), expected)
+
 
 class TestWindowedRunEvents:
     nc_pr = os.path.join(TESTS_DATA, "NRCANdaily", "nrcan_canada_daily_pr_1990.nc")