Lazy indexing in functions using output of xr.argmax and others #1276
Description
Setup Information
- Xclim version: 0.39.13-beta
- Python version: 3.10.8
- Operating System: Ubuntu 21.04
Description
I'm working on a project with xscen, I want to compute xci.snw_max_doy. Starting from processing_level==extracted and beyond, I am not able to compute this function. I prepared a minimal example that has nothing to do with xscen that isolates the problem. Lazy computing an index used to select the subset of an array seems to be the problem. The workaround described is to compute/load arrays before using them in snw_max_doy, and it has been confirmed to work in the more complex case with xscen
Steps To Reproduce
Using xci.snw_max_doy on a snw array that is chunked only works if chunking is performed before using xci.snw_max_doy
success = xci.snw_max_doy(snw.chunk({"rlon":4}).load())
fail = xci.snw_max_doy(snw.chunk({"rlon":4}))The failing error is:
TypeError: unexpected indexer type for VectorizedIndexer: dask.array<nanarg_agg-aggregate, shape=(10, 10), dtype=int64, chunksize=(10, 4), chunktype=numpy.ndarray>
Additional context
This is issue comes from generic.select_resample_op(snw, op=generic.doymax, freq=freq) and
i = da.argmax(dim="time")
out = da.time.dt.dayofyear.isel(time=i, drop=True)If the chunking is hanging, it produces an index i which remains to be computed, which does not work in the last line. This is apparently linked to an open issue in xarray: pydata/xarray#2511
According to @aulemahal, this problem is avoided in run_length with lazy_indexing, so it might be useful to expand the scope of this function.
Contribution
- I would be willing/able to open a Pull Request to address this bug.
Code of Conduct
- I agree to follow this project's Code of Conduct