Skip to main content
 首页 » 编程设计

python之寻找改进的乏味循环

2025年05月04日12webabcd

在我的代码中,我需要多次计算一个向量的值,这些值是另一个数组的不同补丁的平均值。 这是我的代码示例,展示了我是如何做到的,但我发现它在运行时效率太低......

import numpy as np 
vector_a = np.zeros(10) 
array_a = np.random.random((100,100)) 
for i in range(len(vector_a)): 
    vector_a[i] = np.mean(array_a[:,i+20:i+40] 

有什么办法可以提高效率吗?非常欢迎任何意见或建议!非常感谢!

-是的,20 和 40 是固定的。

请您参考如下方法:

编辑:

实际上,您可以更快地完成此操作。可以通过像这样对求和列进行操作来改进以前的功能:

def rolling_means_faster1(array_a, n, first, size): 
    # Sum each relevant columns 
    sum_a = np.sum(array_a[:, first:(first + size + n - 1)], axis=0) 
    # Reshape as before 
    strides_b = (sum_a.strides[0], sum_a.strides[0]) 
    array_b = np.lib.stride_tricks.as_strided(sum_a, (n, size), (strides_b)) 
    # Average 
    v = np.sum(array_b, axis=1) 
    v /= (len(array_a) * size) 
    return v 

另一种方法是使用累加和,根据需要为每个输出元素添加和删除。

def rolling_means_faster2(array_a, n, first, size): 
    # Sum each relevant columns 
    sum_a = np.sum(array_a[:, first:(first + size + n - 1)], axis=0) 
    # Add a zero a the beginning so the next operation works fine 
    sum_a = np.insert(sum_a, 0, 0) 
    # Sum the initial `size` elements and add and remove partial sums as necessary 
    v = np.sum(sum_a[:size]) - np.cumsum(sum_a[:n]) + np.cumsum(sum_a[-n:]) 
    # Average 
    v /= (size * len(array_a)) 
    return v 

与以前的解决方案进行基准测试:

import numpy as np 
 
np.random.seed(100) 
array_a = np.random.random((1000, 1000)) 
n = 100 
first = 100 
size = 200 
 
%timeit rolling_means_orig(array_a, n, first, size) 
# 12.7 ms ± 55.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 
%timeit rolling_means(array_a, n, first, size) 
# 5.49 ms ± 43.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 
%timeit rolling_means_faster1(array_a, n, first, size) 
# 166 µs ± 874 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each) 
%timeit rolling_means_faster2(array_a, n, first, size) 
# 182 µs ± 2.04 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each) 

所以这最后两个在性能上似乎非常接近。这可能取决于输入的相对大小。


这是一个可能的向量化解决方案:

import numpy as np 
 
# Data 
np.random.seed(100) 
array_a = np.random.random((100, 100)) 
 
# Take all the relevant columns 
slice_a = array_a[:, 20:40 + 10] 
# Make a "rolling window" with stride tricks 
strides_b = (slice_a.strides[1], slice_a.strides[0], slice_a.strides[1]) 
array_b = np.lib.stride_tricks.as_strided(slice_a, (10, 100, 20), (strides_b)) 
# Take mean 
result = np.mean(array_b, axis=(1, 2)) 
 
# Original method for testing correctness 
vector_a = np.zeros(10) 
idv1 = np.arange(10) + 20 
idv2 = np.arange(10) + 40 
for i in range(len(vector_a)): 
    vector_a[i] = np.mean(array_a[:,idv1[i]:idv2[i]]) 
print(np.allclose(vector_a, result)) 
# True 

这是 IPython 中的快速基准测试(为了欣赏而增加了尺寸):

import numpy as np 
 
def rolling_means(array_a, n, first, size): 
    slice_a = array_a[:, first:(first + size + n)] 
    strides_b = (slice_a.strides[1], slice_a.strides[0], slice_a.strides[1]) 
    array_b = np.lib.stride_tricks.as_strided(slice_a, (n, len(array_a), size), (strides_b)) 
    return np.mean(array_b, axis=(1, 2)) 
 
def rolling_means_orig(array_a, n, first, size): 
    vector_a = np.zeros(n) 
    idv1 = np.arange(n) + first 
    idv2 = np.arange(n) + (first + size) 
    for i in range(len(vector_a)): 
        vector_a[i] = np.mean(array_a[:,idv1[i]:idv2[i]]) 
    return vector_a 
 
np.random.seed(100) 
array_a = np.random.random((1000, 1000)) 
n = 100 
first = 100 
size = 200 
 
%timeit rolling_means(array_a, n, first, size) 
# 5.48 ms ± 26.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 
%timeit rolling_means_orig(array_a, n, first, size) 
# 32.8 ms ± 762 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)