python - 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 的开始(包括)和结束(不包括)索引，对应于 i 的第 arr2。换句话说，对于任何给定的 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只需要排序一次。