声明Numba向量化返回两个变量

Question

为了提高循环的性能，我使用了Numba向量化方法。

s1 = pd.Series([1,3,5,6,8,10,1,1,1,1,1,1]) 
s2 = pd.Series([4,5,6,8,10,1,7,1,6,5,4,3]) 

ding=pd.DataFrame({'A':s1,'B':s2}) 

@numba.vectorize(['float64(int16,int16)']) 
def sumd(a,b):  
    if a==1: 
     return (a+b) 
    else: 
     return 0 

ding['sum']=sumd(ding.A,ding.B) 

现在我想返回的附加变量，它是COLS A和产品B即我的目标是使用矢量化方法的函数返回的两个变量。我不知道如何初始化numba.vectorize方法。请帮帮我。我愿意倾听任何其他方式来提高方法的效率。

我尝试的一种替代方法是以下，但这对我来说似乎有点复杂。我正在寻找更简单的方法来优化功能。提前致谢。

s1 = pd.Series([1,3,5,6,8,10,1,1,1,1,1,1]) 
s2 = pd.Series([4,5,6,8,10,1,7,1,6,5,4,3]) 

ding=pd.DataFrame({'A':s1,'B':s2}) 

@numba.vectorize(['float64(int16,int16)']) 
def sumd(a,b):  
    if a==1: 
     sumarr.append((a+b)) 
     prodarr.append(a*b) 
     return 1 
    else: 
     sumarr.append(0) 
     prodarr.append(0) 
     return 1 

sumarr=[] 
prodarr=[] 
sumd(ding.A,ding.B) 
ding['sum']=sumarr 
ding['prod']=prodarr 

Answer 1

不能从vectorize返回多个值，并使用全局列表是行不通的。我只想用一个标准的JIT函数：

@nb.jit(nopython=True) 
def sumd(a, b): 
    sumx = np.zeros_like(a, dtype=np.float64) 
    prodx = np.zeros_like(a, dtype=np.float64) 

    for i in range(a.shape[0]): 
     if a[i] == 1: 
      sumx[i] = a[i] + b[i] 
      prodx[i] = a[i] * b[i] 

    return sumx, prodx 

sumx, prodx = sumd(ding.A.values, ding.B.values) 
ding['sum'] = sumx 
ding['prod'] = prodx 

注意，我传递中的每一列的values，这样我可以在nopython模式下使用numba因为这始终是更有效的。