python - 优雅的 numpy 数组移位和 NaN 填充？

Question

我这里有一个特定的性能问题。我正在处理气象预报时间序列，我将其编译成一个 numpy 二维数组，这样

• dim0 = 预测系列开始的时间
• dim1 = 预测范围，例如。 0 到 120 小时

现在，我希望 dim0 以小时为间隔，但某些来源仅每 N 小时产生一次预测。例如，假设 N=3 并且 dim1 中的时间步长为 M=1 小时。然后我得到类似的东西

12:00  11.2  12.2  14.0  15.0  11.3  12.0
13:00  nan   nan   nan   nan   nan   nan
14:00  nan   nan   nan   nan   nan   nan
15:00  14.7  11.5  12.2  13.0  14.3  15.1

但是当然也有 13:00 和 14:00 的信息，因为它可以从 12:00 的预测运行开始填写。所以我想以这样的方式结束:

12:00  11.2  12.2  14.0  15.0  11.3  12.0
13:00  12.2  14.0  15.0  11.3  12.0  nan
14:00  14.0  15.0  11.3  12.0  nan   nan
15:00  14.7  11.5  12.2  13.0  14.3  15.1

假设 dim0 的顺序为 1e4，dim1 的顺序为 1e2，最快的方法是什么？现在我正在一行一行地做，但是速度很慢:

nRows, nCols = dat.shape
if N >= M:
    assert(N % M == 0)  # must have whole numbers
    for i in range(1, nRows):
        k = np.array(np.where(np.isnan(self.dat[i, :])))
        k = k[k < nCols - N]  # do not overstep
        self.dat[i, k] = self.dat[i-1, k+N]

我确定一定有更优雅的方法来做到这一点？任何提示将不胜感激。

Answer 1

看， bool 索引的威力!!!

def shift_nans(arr) :
    while True:
        nan_mask = np.isnan(arr)
        write_mask = nan_mask[1:, :-1]
        read_mask = nan_mask[:-1, 1:]
        write_mask &= ~read_mask
        if not np.any(write_mask):
            return arr
        arr[1:, :-1][write_mask] = arr[:-1, 1:][write_mask]

我认为命名是对正在发生的事情的不言自明。正确切片很痛苦，但它似乎有效:

In [214]: shift_nans_bis(test_data)
Out[214]: 
array([[ 11.2,  12.2,  14. ,  15. ,  11.3,  12. ],
       [ 12.2,  14. ,  15. ,  11.3,  12. ,   nan],
       [ 14. ,  15. ,  11.3,  12. ,   nan,   nan],
       [ 14.7,  11.5,  12.2,  13. ,  14.3,  15.1],
       [ 11.5,  12.2,  13. ,  14.3,  15.1,   nan],
       [ 15.7,  16.5,  17.2,  18. ,  14. ,  12. ]])

对于计时:

tmp1 = np.random.uniform(-10, 20, (1e4, 1e2))
nan_idx = np.random.randint(30, 1e4 - 1,1e4)
tmp1[nan_idx] = np.nan
tmp1 = tmp.copy()

import timeit

t1 = timeit.timeit(stmt='shift_nans(tmp)',
                   setup='from __main__ import tmp, shift_nans',
                   number=1)
t2 = timeit.timeit(stmt='shift_time(tmp1)', # Ophion's code
                   setup='from __main__ import tmp1, shift_time',
                   number=1)

In [242]: t1, t2
Out[242]: (0.12696346416487359, 0.3427293070417363)