运行 列表列表计算的最快方法

Fastest way to run calculations on a list of lists

我有一个这样的列表:

import numpy as np
import random
import time
import itertools

N = 1000
x =np.random.random((N,N))
y = np.zeros((N,N))
z = np.random.random((N,N))

list_of_lists = [[x, y], [y,z], [z,x]]

并且对于每个子列表,我想计算非零的数量、平均值和标准偏差。

我是这样做的:

distribution = []
alb_mean = []
alb_std = []


start = time.time()

for i in range(len(list_of_lists)):

    one_mean = []
    non_zero_l = []
    one_list = list_of_lists[i]

    for n in one_list:


        #count non_zeros
        non_zero_count = np.count_nonzero(n)
        non_zero_l.append(non_zero_count)

        #assign nans
        n = n.astype(float)
        n[n == 0.0] = np.nan

        #flatten the matrix
        n = np.array(n.flatten())
        one_mean.append(n)

    #append means and stds
    distribution.append(sum(non_zero_l))
    alb_mean.append(np.nanmean(one_mean))
    alb_std.append(np.nanstd(one_mean))


end = time.time()
print "Loop took {} seconds".format((end - start))

这需要 0.23 秒。

我尝试使用第二个选项来加快速度:

distribution = []
alb_mean = []
alb_std = []


start = time.time()

for i in range(len(list_of_lists)):

    for_mean = []

    #get one list
    one_list = list_of_lists[i]

    #flatten the list
    chain = itertools.chain(*one_list)
    flat = list(chain)

    #count non_zeros
    non_zero_count = np.count_nonzero(flat)
    distribution.append(non_zero_count)

    #remove zeros
    remove_zero = np.setdiff1d(flat ,[0.0])
    alb_mean.append(np.nanmean(remove_zero))
    alb_std.append(np.nanstd(remove_zero))

end = time.time()
print "Loop took {} seconds".format((end - start))

这实际上更慢,需要 0.88 秒。

大量的循环让我想到有更好的方法来做到这一点。我试过 numba 但它不喜欢在函数中附加。

版本 #1

好吧,在您使用循环解决方案的示例中,您正在循环两个循环 - 一个循环 3 次迭代,另一个循环 2 次迭代。所以,它已经接近矢量化了。唯一的瓶颈是 append 步骤。

完全矢量化,这是一种方法 -

a = np.array(list_of_lists, dtype=float)
zm = a!=0
avgs = np.einsum('ijkl,ijkl->i',zm,a)/zm.sum(axis=(1,2,3)).astype(float)

a[~zm] = np.nan
stds = np.nanstd(a, axis=(1,2,3))

使用与问题中相同的设置,这是我得到的时间 -

Loop took 0.150925159454 seconds
Proposed solution took 0.121352910995 seconds

版本 #2

我们可以使用 average 计算 std,从而重新使用 avgs 进一步提升:

因此,修改后的版本将是 -

a = np.asarray(list_of_lists)
zm = a!=0
N = zm.sum(axis=(1,2,3)).astype(float)
avgs = np.einsum('ijkl,ijkl->i',zm,a)/N

diffs = ((a-avgs[:,None,None,None])**2)
stds = np.sqrt(np.einsum('ijkl,ijkl->i',zm,diffs)/N)

更新时间 -

Loop took 0.155035018921 seconds
Proposed solution took 0.0648851394653 seconds