我可以使用`xarray.apply_ufunc`并行化`numpy.bincount`吗?
Posted
tags:
篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了我可以使用`xarray.apply_ufunc`并行化`numpy.bincount`吗?相关的知识,希望对你有一定的参考价值。
我想使用numpy.bincount
的apply_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 吗?
HTML5 本地存储 - 我可以为每个键存储多个值,如果不能,我可以使用啥替代方案?