BaseAdjustment _check_inputs too restrictive for chunking along the main dimension

[saschahofmann]

Addressing a Problem?

I think for at least some adjustments like the Scaling class there is no need to prevent a chunked main dimension. Since the mean works really well over a chunked dimension, I think this restriction is more harmful (memory) than helpful. Unless there is an advantage of having the dimension unchunked that I am not seeing.

Potential Solution

A simple _allow_chunked_main_dimension = False similar to _allow_diff_calendars could do the trick quite easily. Happy to implement this if you're happy with this.

Additional context

Reproducible example

import xarray as xr
from xsdba import Scaling
import numpy as np

time = xr.date_range('2000-01-01', '2010-12-31')
hist = xr.DataArray(np.random.rand(time.size), coords=dict(time=('time', time))).chunk(time=100)
ref = xr.DataArray(np.random.rand(time.size), coords=dict(time=('time', time)))


Scaling.train(ref, hist)

Fails with on adjustment.py:93

ValueError: Multiple chunks along the main adjustment dimension time is not supported.

Contribution

• I would be willing/able to open a Pull Request to contribute this feature.

[coxipi]

Good idea, you can go ahead with this, thanks!

This is a restriction with quantile methods which don't apply to Scaling. I was unsure, maybe this was also needed because of Grouper / Map_groups, but it doesn't seem to be the case.

[saschahofmann]

I just monkeypatched the _check_inputs function and looks like you're right the map_blocks decorator does some check on its own that makes the operation fail.

Still looking whether this is easy to fix as well. It fails in base.py:624.

ValueError: The dimension(s) over which we group, reduce or interpolate cannot be chunked

[saschahofmann]

scaling_train is wrapped with map_groups which wraps map_blocks.

[saschahofmann]

This seems a little hard to solve without removing the map_groups wrapping. Or do you see a way to change this behaviour?

[coxipi]

Unfortunately no. There is an idea to remove map_groups / map_blocks altogether #147 , but it is still an idea.

Although, let's ask @aulemahal : In base.py::map_blocks (line 666), there is an error raised if the main dimension is chunked. Is that a true limitation related to Grouper/map_blocks, or just a safeguard because we expect Grouper should be used with quantile-based methods?

@saschahofmann you can try removing the Error raised in map_blocks and see if Scaling works. In this case, we could let the train/adjustment classes handle the errors on chunked dimensions on their own

[saschahofmann]

I commented out raising the error and it now the training on the lazy loaded data takes forever. So I think something isn't working correctly when I do that.

[aulemahal]

I sadly don't remember my full thought process for implementing that check in map_groups. The guess that it was in prevision for quantile-based method is probable. We could indeed relax this.

Testing with xclim's testdata, I see no real differences for Scaling when chunking along time or not. But my data is too small to encounter any memory issue.

While not being a good reason for raising an error, I'd want to note that when optimizing sdba for large datasets, we were trying to circumvent issues with having too many tasks in the tree, not only RAM issues. And having chunks along the grouped dimension usually increases the number of tasks more then having similar sized chunks along the non-grouped dims.

[saschahofmann]

How did you make Scaling run with chunked time? I disabled the check_inputs function and removed raising the error in map_blocks but as mentioned above suddenly it started loading the whole dataset even when it should have been lazily loaded.

[aulemahal]

I just took some time to think about this again and indeed, it is the map_blocks that blocks implementation of a Scaling that would support chunking along time. The method (like almost all others in xsdba) is coded so that a non-dask function takes care of most of the computation (scaling_train, scaling_adjust). The idea was to ease the process of implementing a method, while also supporting dask-backed data.

As discussed in #147, I'm not sure the goals were achieved. But anyway, we can't just remove the check. If data is sent to map_blocks with chunks along "time" but the wrapped function reduces "time", that doesn't make sense. To support this for Scaling we would need to reimplement it without map_blocks. Which means loosing the benefits it gives for very large dataset.

@coxipi , I would argue this is a "won't fix" in itself. It would be taken care of with #147.