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

如何解决不使用嵌套循环的 3D 数组中的 Python 时间序列计算？

我有一个很大的 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,:])

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

解决方法

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

结果：