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

Posted

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了我可以使用`xarray.apply_ufunc`并行化`numpy.bincount`吗?相关的知识,希望对你有一定的参考价值。

我想使用numpy.bincountapply_ufunc API并行化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)

但是我收到以下错误:

---------------------------------------------------------------------------
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

我有点迷失在错误产生的地方,非常感谢帮助......

答案

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

解决此问题的一种方法是在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

以上是关于我可以使用`xarray.apply_ufunc`并行化`numpy.bincount`吗?的主要内容,如果未能解决你的问题,请参考以下文章

一旦我可以使用 XMPP 列出用户,我如何在这些用户之间实现聊天?

我可以在 ScriptEditor、SharePoint 中使用 Json 吗?

我可以使用啥图表来显示 Html 事件和操作

HTML5 本地存储 - 我可以为每个键存储多个值,如果不能,我可以使用啥替代方案?

我可以从 WPF 窗口获取 HDC 以便我可以使用 Gdi+ API 绘制它吗?

我可以使用 Rails 制作 phongap 应用吗