在不使用嵌套循环的情况下在3D数组中进行Python时间序列计算？

Question

我有一个很大的3D数组，每一层都是同一时间的数据。我想通过时间计算每个网格单元的时间序列统计。下面的代码块是一个非常简单的版本。

我古老的FORTRAN/C背景告诉我应该使用嵌套循环，但我不禁想，一定有一种更"Pythonic“的方法来做这件事(也许也要快得多？)。有什么帮助吗？

谢谢

# Brute-force time-series stats

import numpy as np

rows = 2
cols = 4
dates = 3

time_series = np.empty([rows, cols, dates])
samples = np.zeros([rows, cols])
means = np.zeros([rows, cols])

for r in range(rows):
    for c in range(cols):
        
        samples[r, c] = np.count_nonzero(~np.isnan(time_series[r, c, :]))
        means[r, c] = np.nanmean(time_series[r, c, :])

print("Means at each grid cell:", means)
print("Samples:", samples)

Answer 1

numpy.count_nonzero和numpy.nanmean都接受轴参数来指定它应该沿着哪个轴计算这些值，所以如果你向每个轴传递axis = 2 (2是第三个轴，即(0，1，2))，你就不需要循环了。

samples = np.count_nonzero(~np.isnan(time_series), axis = 2)
means = np.nanmean(time_series, axis = 2)

将给出与您的双循环相同的结果。

在我的jupyter笔记本上使用%%timeit和使用numpy轴参数都进行了检查，确实比循环快至少5倍(这将取决于您的数组的复杂性，但从我的几次尝试中，性能时间的增益随着复杂性的增加而变得更好)。

结果：

​

​