python-2.7 - 对列表列表运行计算的最快方法

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

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