map_blocks: Allow passing dask-backed objects in args

[dcherian]

• Tests added
• Passes isort -rc . && black . && mypy . && flake8
• Fully documented, including whats-new.rst for all changes and api.rst for new API

It parses args and breaks any xarray objects into appropriate blocks before passing them on to the user function.

e.g.

da1 = xr.DataArray(
    np.ones((10, 20)), dims=["x", "y"], coords={"x": np.arange(10), "y": np.arange(20)}
).chunk({"x": 5, "y": 4})
da1

def sumda(da1, da2):
    #print(da1.shape)
    #print(da2.shape)
    return da1 - da2

da3 = (da1 + 1).isel(x=1, drop=True).rename({"y": "k"})
mapped = xr.map_blocks(sumda, da1, args=[da3])
xr.testing.assert_equal(da1-da3, mapped) # passes

[jhamman]

I've started testing this out and have run into one problem. Here's a simple example that uses template and Dask backed DataArrays in args

ds = xr.tutorial.load_dataset('air_temperature')


def func(X, y):
    ''' a simple reduction (assume the output can't be inferred automatically) '''
    return X.sum('time') + y.min('time')


ds = ds.chunk({'lat': 10, 'lon': 10})
X = ds['air']
y = ds['air'] **2
template = X.sum('time')

expected = func(X, y)
actual = xr.map_blocks(func, X, args=[y], template=template)
xr.testing.assert_identical(actual, expected)

This raises:

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-9-5491345c8970> in <module>
      5 
      6 expected = func(X, y)
----> 7 actual = xr.map_blocks(func, X, args=[y], template=template)
      8 xr.testing.assert_identical(actual, expected)

/srv/conda/envs/notebook/lib/python3.7/site-packages/xarray/core/parallel.py in map_blocks(func, obj, args, kwargs, template)
    424         # even if length of dimension is changed by the applied function
    425         expected["shapes"] = {
--> 426             k: output_chunks[k][v] for k, v in input_chunk_index.items()
    427         }
    428         expected["data_vars"] = set(template.data_vars.keys())  # type: ignore

/srv/conda/envs/notebook/lib/python3.7/site-packages/xarray/core/parallel.py in <dictcomp>(.0)
    424         # even if length of dimension is changed by the applied function
    425         expected["shapes"] = {
--> 426             k: output_chunks[k][v] for k, v in input_chunk_index.items()
    427         }
    428         expected["data_vars"] = set(template.data_vars.keys())  # type: ignore

KeyError: 'time'

[pep8speaks]

Hello @dcherian! Thanks for updating this PR. We checked the lines you've touched for PEP 8 issues, and found:

There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻

Comment last updated at 2020-05-29 17:44:28 UTC

[dcherian]

I fixed that on the template branch. After a rebase, this example works with assert_equal

[dcherian]

Is there a better way to do this assignment?

np.array(args) ends up computing things.

[jhamman]

@dcherian - can you resolve conflicts here?

[dcherian]

converting to object array so that we can use boolean indexing to pull out xarray objects

[dcherian]

This is ready for review.

I've minimized the diff. Once this is merged, I'll do some refactoring.

[jhamman]

Thanks @dcherian, this is quite close. Would love to get either @shoyer or @TomAugspurger to look this over but everything seems good to me. Just a series of small comments.

[jhamman]

Things seem to have died down here. I suggest we merge this as is. As a reminder, the map_blocks function is still marked as an experimental feature so I'm not too concerned about breaking things in the wild. Better to get some early feedback and iterate on the design.

[shoyer]

+1

…

On Tue, Jun 2, 2020 at 11:44 AM Joe Hamman ***@***.***> wrote: Things seem to have died down here. I suggest we merge this as is. As a reminder, the map_blocks function is still marked as an experimental feature so I'm not too concerned about breaking things in the wild. Better to get some early feedback and iterate on the design. — You are receiving this because you were mentioned. Reply to this email directly, view it on GitHub <#3818 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVVSE6VIHAUV4IUX6H3RUVCARANCNFSM4K7UHXFA> .