numpy/pandas 矢量化自定义 for 循环

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, :] 是arr 的开始（包括）和结束（不包括）索引，对应于arr2 的ith 元素。换句话说，对于任何给定的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只需要排序一次。

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)

请注意，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，向量化时间也会再次爆炸。

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