numpy/pandas 向量化自定义 for 循环
numpy/pandas vectorize custom for loop
我创建了一些示例代码来模仿我得到的代码：
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 组合起来以矢量化的方式进行，而不使用显式循环。我尝试了很多方法，但没有想到什么。
解决这个问题的一种方法是将所有不符合您条件的数字设置为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
如果您要处理大型数组，我会推荐一种完全不同的方法。现在，您正在整个 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, :] 是 arr 的开始（包括）和结束（不包括）索引，对应于 arr2 的第 i 个元素。换句话说，对于任何给定的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()
一个容易出现舍入错误的快速而肮脏的解决方案使用 np.cumsum 和花哨的索引使用 ind:
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)
对于具有 1k 个 bin 的 100 个元素，原始方法比使用 np.tile 慢约 10 倍。使用列表理解仅比原始方法好 2 倍。虽然 np.cumsum 方法似乎比 np.add.reduce 快一点，但它在数值上可能不太稳定。
使用我建议的方法的另一个好处是你可以任意改变arr2，而arr只需要排序一次。
除非 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)
请注意，vectorized2 和 vectorized3 实际上并不是计算的矢量化。它们只是隐藏了正在执行的循环。 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，向量化时间也会再次爆炸。
我创建了一些示例代码来模仿我得到的代码：
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 组合起来以矢量化的方式进行，而不使用显式循环。我尝试了很多方法，但没有想到什么。
解决这个问题的一种方法是将所有不符合您条件的数字设置为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
如果您要处理大型数组，我会推荐一种完全不同的方法。现在，您正在整个 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, :] 是 arr 的开始（包括）和结束（不包括）索引，对应于 arr2 的第 i 个元素。换句话说，对于任何给定的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()
一个容易出现舍入错误的快速而肮脏的解决方案使用 np.cumsum 和花哨的索引使用 ind:
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)
对于具有 1k 个 bin 的 100 个元素，原始方法比使用 np.tile 慢约 10 倍。使用列表理解仅比原始方法好 2 倍。虽然 np.cumsum 方法似乎比 np.add.reduce 快一点，但它在数值上可能不太稳定。
使用我建议的方法的另一个好处是你可以任意改变arr2，而arr只需要排序一次。
除非 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)
请注意，vectorized2 和 vectorized3 实际上并不是计算的矢量化。它们只是隐藏了正在执行的循环。 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，向量化时间也会再次爆炸。