numpy/pandas 矢量化自定义循环

Question

我创建了一些示例代码来模仿我得到的代码：

import numpy as np 
arr = np.random.random(100)
arr2 = np.linspace(0, 1, 20)
arr3 = np.zeros(20) # this is the array i want to store the result in
for index, num in enumerate(list(arr2)):
arr3[index] = np.mean(arr[np.abs(num - arr) < 0.2])
>>> arr3
array([0.10970893, 0.1132479 , 0.14687451, 0.17257954, 0.19401919,
0.23852137, 0.29151448, 0.35715096, 0.43273118, 0.45800796,
0.52940421, 0.60345354, 0.63969432, 0.67656363, 0.72921913,
0.78330793, 0.82693675, 0.83717402, 0.86651827, 0.89782569])

我的问题是这段代码在更大的数据上运行。我想知道是否可以在不使用显式循环的情况下以矢量化的方式进行 numpy 或 pandas 的某种组合。我尝试了很多方法，但脑海中没有一些东西。

Answer 1

如果您正在处理大型数组，我会推荐一种完全不同的方法。现在，您正在整个arr中搜索arr2中的每个元素。这显然是矫枉过正。相反，您可以对排序arr进行操作，只需在从np.searchsorted获得的插入点之间求和。

如果可以，请就地arr排序：

arr.sort()

您知道间隔的宽度，因此请找到边界值。我正在使数组形状(20, 2)更容易地匹配边界：

bounds = arr2.reshape(-1, 1) + [-0.2, 0.2]

现在找到插入索引：

ind = np.searchsorted(arr, bounds)

ind的形状与bounds相同。ind[i, :]是将开始(包含)和结束(排除)索引转换为对应于arr2的第i个元素的arr。换句话说，对于任何给定的i，原始问题中的arr3[i]是arr[ind[i, 0]:ind[i, 1]].mean()。您可以将其直接用于非矢量化解决方案：

result = np.array([arr[slice(*i)].mean() for i in ind])

有几种方法可以矢量化解决方案。无论哪种情况，您都需要每次运行中的元素数：

n = np.diff(ind, axis=1).ravel()

一个容易出现舍入错误的快速而肮脏的解决方案使用indnp.cumsum和花哨的索引：

cumulative = np.r_[0, np.cumsum(arr)]
sums = np.diff(cumulative[ind], axis=1).ravel()
result = sums / n

更强大的解决方案将使用np.add.reduceat仅提取您实际需要的总和：

arr = np.r_[arr, 0]  # compensate for index past the end
sums = np.add.reduceat(arr, ind.ravel())[::2]
result = sums / n

您可以将两种方法的结果与问题中计算arr3进行比较，以验证第二种方法是否明显更准确，即使您的玩具示例也是如此。

定时

def original(arr, arr2, d):
arr3 = np.empty_like(arr2)
for index, num in enumerate(arr2):
arr3[index] = np.mean(arr[np.abs(num - arr) < d])
return arr3
def ananda(arr, arr2, d):
arr_tile = np.tile(arr, (len(arr2), 1))
arr_tile[np.abs(arr - arr2[:, None]) >= d] = np.nan
return np.nanmean(arr_tile, axis=1)
def mad_0(arr, arr2, d):
arr.sort()
ind = np.searchsorted(arr, arr2.reshape(-1, 1) + [-d, d])
return np.array([arr[slice(*i)].mean() for i in ind])
def mad_1(arr, arr2, d):
arr.sort()
ind = np.searchsorted(arr, arr2.reshape(-1, 1) + [-d, d])
n = np.diff(ind, axis=1).ravel()
sums = np.diff(np.r_[0, np.cumsum(arr)][ind], axis=1).ravel()
return sums / n
def mad_2(arr, arr2, d):
arr.sort()
ind = np.searchsorted(arr, arr2.reshape(-1, 1) + [-d, d])
n = np.diff(ind, axis=1).ravel()
arr = np.r_[arr, 0]
sums = np.add.reduceat(arr, ind.ravel())[::2]
return sums / n

输入(每次运行复位)：

np.random.seed(42)
arr = np.random.rand(100)
arr2 = np.linspace(0, 1, 1000)

结果：

original: 25.5 ms ± 278 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
ananda: 2.66 ms ± 35.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
mad_0: 14.5 ms ± 48.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
mad_1:  211 µs ± 1.41 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
mad_2:  242 µs ± 1.93 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

对于 100 个具有 1k 箱的元素，原始方法比使用np.tile慢 ~10 倍。使用列表理解仅比原始方法好 2 倍。虽然np.cumsum方法似乎比np.add.reduce快一点，但它在数值上可能不那么稳定。

使用我建议的方法的另一个优点是您可以任意更改arr2，而arr只需要排序一次。

Answer 2

除非arr2大于arr，否则通过矢量化该循环，您将获得很少的性能改进(如果有的话)。 所有矢量化方法都需要将 arr 板播到 arr2，从而给出一个大小将抵消好处的 NxM 结构。

我测量了各种方法：

import numpy as np
def original(arr,arr2):
arr3 = np.zeros(arr2.size) # this is the array i want to store the result in
for index, num in enumerate(list(arr2)):
arr3[index] = np.mean(arr[np.abs(num - arr) < 0.2])
return arr3

def vectorized1(arr,arr2):
delta  = (np.abs(arr2[:,None]-arr)<0.2).astype(np.int)
return np.sum(arr*delta,axis=1)/np.sum(delta,axis=1)
def vectorized2(arr,arr2):
return np.fromiter((np.mean(arr[np.abs(num-arr)<0.2]) for num in arr2),np.float64)
def vectorized3(arr,arr2):
return np.apply_along_axis(lambda num:np.mean(arr[np.abs(num-arr)<0.2]),1,arr2[:,None])
def vectorized4(arr,arr2):
select  = np.array([np.nan,1])[(np.abs(arr2[:,None]-arr)<0.2).astype(np.int)]
return np.nanmean(select*arr,axis=1)
def vectorized5(arr,arr2):
arr_tile = np.tile(arr, (len(arr2), 1))
arr_tile[np.abs(arr - arr2[:, None]) >= 0.2] = np.NaN
return np.nanmean(arr_tile, axis=1)

请注意，矢量化 2 和矢量化 3 实际上并不是计算的矢量化。它们只是隐藏正在执行的循环。Vectorized5是Ananda的解决方案

当 arr 大于 arr2 时的结果：

from timeit import timeit
count = 1
arr  = np.random.random(10000)
arr2 = np.linspace(0, 1, 2000)
t = timeit(lambda:original(arr,arr2),number=count)
print("original    time:",t)
t = timeit(lambda:vectorized1(arr,arr2),number=count)
print("vectorized1 time:",t)
t = timeit(lambda:vectorized2(arr,arr2),number=count)
print("vectorized2 time:",t)
t = timeit(lambda:vectorized3(arr,arr2),number=count)
print("vectorized3 time:",t)
t = timeit(lambda:vectorized4(arr,arr2),number=count)
print("vectorized4 time:",t)
t = timeit(lambda:vectorized5(arr,arr2),number=count)
print("vectorized5 time:",t)
original    time: 0.14478049999999998
vectorized1 time: 0.3868172580000001
vectorized2 time: 0.14587923599999986
vectorized3 time: 0.15062318699999988
vectorized4 time: 0.6438709420000002
vectorized5 time: 0.543624409

当 arr2 大于 arr 时的结果：

arr  = np.random.random(100)
arr2 = np.linspace(0, 1, 200000)
print()
print(f"arr {arr.size}, arr2 {arr2.size}")
t = timeit(lambda:original(arr,arr2),number=count)
print("original    time:",t)
t = timeit(lambda:vectorized1(arr,arr2),number=count)
print("vectorized1 time:",t)
t = timeit(lambda:vectorized2(arr,arr2),number=count)
print("vectorized2 time:",t)
t = timeit(lambda:vectorized3(arr,arr2),number=count)
print("vectorized3 time:",t)
t = timeit(lambda:vectorized4(arr,arr2),number=count)
print("vectorized4 time:",t)
t = timeit(lambda:vectorized5(arr,arr2),number=count)
print("vectorized5 time:",t)
original    time: 1.7699030359999997
vectorized1 time: 0.38871579499999953
vectorized2 time: 1.782099327
vectorized3 time: 2.443001885
vectorized4 time: 0.5951444290000012
vectorized5 time: 0.4536258110000002

请注意，这种明显的改进只适用于某一点，因为当超过内存时，即使 ARR2 大于 ARR，矢量化时间也会再次爆炸。

Answer 3

解决此问题的一种方法是将所有不符合您条件的数字设置为nan并取其余数字的平均值。

import numpy as np 
arr = np.random.random((100))
arr2 = np.linspace(0,1,20)
arr3 = np.zeros(20) # this is the array i want to store the result in...
for index,num in enumerate(list(arr2)): 
arr3[index] = np.mean(arr[np.abs(num - arr) < 0.2])

arr_tile = np.tile(arr, (len(arr2), 1))
arr_tile[np.abs(arr - arr2[:, None]) >= 0.2] = np.NaN
res = np.nanmean(arr_tile, axis=1)
np.allclose(res, arr3)

这给了True