pydata / pydata/xarray

Pass DataArrays in apply_ufunc kwargs as numpy arrays

Open
#7,064 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.