Python 3D Arrays 中的时间序列计算不使用嵌套循环？

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 倍（这将取决于 arrays 的复杂性，但从我的几次尝试中，性能时间的增益随着增加而变得更好复杂）。

结果：