Repository: Ouranosinc/xclim

Issue:
rle fails when first dimension is not 'time' on dask arrays
* xclim version: 0.17.0 & @master
* Python version: 3.6 & 3.8
* Operating System: boréas and Linux

### Description
<!--Describe what you were trying to get done.
Tell us what happened, what went wrong, and what you expected to happen.-->
When `xclim.indices.run_length.rle` is called with an array who's first dimension is not the specified `dim` (usually 'time'), it would fail if the array is stored in dask arrays. The reason for this is that we construct an array along `dim` and then broadcast is with `da`, but rechunk the new array with the chunks of `da` _instead of the chunks of broadcasted `da`_. As the broadcasting might change the orders of dimensions, it sometimes fail to align the requested chunks and the array's dimensions.

### What I Did
<!--Paste the command(s) you ran and the output.
If there was a crash, please include the traceback below.-->
```
import xclim as xc
import xarray as xr
from xclim import ensembles as xcens

ds = xr.tutorial.open_dataset('air_temperature').resample(time='D').mean()
ens = xcens.create_ensemble([ds, ds + 1, ds * 1.1])
ensc = ens.chunk({'time': 365, 'realization': 1})

xc.indices.run_length.rle(da=ensc.air > 275)
```

### What I Received
<!--Paste the output or the stack trace of the problem you experienced here.-->
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-8-c5b2c6ff097b> in <module>
----> 1 xc.indices.run_length.rle(da=ensc.air > 275)

... Shortened traceback ...

ValueError: Chunks do not add up to shape. Got chunks=((1, 1, 1), (365, 365), (25,), (53,)), shape=(730, 3, 25, 53)
```


Discussion:

Patch:
diff --git a/xclim/indices/run_length.py b/xclim/indices/run_length.py
index e1dc8f33..9b69522c 100644
--- a/xclim/indices/run_length.py
+++ b/xclim/indices/run_length.py
@@ -43,8 +43,11 @@ def get_npts(da: xr.DataArray) -> int:
 
 def rle(da: xr.DataArray, dim: str = "time", max_chunk: int = 1_000_000):
     n = len(da[dim])
-    i = xr.DataArray(np.arange(da[dim].size), dims=dim).chunk({"time": 1})
-    ind = xr.broadcast(i, da)[0].chunk(da.chunks)
+    # Need to chunk here to ensure the broadcasting is not made in memory
+    i = xr.DataArray(np.arange(da[dim].size), dims=dim).chunk({"time": -1})
+    ind, da = xr.broadcast(i, da)
+    # Rechunk, but with broadcasted da
+    ind = ind.chunk(da.chunks)
     b = ind.where(~da)  # find indexes where false
     end1 = (
         da.where(b[dim] == b[dim][-1], drop=True) * 0 + n


Test patch:
diff --git a/tests/test_run_length.py b/tests/test_run_length.py
index 90e2a0ae..03b853ed 100644
--- a/tests/test_run_length.py
+++ b/tests/test_run_length.py
@@ -26,6 +26,22 @@ class TestRLE:
         np.testing.assert_array_equal(l, [1, 10, 354])
         np.testing.assert_array_equal(p, [0, 1, 11])
 
+    @pytest.mark.parametrize("use_dask", [True, False])
+    def test_dataarray_nd(self, use_dask):
+        values = np.zeros((10, 365, 4, 4))
+        time = pd.date_range("2000-01-01", periods=365, freq="D")
+        values[:, 1:11, ...] = 1
+        da = xr.DataArray(values, coords={"time": time}, dims=("a", "time", "b", "c"))
+
+        if use_dask:
+            da = da.chunk({"a": 1, "b": 2})
+
+        out = rl.rle(da != 0).mean(["a", "b", "c"])
+        expected = np.zeros(366)
+        expected[1] = 10
+        expected[2:12] = np.nan
+        np.testing.assert_array_equal(out, expected)
+
 
 class TestLongestRun:
     nc_pr = os.path.join(TESTS_DATA, "NRCANdaily", "nrcan_canada_daily_pr_1990.nc")