Merge pull request #495 from Ouranosinc/fix-ND-rle-chunks

Fix alignment bug when RLE is called with first dim not time

tests/test_run_length.py

@@ -26,6 +26,22 @@ def test_dataarray(self):
         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")

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