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

Question

我想使用 xarray 的 apply_ufunc API 并行化 numpy.bincount 函数，以下代码是我尝试过的:

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)

但我收到以下错误:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-203-974a8f0a89e8> in <module>()
     12 
     13 func = lambda a, b: np.bincount(a, weights=b)
---> 14 xr.apply_ufunc(func, xr.DataArray(ridx,dims=['idx']), f)

~/anaconda/envs/uptodate/lib/python3.6/site-packages/xarray/core/computation.py in apply_ufunc(func, *args, **kwargs)
    979                                      signature=signature,
    980                                      join=join,
--> 981                                      exclude_dims=exclude_dims)
    982     elif any(isinstance(a, Variable) for a in args):
    983         return variables_ufunc(*args)

~/anaconda/envs/uptodate/lib/python3.6/site-packages/xarray/core/computation.py in apply_dataarray_ufunc(func, *args, **kwargs)
    208 
    209     data_vars = [getattr(a, 'variable', a) for a in args]
--> 210     result_var = func(*data_vars)
    211 
    212     if signature.num_outputs > 1:

~/anaconda/envs/uptodate/lib/python3.6/site-packages/xarray/core/computation.py in apply_variable_ufunc(func, *args, **kwargs)
    558             raise ValueError('unknown setting for dask array handling in '
    559                              'apply_ufunc: {}'.format(dask))
--> 560     result_data = func(*input_data)
    561 
    562     if signature.num_outputs == 1:

<ipython-input-203-974a8f0a89e8> in <lambda>(a, b)
     11 ridx = np.digitize(np.ravel(r), ri)
     12 
---> 13 func = lambda a, b: np.bincount(a, weights=b)
     14 xr.apply_ufunc(func, xr.DataArray(ridx,dims=['idx']), f)

ValueError: object too deep for desired array

我有点迷失了错误的根源，非常感谢帮助......

Answer 1

问题是 apply_along_axis 迭代所应用函数的第一个参数的一维切片，而不是其他任何参数。如果我正确理解您的用例，您实际上想要迭代权重的一维切片 (weights in the np.bincount signature) ，不是整数数组(np.bincount 签名中的x)。

解决此问题的一种方法是围绕 np.bincount 编写一个瘦包装函数，该函数只需切换参数的顺序:

def wrapped_bincount(weights, x):
    return np.bincount(x, weights=weights)

然后我们可以将 np.apply_along_axis 与此函数一起用于您的用例:

def apply_bincount_along_axis(x, weights, axis=-1):
    return np.apply_along_axis(wrapped_bincount, axis, weights, x)

最后，我们可以使用 apply_ufunc 包装这个新函数以与 xarray 一起使用。 ，注意它可以自动与 dask 并行(另请注意，我们不需要提供 axis 参数，因为 xarray 会自动将输入核心维度 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})

将此函数应用于您的示例，如下所示:

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

根据需要，它也可以与 dask 数组一起使用:

xbincount(ridx, f.chunk({'time': 1}))

<xarray.DataArray (time: 2, bin: 5)>
dask.array<shape=(2, 5), dtype=float64, chunksize=(1, 5)>
Dimensions without coordinates: time, bin