Dask-friendly nan check in xr.corr() and xr.cov()

[AndrewILWilliams]

Was reading the discussion here and thought I'd draft a PR to implement some of the changes people suggested. It seems like the main problem is that currently in computation.py, if not valid_values.all() is not a lazy operation, and so can be a bottleneck for very large dataarrays. To get around this, I've lifted some neat tricks from #4559 so that valid_values remains a dask array.

---

• Closes Improve performance of xarray.corr() on big datasets #4804
• Tests added
• Passes pre-commit run --all-files
• User visible changes (including notable bug fixes) are documented in whats-new.rst

[max-sixty]

Hi @AndrewWilliams3142 , thanks for another PR!

This looks good. Could we add a test like https://github.com/pydata/xarray/pull/4559/files#diff-74d2dc289aa601b2de094fb3a3b687fd65963401b51b95cc5e0afcd06cc4cb82R45? And maybe reference this or #4559 as an explanation of what's going on.

[dcherian]

Thanks @AndrewWilliams3142 .

Unfortunately I don't think you can avoid this for dask arrays. Fundamentally bool(dask_array) will call compute. So this line if not _are_there_nans will call compute.

I think the only solution here is to special case this optimization for non-dask DataArray, and always call .where(valid_values) for dask arrays but maybe someone else has an idea.

[dcherian]

OK I guess you could stick the

valid_values = ...
if valid_values.any():
    ...

in a helper function and map_blocks that.

[AndrewILWilliams]

Hmm, ok. Silly question but does it matter that _are_there_nans is a scalar reduction of valid_values? So valid_values is still left uncomputed. Would this help the bottleneck at all?

[dcherian]

No because the program needs to know valid_blocks.any() to decide which branch of the if-else statement to execute. So it gets computed.

[AndrewILWilliams]

Hi @dcherian , just thinking about your suggestion for using map_blocks on the actual valid_values check. I've tested this and was wondering if you could maybe point to where I'm going wrong? It does mask out some of the values in a lazy way, but not the correct ones.

da_a = xr.DataArray(
        np.array([[1, 2, 3, 4], [1, 0.1, 0.2, 0.3], [2, 3.2, 0.6, 1.8]]),
        dims=("space", "time"),
        coords=[
            ("space", ["IA", "IL", "IN"]),
            ("time", pd.date_range("2000-01-01", freq="1D", periods=4)),
        ],
    ).chunk({'time':1})

da_b = xr.DataArray(
        np.array([[0.2, 0.4, 0.6, 2], [15, 10, 5, 1], [1, 3.2, np.nan, 1.8]]),
        dims=("space", "time"),
        coords=[
            ("space", ["IA", "IL", "IN"]),
            ("time", pd.date_range("2000-01-01", freq="1D", periods=4)),
        ],
    ).chunk({'time':1})

print(da_a)
>>> <xarray.DataArray (space: 3, time: 4)>
array([[1. , 2. , 3. , 4. ],
       [1. , 0.1, 0.2, 0.3],
       [2. , 3.2, 0.6, 1.8]])
Coordinates:
  * space    (space) <U2 'IA' 'IL' 'IN'
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04

print(da_b)
>>> <xarray.DataArray (space: 3, time: 4)>
array([[ 0.2,  0.4,  0.6,  2. ],
       [15. , 10. ,  5. ,  1. ],
       [ 1. ,  3.2,  nan,  1.8]])
Coordinates:
  * space    (space) <U2 'IA' 'IL' 'IN'
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04


# Define function to use in map_blocks
def _get_valid_values(da, other):
    da1, da2 = xr.align(da, other, join="inner", copy=False)
    
    # 2. Ignore the nans
    missing_vals = np.logical_or(da1.isnull(), da2.isnull())
    
    if missing_vals.any():
        da = da.where(missing_vals)
        return da
    else:
        return da


# test
outp = da_a.map_blocks(_get_valid_values, args=[da_b])


print(outp.compute())
>>> <xarray.DataArray (space: 3, time: 4)>
array([[1. , 2. , nan, 4. ],
       [1. , 0.1, nan, 0.3],
       [2. , 3.2, 0.6, 1.8]])
Coordinates:
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04
  * space    (space) object 'IA' 'IL' 'IN'

[keewis]

#4559 and #4668 fixed a similar issue using map_blocks, maybe you can use that as a reference?

[dcherian]

Well that was confusing! if missing_vals.any(): will not be triggered if all the values in a block are valid. With .chunk({"time": 1}), some of the blocks are all valid.

I discovered this with print debugging and the single-threaded scheduler which loops over blocks in a for-loop

# Define function to use in map_blocks
def _get_valid_values(da, other):
    da1, da2 = xr.align(da, other, join="inner", copy=False)
    
    # 2. Ignore the nans
    missing_vals = np.logical_or(da1.isnull(), da2.isnull())
    print(missing_vals)
    
    if missing_vals.any():
        da = da.where(missing_vals)
    return da

da_a.map_blocks(_get_valid_values, args=[da_b]).compute(scheduler="sync")

For better performance, we should try a dask.array.map_blocks approach with duck_array_ops.isnull

[AndrewILWilliams]

Thanks for that @dcherian ! I didn't know you could use print debugging on chunked operations like this!

One thing actually: If I change da = da.where(missing_vals) to da = da.where(~missing_vals) then we get the results we'd expect. Do you think this fixes the problem?

def _get_valid_values(da, other):
    da1, da2 = xr.align(da, other, join="outer", copy=False)
    
    # 2. Ignore the nans
    missing_vals = np.logical_or(da1.isnull(), da2.isnull())
    
    if missing_vals.any():
        da = da.where(~missing_vals)
        return da
    else:
        return da

print(da_a.map_blocks(_get_valid_values, args=[da_b]).compute())
<xarray.DataArray (space: 3, time: 4)>
array([[1. , 2. , 3. , 4. ],
       [1. , 0.1, 0.2, 0.3],
       [2. , 3.2, nan, 1.8]])
Coordinates:
  * time     (time) datetime64[ns] 2000-01-01 2000-01-02 2000-01-03 2000-01-04
  * space    (space) object 'IA' 'IL' 'IN'
  * 

[AndrewILWilliams]

Hey both, I've added a test to check that dask doesn't compute when calling either xr.corr() or xr.cov(), and also that the end result is still a dask array. Let me know if there's anything I've missed though! thanks for the help :)

@dcherian, regarding the apply_ufunc approach, I might leave that for now but as you said it can always be a future PR

[dcherian]

Thanks @AndrewWilliams3142 !

[max-sixty]

Please feel free to add a whatsnew @AndrewWilliams3142

[max-sixty]

Thank you very much @AndrewWilliams3142 !