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

Question

我有一个列表如下：

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，但它并不像在函数中追加一样。

Answer 1

版本＃1

在使用loopy解决方案的示例中，您正在循环使用两个循环 - 一个具有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

我们可以使用std计算average，从而重新使用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