python - 如何使用最接近的值在 3D numpy 数组中组合 NaN 的前向和后向填充？

Question

我正在处理时间序列数据（卫星图像），并希望用最接近的非缺失值填充缺失值（例如由于云）。我已经找到了非常有用的帖子最有效的方法来在 numpy 数组中前向填充 NaN 值，它回答了许多问题。数组的前向和后向填充效果很好，速度也很快。我现在想将这两种方式组合成一个方式，其中只选择“最近”的值。

import numpy as np

def np_ffill(arr, axis):
    idx_shape = tuple([slice(None)] + [np.newaxis] * (len(arr.shape) - axis - 1))
    fwd_idx = np.where(~np.isnan(arr), np.arange(arr.shape[axis])[idx_shape], 0)
    fwd_idx = np.maximum.accumulate(fwd_idx, axis=axis)
    slc = [np.arange(k)[tuple([slice(None) if dim==i else np.newaxis
        for dim in range(len(arr.shape))])]
        for i, k in enumerate(arr.shape)]
    slc[axis] = fwd_idx
    return arr[tuple(slc)]

def np_bfill(arr, axis):
    idx_shape = tuple([slice(None)] + [np.newaxis] * (len(arr.shape) - axis - 1))
    bwd_idx = np.where(~np.isnan(arr), np.arange(arr.shape[axis])[idx_shape], arr.shape[axis] - 1)
    bwd_idx = np.minimum.accumulate(bwd_idx[:,:,::-1], axis=axis)[:,:,::-1]
    slc = [np.arange(k)[tuple([slice(None) if dim==i else np.newaxis
        for dim in range(len(arr.shape))])]
        for i, k in enumerate(arr.shape)]
    slc[axis] = bwd_idx
    return arr[tuple(slc)]

def random_array(shape):
    choices = [1, 2, 3, 4, np.nan]
    out = np.random.choice(choices, size=shape)
    return out
    
ra = random_array((10, 10, 5)) # for testing, I assume 5 images with the size of 10x10 pixels
ffill = np_ffill(ra,2) # the filling should only be applied on the last axis (2)
bfill = np_bfill(ra,2)

到目前为止，我唯一的想法是比较索引fwd_idx并bwd_idx确定哪个位置更接近要填充的位置。但是，这将意味着再次创建一个 FOR 循环。是否还有一种矢量化的numpy方法？非常感谢您的帮助。

标签： pythonarraysperformancenumpynan