使用scipy.stats来适合xarray DataArray

Question

I want to compute the parameters of a statistical distribution fitted over the time dimension of an xarray.DataArray.

I'd like to create a function that does something like:

from scipy import stats
import xarray as xr

def fit(arr):
    return xr.apply_ufunc(stats.norm.fit, arr, ...)

that returns a new DataArray storing the two parameters of the distribution computed over the time dimension. So if an input has dimensions (time, lat, lon), fit would return a DataArray with dimensions (params, lat, lon). The next step would be to use these parameters to compute various percentiles (eg stats.norm.ppf).

After many unsuccessful trials, I'm doubting apply_ufunc supports this use case and that I should rather do the computation using

params = np.apply_along_axis(stats.norm.fit, arr.get_axis_num('time'), arr.data)

then create the DataArray manually, copying dimensions and attributes.

Thoughts? Suggestions?

---

Here is what I ended up doing, which feels a bit like a hack:

# Fit the parameters (lazy computation)
data = dask.array.apply_along_axis(dc.fit, arr.get_axis_num('time'), arr)

# Create a DataArray with the desired dimensions to copy them over to the parameter array.
mean = arr.mean(dim='time', keep_attrs=True)
coords = dict(mean.coords.items())
coords['dparams'] = ([] if dc.shapes is None else dc.shapes.split(',')) + ['loc', 'scale']
out = xr.DataArray(data=data, coords=coords, dims=(u'dparams',) + mean.dims)
out.attrs = arr.attrs

Answer 1

Dask array includes an analogue of apply_along_axis , may be the most obvious place to start. Note that each variable of an xarray that has chunks set automatically encapsulate a dask array in the .data attribute. You may even be able to pass the xarray variable directly.