-
-
Notifications
You must be signed in to change notification settings - Fork 1.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Internal refactor: create a generic function for applying ufuncs-like functions to xarray objects #770
Comments
I've started working on this in master...shoyer:apply_ufunc |
I had a go implementing I've really struggled at working out how the layers work - between Potentially this might be worth waiting for the more generic function before implementing this? Or should I persevere? |
@MaximilianR Take a look at my branch above. I'm entirely rewriting the guts of the computation code. Once this is in, writing |
But I totally agree -- the current code involves way too many layers of indirection and anonymous functions. |
It would be awesome to have a generic function for making functions that act like NumPy's generalized universal functions "xarray aware".
What would
xarray.apply_ufunc(func, objs, join='inner', agg_dims=None, drop_dims=None, kwargs=None)
do?Dataset
orGroupBy
instances, dispatch to specialized loops that call the remainder ofapply_ufunc
repeatedly.align
all objects along shared labels using the indicatedjoin
(for some operations, e.g.,where
, a left join is appropriate rather than an inner join).broadcast
all objects against each other to expand dimensionality along all dimensions except (optionally) those listed inagg_dims
/drop_dims
.drop_dims
should be moved to the end, for consistency with gufunc signatures.agg_dims
(if provided) into anaxis
argument usingget_axis_num
and insert it intokwargs
.func
to thedata
argument of each array to calculate the result using the providedkwargs
. The result is expected to have all the same dimensions in the provided arrays, except any listed in thedims
anddrop_dims
arguments.merge
all coordinate data together (i.e., with an n-ary version of theCoordinate.merge
method) and add these to the result array.If any of
args
are not xarray objects (e.g., they're NumPy or dask arrays), they should be skipped in operations that don't apply to them.xarray.Variable
don't align or have coordinates, for example.A concrete example of similar functionality in dask.array is
atop
. The most similar thing to this that we currently have in xarray are the_unary_op
and_binary_op
staticmethods (e.g., on DataArray), but these only handle one or two arguments, don't handle aggregated dimensions and most importantly, are difficult to apply to new operations.Here are a few concrete examples of how this could work:
The text was updated successfully, but these errors were encountered: