大型3D数组中快速进行1D线性np.NaN插值

6
我有一个3D数组(z, y, x),其shape=(92, 4800, 4800),其中沿axis 0的每个值代表不同时间点。在某些情况下,时间域内的值获取失败,导致一些值为np.NaN。在其他情况下,没有获取到任何值,所有z上的值都是np.NaN

最有效的方法是什么,可以使用线性插值来填充axis 0上的np.NaN,而忽略所有值都是np.NaN的情况?

这里是一个工作示例,它使用pandas包装器来调用scipy.interpolate.interp1d。这需要大约2秒钟来处理原始数据集中的每个切片,这意味着整个数组需要2.6小时才能处理完毕。缩小大小后的示例数据集需要大约9.5秒钟。

import numpy as np
import pandas as pd

# create example data, original is (92, 4800, 4800)
test_arr = np.random.randint(low=-10000, high=10000, size=(92, 480, 480))
test_arr[1:90:7, :, :] = -32768  # NaN fill value in original data
test_arr[:, 1:90:6, 1:90:8] = -32768

def interpolate_nan(arr, method="linear", limit=3):
    """return array interpolated along time-axis to fill missing values"""
    result = np.zeros_like(arr, dtype=np.int16)

    for i in range(arr.shape[1]):
        # slice along y axis, interpolate with pandas wrapper to interp1d
        line_stack = pd.DataFrame(data=arr[:,i,:], dtype=np.float32)
        line_stack.replace(to_replace=-37268, value=np.NaN, inplace=True)
        line_stack.interpolate(method=method, axis=0, inplace=True, limit=limit)
        line_stack.replace(to_replace=np.NaN, value=-37268, inplace=True)
        result[:, i, :] = line_stack.values.astype(np.int16)
    return result

使用示例数据集在我的计算机上的性能表现:

%timeit interpolate_nan(test_arr)
1 loops, best of 3: 9.51 s per loop

编辑:

我需要澄清的是,该代码正在生成我期望的结果。问题是-我如何优化这个过程?


1
运行示例在我的机器上大约需要9.5秒,但test_arr的形状是(92, 480, 480)。如果将其增加到真实数据集的大小(92, 4800, 4800)并使用更多的NaN进行传播,则此方法需要更长时间。 - Kersten
3个回答

7

最近我使用 numba 成功地解决了我的特定问题,并且还写了一篇相关文章。该技术与IT有关。

from numba import jit

@jit(nopython=True)
def interpolate_numba(arr, no_data=-32768):
    """return array interpolated along time-axis to fill missing values"""
    result = np.zeros_like(arr, dtype=np.int16)

    for x in range(arr.shape[2]):
        # slice along x axis
        for y in range(arr.shape[1]):
            # slice along y axis
            for z in range(arr.shape[0]):
                value = arr[z,y,x]
                if z == 0:  # don't interpolate first value
                    new_value = value
                elif z == len(arr[:,0,0])-1:  # don't interpolate last value
                    new_value = value

                elif value == no_data:  # interpolate

                    left = arr[z-1,y,x]
                    right = arr[z+1,y,x]
                    # look for valid neighbours
                    if left != no_data and right != no_data:  # left and right are valid
                        new_value = (left + right) / 2

                    elif left == no_data and z == 1:  # boundary condition left
                        new_value = value
                    elif right == no_data and z == len(arr[:,0,0])-2:  # boundary condition right
                        new_value = value

                    elif left == no_data and right != no_data:  # take second neighbour to the left
                        more_left = arr[z-2,y,x]
                        if more_left == no_data:
                            new_value = value
                        else:
                            new_value = (more_left + right) / 2

                    elif left != no_data and right == no_data:  # take second neighbour to the right
                        more_right = arr[z+2,y,x]
                        if more_right == no_data:
                            new_value = value
                        else:
                            new_value = (more_right + left) / 2

                    elif left == no_data and right == no_data:  # take second neighbour on both sides
                        more_left = arr[z-2,y,x]
                        more_right = arr[z+2,y,x]
                        if more_left != no_data and more_right != no_data:
                            new_value = (more_left + more_right) / 2
                        else:
                            new_value = value
                    else:
                        new_value = value
                else:
                    new_value = value
                result[z,y,x] = int(new_value)
    return result

这比我最初的代码快了大约20倍


很有趣,可以看一下Cython的比较。 - Moritz
如果有人能够/愿意用Cython编写该函数,我也很想看到直接的速度比较。 - Kersten
感谢您的提问和回答。我对线性插值的逻辑有了一些新的想法。我已经发布了一个新的答案,感谢您的关注! - Fei Yao

3
提问者利用 `numba` 给出了一个出色的答案,我非常感激,但是我不能完全同意 `interpolate_numba` 函数内部的内容。我认为线性插值在特定点上的逻辑不是找到其左右邻居的平均值。举例来说,假设我们有一个数组 [1,nan,nan,4,nan,6],上面的 `interpolate_numba` 函数可能会返回 [1,2.5,2.5,4,5,6](仅理论推导),而 `pandas` 包装器肯定会返回 [1,2,3,4,5,6]。相反,我认为线性插值在特定点上的逻辑是找到其左右邻居,使用它们的值确定一条直线(即斜率和截距),最后计算插值值。下面是我的代码。为了简单起见,我假设输入数据是包含 NaN 值的三维数组。我规定第一个和最后一个元素等于它们的右邻居和左邻居(即 `pandas` 中的 `limit_direction='both'`)。我不指定连续插值的最大次数(即 `pandas` 中没有 `limit`)。
import numpy as np
from numba import jit
@jit(nopython=True)
def f(arr_3d):
    result=np.zeros_like(arr_3d)
    for i in range(arr_3d.shape[1]):
        for j in range(arr_3d.shape[2]):
            arr=arr_3d[:,i,j]
            # If all elements are nan then cannot conduct linear interpolation.
            if np.sum(np.isnan(arr))==arr.shape[0]:
                result[:,i,j]=arr
            else:
                # If the first elemet is nan, then assign the value of its right nearest neighbor to it.
                if np.isnan(arr[0]):
                    arr[0]=arr[~np.isnan(arr)][0]
                # If the last element is nan, then assign the value of its left nearest neighbor to it.
                if np.isnan(arr[-1]):
                    arr[-1]=arr[~np.isnan(arr)][-1]
                # If the element is in the middle and its value is nan, do linear interpolation using neighbor values.
                for k in range(arr.shape[0]):
                    if np.isnan(arr[k]):
                        x=k
                        x1=x-1
                        x2=x+1
                        # Find left neighbor whose value is not nan.
                        while x1>=0:
                            if np.isnan(arr[x1]):
                                x1=x1-1
                            else:
                                y1=arr[x1]
                                break
                        # Find right neighbor whose value is not nan.
                        while x2<arr.shape[0]:
                            if np.isnan(arr[x2]):
                                x2=x2+1
                            else:
                                y2=arr[x2]
                                break
                        # Calculate the slope and intercept determined by the left and right neighbors.
                        slope=(y2-y1)/(x2-x1)
                        intercept=y1-slope*x1
                        # Linear interpolation and assignment.
                        y=slope*x+intercept
                        arr[x]=y
                result[:,i,j]=arr
    return result

初始化一个包含一些NaN值的三维数组,我已经检查过我的代码,可以得到与 pandas 包装器相同的答案。阅读 pandas 包装器代码可能会有点困惑,因为 pandas 只能处理二维数据。

使用我的代码

y1=np.ones((2,2))
y2=y1+1
y3=y2+np.nan
y4=y2+2
y5=y1+np.nan
y6=y4+2
y1[1,1]=np.nan
y2[0,0]=np.nan
y4[1,1]=np.nan
y6[1,1]=np.nan
y=np.stack((y1,y2,y3,y4,y5,y6),axis=0)
print(y)
print("="*10)
f(y)

使用pandas包装器
import pandas as pd
y1=np.ones((2,2)).flatten()
y2=y1+1
y3=y2+np.nan
y4=y2+2
y5=y1+np.nan
y6=y4+2
y1[3]=np.nan
y2[0]=np.nan
y4[3]=np.nan
y6[3]=np.nan
y=pd.DataFrame(np.stack([y1,y2,y3,y4,y5,y6],axis=0))
y=y.interpolate(method='linear', limit_direction='both', axis=0)
y_numpy=y.to_numpy()
y_numpy.shape=((6,2,2))
print(np.stack([y1,y2,y3,y4,y5,y6],axis=0).reshape(6,2,2))
print("="*10)
print(y_numpy)

输出结果将会相同。
[[[ 1.  1.]
  [ 1. nan]]

 [[nan  2.]
  [ 2.  2.]]

 [[nan nan]
  [nan nan]]

 [[ 4.  4.]
  [ 4. nan]]

 [[nan nan]
  [nan nan]]

 [[ 6.  6.]
  [ 6. nan]]]
==========
[[[1. 1.]
  [1. 2.]]

 [[2. 2.]
  [2. 2.]]

 [[3. 3.]
  [3. 2.]]

 [[4. 4.]
  [4. 2.]]

 [[5. 5.]
  [5. 2.]]

 [[6. 6.]
  [6. 2.]]]

使用增加到 (92,4800,4800) 的 test_arr 数据作为输入,我发现完成插值只需要约40秒!
test_arr = np.random.randint(low=-10000, high=10000, size=(92, 4800, 4800))
test_arr[1:90:7, :, :] = np.nan  # NaN fill value in original data
test_arr[2,:,:] = np.nan
test_arr[:, 1:479:6, 1:479:8] = np.nan
%time f(test_arr)

输出

CPU times: user 32.5 s, sys: 9.13 s, total: 41.6 s
Wall time: 41.6 s

很棒的答案!我的初始问题只包含给定数字系列中缺失的1个观察值。然而,你的解决方案是一个不错的解决方案,可以适用于更多的应用场景。 - Kersten
那很有道理! - Fei Yao

2
这要看情况。你需要拿出一张纸,计算如果你使用零填充这些NaN而不进行插值的话,整体统计数据会出现多大的误差。
除此之外,我认为你的插值方法有些过头了。只需找到每个NaN,然后线性插值到相邻的四个值(即在(y +- 1,x +- 1)处求和)-- 这将严格限制你的误差(自己计算!),而且你不必使用任何复杂的方法进行插值(你没有定义“method”)。
你可以尝试预先计算每个z值的一个“平均”4800x4800矩阵 -- 这应该不会太久 -- 通过在矩阵上应用一个十字形核(这非常类似于图像处理)。如果有NaN,平均值中的一些值将为NaN(每个平均像素周围都有NaN),但无所谓 -- 除非有两个相邻的NaN,否则你要替换原始矩阵中所有NaN单元格的值都是实数。
然后你只需用平均矩阵中的值替换所有NaN即可。
比较这种方法的速度与手动计算每个NaN的邻域平均值的速度。

1
插值实际上只是根据时间域中的信号进行分类的先决条件,因此不适用于0填充。目前我正在使用线性插值method="linear"。如果线性插值失败,我的备选方案是用其z轴均值替换每个NaN - Kersten
1
嗯,周围两个z值的平均值就是确切地线性插值。 - Marcus Müller
我并不是很理解 - 我在之前的评论中提到的线性插值应该相当快,并且它是一种有效的插值方法。 - Marcus Müller
如果我正确理解了你的回答,你建议使用2D线性插值。我选择进行1D插值(对于每个NaN从z+-1进行插值)。这也是pandas包装器到scipy.signal.interp1d所做的。我的回答更多是关于代码优化而不是插值选择。除非我误解了你的回答并且它更有效 - 如果是这种情况:能否用代码示例解释一下? - Kersten
啊,我明白我引起了混淆的地方。抱歉。我的答案建议使用2D插值。我的评论建议您只是使用相同的方法进行1D插值:在z方向上取前一个和后一个值并计算平均值。这是线性插值给出的值。您的Pandas方法似乎以某种方式考虑了所有92个z值(至少从运行时间来看是这样)。 - Marcus Müller
显示剩余5条评论

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接