issue_comments

@willirath this is cool, but I think it doesn't explain why the tests fail. Currently da_a.mean() and the da_b.mean() calls do know about each other's missing data! That's what we're doing in these lines.

@dcherian, I think I've got it to work, but you need to account for the length(s) of the dimension you're calculating the correlation over. (i.e. (da-da.mean('time')).sum('time') is not the same as da.sum('time') - da.mean('time') because you should actually do da.sum('time') - da.mean('time')*length_of_time_dim)

This latest commit does this, but I'm not sure whether the added complication is worth it yet? Thoughts welcome. ```python3 def _mean(da): return (da.sum(dim=dim, skipna=True, min_count=1) / (valid_count))

dim_length = da_a.notnull().sum(dim=dim, skipna=True) def _mean_detrended_term(da): return (dim_length * da / (valid_count))

cov = _mean(da_a * da_b) - _mean_detrended_term(da_a.mean(dim=dim) * da_b.mean(dim=dim)) ```

I think the problem with

cov = _mean(da_a * da_b) - da_a.mean(dim=dim) * da_b.mean(dim=dim)

is that the da_a.mean() and the da_b.mean() calls don't know about each other's missing data.

Shouldn't the following do? python cov = ( (da_a * da_b).mean(dim) - ( da_a.where(da_b.notnull()).mean(dim) * da_b.where(da_a.notnull()).mean(dim) ) ) (See here: https://nbviewer.jupyter.org/gist/willirath/cfaa8fb1b53fcb8dcb05ddde839c794c )

is it just

cov = _mean(da_a * da_b) - da_a.mean(dim=dim) * da_b.mean(dim=dim)

I think you'd still have to normalize the second term by 1 / (valid_count). However, I just tried both of these approaches and neither pass the test suite, so we may need to do more thinking...

@willirath this is great stuff, thanks again! So generally it looks like the graph is more efficient when doing operations of the form:

python3 (X * Y).mean('time') - (X.mean('time') * Y.mean('time'))

than doing python3 ((X - X.mean('time')) * (Y-Y.mean('time'))).mean('time')

or like what I've implemented (see screenshot)? ```python3 intermediate = (X * Y) - (X.mean('time') * Y.mean('time'))

intermediate.mean('time') ```

If so, it seems like the most efficient(?) way to do the computation in _cov_corr() is to combine it all into one line? I can't think of how to do this though...

@AndrewWilliams3142 @dcherian Looks like I broke the first Gist. :(

Your Example above does not quite get there, because the xr.DataArray(np...).chunk() just leads to one chunk per data array.

Here's a Gist that explains the idea for the correlations: https://nbviewer.jupyter.org/gist/willirath/c5c5274f31c98e8452548e8571158803

With ```python X = xr.DataArray( darr.random.normal(size=array_size, chunks=chunk_size), dims=("t", "y", "x"), name="X", )

Y = xr.DataArray( darr.random.normal(size=array_size, chunks=chunk_size), dims=("t", "y", "x"), name="Y", ) the "bad" / explicit way of calculating the correlationpython corr_exp = ((X - X.mean("t")) * (Y - Y.mean("t"))).mean("t") ``` leads to a graph like this:

Dask won't release any of the tasks defining X and Y until the marked substraction tasks are done.

The "good" / aggregating way of calculting the correlation python corr_agg = (X * Y).mean("t") - X.mean("t") * Y.mean("t") has the following graph where the marked multiplication and mean_chunk tasks are acting on only pairs of chunks and individual chunks and then release the original chunks of X and Y. This graph can be evaluated with a much smaller memory foot print than the other one. (It's not certain that this is always leading to lower memory use, however. But this is a different issue ...)

@willirath , thanks for your example notebook! I'm still trying to get my head around this a bit though.

Say you have da_a and da_b defined as:

```python3 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()

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() ```

The original computation in _cov_corr has a graph something like:

Whereas my alteration now has a graph more like this:

Am I correct in thinking that this is a 'better' computational graph? Because the original chunks are not passed onto later points in the computation?