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

比如说，我有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

必须包含

f

对

S

的每个分割的结果

len（结果）=len（S）

---

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

首先，让我们让函数更高效一些。您正在执行一些不必要的索引操作：而不是

y\u left[lcond].shape[0]

您只需要

lcond.sum（）

，或者

len（lcond.nonzero（）[0]）

，这似乎更快

下面是代码的改进循环版本（包括伪输入）：

这些变化相当简单

现在，事实证明，我们确实可以向量化你的循环。为此，我们必须同时使用

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)                                           

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

>>> %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版本的两倍，矢量化版本的速度是原来的三倍。当然，后一种改进的代价是增加了内存需求：不再使用shape

（n，）

数组，而是使用shape

（nS，n）

数组，如果输入数组太大，这些数组可能会变得很大。但是，正如他们所说，没有免费的午餐，使用矢量化，您常常用运行时换取内存。

无法使用任意函数神奇地进行矢量化。您必须实现函数本身，以便在其中使用矢量化操作（通常是多维数组上的算术操作）。你的功能是什么？我编辑了代码和问题文本

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