Pass DataArrays in apply_ufunc kwargs as numpy arrays
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 4.2k
- Forks
- 1.4k
- Avg merge
- 2d 15h
- Merged PRs (30d)
- 14
Description
Is your feature request related to a problem?
apply_ufunc( func, *args, kwargs={kwargs})
I am using a vectorized function that takes as a kwarg either a numpy array of dew point temperatures or a numpy array of vapor pressures. It seems that unlike the args of apply_ufunc, DataArray kwargs are not passed to the func as numpy arrays.
Example:
When I run:
xr.apply_ufunc(thermofeel.calculate_utci, ds.temp2m, ds.windSpeed10m, ds.meanRadTemp, kwargs={'td_k' : ds.dewPoint2m})
I get the error: Cannot determine Numba type of <class 'xarray.core.dataarray.DataArray'>
If I use a numpy array for the kwarg:
xr.apply_ufunc(thermofeel.calculate_utci, ds.temp2m, ds.windSpeed10m, ds.meanRadTemp, kwargs={'td_k' : ds.dewPoint2m.to_numpy()})
it works as expected.
Describe the solution you'd like
The example above has an easy fix because the args and kwargs arrays are the same shape; however, it seems to me it might be useful to handle DataArray kwargs the same way DataArray args are handled so that vectorized functions that require array kwargs can easily work with DataArrays of different shapes using apply_ufunc.
Describe alternatives you've considered
Broadcast args and kwargs to numpy arrays manually, run function, create DataArray with result
Additional context
No response
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the apply_ufunc entry point and compare how DataArray positional arguments and kwargs are converted and aligned. Reproduce the thermofeel example with a DataArray kwarg, including differing shapes, then add coverage showing kwargs reach the function as NumPy arrays with expected broadcasting; done when the example no longer raises the Numba type error.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- numpy, python
- Domain
- data
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 43/100