Question

我想使用numpy.bincount的{{1}} API来并行化apply_ufunc函数，下面的代码是我尝试过的：

xarray

但出现以下错误：

import numpy as np
import xarray as xr
da = xr.DataArray(np.random.rand(2,16,32),
                  dims=['time', 'y', 'x'],
                  coords={'time': np.array(['2019-04-18', '2019-04-19'],
                                          dtype='datetime64'), 
                         'y': np.arange(16), 'x': np.arange(32)})

f = xr.DataArray(da.data.reshape((2,512)),dims=['time','idx'])
x = da.x.values
y = da.y.values
r = np.sqrt(x[np.newaxis,:]**2 + y[:,np.newaxis]**2)
nbins = 4
if x.max() > y.max():
    ri = np.linspace(0., y.max(), nbins)
else:
    ri = np.linspace(0., x.max(), nbins)

ridx = np.digitize(np.ravel(r), ri)

func = lambda a, b: np.bincount(a, weights=b)
xr.apply_ufunc(func, xr.DataArray(ridx,dims=['idx']), f)

我很迷失错误的根源，我们将不胜感激...

Answer 1

问题是apply_along_axis遍历第一个参数的一维切片到应用函数，而不是其他任何一个。如果我正确理解了用例，则实际上您想遍历权重(weights in the np.bincount signature)的一维切片，而不是整数数组（{{1}中的x签名）。

解决此问题的一种方法是在np.bincount周围编写一个瘦包装函数，该函数可以简单地切换参数的顺序：

np.bincount

我们可以在您的用例中将此功能与def wrapped_bincount(weights, x): return np.bincount(x, weights=weights)一起使用：

np.apply_along_axis

最后，我们可以使用apply_ufunc将此新函数包装为与xarray一起使用，请注意，它可以与dask自动并行化（还要注意，我们不需要提供def apply_bincount_along_axis(x, weights, axis=-1): return np.apply_along_axis(wrapped_bincount, axis, weights, x)参数，因为xarray会在应用函数之前自动将输入核心尺寸axis移动到dim数组中的最后一个位置）：

weights

将此功能应用于您的示例，如下所示：

def xbincount(x, weights):
    if len(x.dims) != 1:
        raise ValueError('x must be one-dimensional')

    dim, = x.dims
    nbins = x.max() + 1

    return xr.apply_ufunc(apply_bincount_along_axis, x, weights, 
        input_core_dims=[[dim], [dim]],
        output_core_dims=[['bin']], dask='parallelized',
        output_dtypes=[np.float], output_sizes={'bin': nbins})

根据需要，它也可以用于dask数组：

xbincount(ridx, f)

<xarray.DataArray (time: 2, bin: 5)>
array([[  0.      ,   7.934821,  34.066872,  51.118065, 152.769169],
       [  0.      ,  11.692989,  33.262936,  44.993856, 157.642972]])
Dimensions without coordinates: time, bin

我可以使用xarray.apply_ufunc并行化numpy.bincount吗？

1 个答案: