test chunked raise computation #250

import dask.array as da
import numpy as np
import xarray as xr
import xskillscore as xs


class CountingScheduler:
    """Simple dask scheduler counting the number of computes.
    Reference: https://stackoverflow.com/questions/53289286/"""

    def __init__(self, max_computes=0):
        self.total_computes = 0
        self.max_computes = max_computes

    def __call__(self, dsk, keys, **kwargs):
        self.total_computes += 1
        if self.total_computes > self.max_computes:
            raise RuntimeError(
                "Too many computes. Total: %d > max: %d."
                % (self.total_computes, self.max_computes)
            )
        return dask.get(dsk, keys, **kwargs)

def raise_if_dask_computes(max_computes=0):
    # return a dummy context manager so that this can be used for non-dask objects
    scheduler = CountingScheduler(max_computes)
    return dask.config.set(scheduler=scheduler)

value = dask.delayed(np.ones)(10000, 10000)
array = da.from_delayed(value, (100,100), dtype=float)

obs3 = xr.DataArray(
       array,
       dims=["lat", "lon"],
       name='var'
)
fct3 = xr.DataArray(
       array,
       dims=["lat", "lon"],
       name='var'
)
with raise_if_dask_computes():
    xs.rmse(obs3, fct3, skipna=True)

/mnt/c/Users/Solactus/GOOGLE~1/Bash/xskillscore/xskillscore/core/np_deterministic.py in _rmse(a, b, weights, axis, skipna)
    527     sumfunc, meanfunc = _get_numpy_funcs(skipna)
    528     if skipna:
--> 529         a, b, weights = _match_nans(a, b, weights)
    530     weights = _check_weights(weights)
    531 

/mnt/c/Users/Solactus/GOOGLE~1/Bash/xskillscore/xskillscore/core/np_deterministic.py in _match_nans(a, b, weights)
     34 
     35     """
---> 36     if np.isnan(a).any() or np.isnan(b).any():
     37         # Avoids mutating original arrays and bypasses read-only issue.
     38         a, b = a.copy(), b.copy()
...
<ipython-input-27-37dcfba2ae67> in __call__(self, dsk, keys, **kwargs)
     12             raise RuntimeError(
     13                 "Too many computes. Total: %d > max: %d."
---> 14                 % (self.total_computes, self.max_computes)
     15             )
     16         return dask.get(dsk, keys, **kwargs)

RuntimeError: Too many computes. Total: 1 > max: 0.