以矢量化方式计算数组切片上的函数答案

【问题标题】：Calculate function on array slices in a vectorized way以矢量化方式计算数组切片上的函数
【发布时间】：2019-04-15 08:04:55
【问题描述】：

假设，我有 1D numpy 数组 X（特征）和 Y（二进制类）和一个函数 f，它接受 X 和 Y 的两个切片并计算一个数字。

我还有一组索引S，我需要将X 和Y 分开。保证每个切片都不为空。

所以我的代码如下所示：

def f(x_left, y_left, x_right, y_right):
    n = x_left.shape[0] + x_right.shape[0]

    lcond = y_left == 1
    rcond = y_right == 1

    hleft = 1 - ((y_left[lcond].shape[0])**2
                     + (y_left[~lcond].shape[0])**2) / n**2

    hright = 1 - ((y_right[rcond].shape[0])**2
                     + (y_right[~rcond].shape[0])**2) / n**2

    return -(x_left.shape[0] / n) * hleft - (x_right.shape[0] / n) * hright

results = np.empty(len(S))
for i in range(len(S)):
    results[i] = f(X[:S[i]], Y[:S[i]], X[S[i]:], Y[S[i]:])

数组results 必须包含从S 每次拆分时f 的结果。

len(results) == len(S)

我的问题是如何使用 numpy 以矢量化方式执行我的计算，以使这段代码更快？

【问题讨论】：

没有办法使用任意函数神奇地矢量化。您必须自己实现函数，以便可以在其中使用矢量化运算（通常是多维数组上的算术运算）。你的函数是做什么的？
我编辑了代码和问题文本

标签： python arrays python-3.x numpy vectorization

【解决方案1】：

首先，让我们让您的函数更高效一些。您正在执行一些不必要的索引操作：您只需要 lcond.sum() 或 len(lcond.nonzero()[0]) 而不是 y_left[lcond].shape[0]，这似乎更快。

这是您的代码的改进循环版本（包含虚拟输入）：

import numpy as np           

n = 1000                     
X = np.random.randint(0,n,n) 
Y = np.random.randint(0,n,n) 
S = np.random.choice(n//2, n)

def f2(x, y, s):                                     
    """Same loopy solution as original, only faster"""
    n = x.size                                       
    isone = y == 1                                   
    lval = len(isone[:s].nonzero()[0])               
    rval = len(isone[s:].nonzero()[0])               

    hleft = 1 - (lval**2 + (s - lval)**2) / n**2     
    hright = 1 - (rval**2 + (n - s - rval)**2) / n**2

    return - s / n * hleft - (n - s) / n * hright

def time_newloop():                                   
    """Callable front-end for timing comparisons"""   
    results = np.empty(len(S))                        
    for i in range(len(S)):                           
        results[i] = f2(X, Y, S[i])                   
    return results

更改相当简单。

现在，事实证明我们确实可以矢量化您的循环。为此，我们必须同时使用S 的每个元素进行比较。我们可以做到这一点的方法是创建一个形状为(nS, n)（其中S.size == nS）的二维蒙版，它将值截止到S 的相应元素。方法如下：

def f3(X, Y, S):                                     
    """Vectorized solution working on all the data at the same time"""
    n = X.size                                                        
    leftmask = np.arange(n) < S[:,None] # boolean, shape (nS, n)      
    rightmask = ~leftmask # boolean, shape (nS, n)              

    isone = Y == 1 # shape (n,)                                 
    lval = (isone & leftmask).sum(axis=1) # shape (nS,)         
    rval = (isone & rightmask).sum(axis=1) # shape (nS,)        

    hleft = 1 - (lval**2 + (S - lval)**2) / n**2                
    hright = 1 - (rval**2 + (n - S - rval)**2) / n**2           

    return - S / n * hleft - (n - S) / n * hright # shape (nS,) 

def time_vector():                                             
    """Trivial front-end for fair timing"""                    
    return f3(X,Y,S)

将您的原始解决方案定义为time_orig()，我们可以检查结果是否相同：

>>> np.array_equal(time_orig(), time_newloop()), np.array_equal(time_orig(), time_vector())
(True, True)

以及具有上述随机输入的运行时：

>>> %timeit time_orig()
... %timeit time_newloop()
... %timeit time_vector()
... 
... 
19 ms ± 501 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
11.4 ms ± 214 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
3.93 ms ± 37.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

这意味着上面的 loopy 版本几乎是原始 loopy 版本的两倍，而矢量化版本的速度又快了三倍。当然，后一种改进的代价是增加了内存需求：你现在有形状为(nS, n) 的数组，而不是形状为(n,) 的数组，如果你的输入数组很大，它会变得相当大。但正如他们所说，没有免费的午餐，通过矢量化，您通常会以运行时换取内存。

【讨论】：