如何高效插值63×479060大小Numpy数组中的NaN值？

问题背景 #

需要替换一个63×479060的NumPy数组中的所有NaN值，尝试使用griddata进行二维插值，但因数组规模过大未得到预期结果，原代码如下：

import numpy as np
from scipy.interpolate import griddata

a = np.arange(30180780).reshape((63, 479060)).astype(float)
a[np.random.randint(2, size=(63, 479060)).astype(bool)] = np.NaN
x, y = np.indices(a.shape)
interp = np.array(a)
interp[np.isnan(interp)] = griddata(
    (x[~np.isnan(a)], y[~np.isnan(a)]), 
    a[~np.isnan(a)],                    
    (x[np.isnan(a)], y[np.isnan(a)]))   

高效解决方案 #

针对这类行数少、列数极多的大型数组，二维插值（如griddata）效率极低，推荐采用一维逐维度插值的方式，以下是几种可行方案：

1. 使用Pandas interpolate（简洁高效） #

Pandas的插值函数针对大型数组做了优化，能快速处理沿轴的NaN填充：

import pandas as pd
import numpy as np

# 转换为DataFrame并沿列方向线性插值
df = pd.DataFrame(a)
interp_array = df.interpolate(axis=0, method='linear').to_numpy()

• 核心优势：代码简洁，自动处理边缘NaN（默认用首尾值填充），内部实现优化，可轻松应对百万级列规模。
• 可选配置：method参数可指定'nearest'（最近邻）、'quadratic'（二次插值）等，axis可设为1改为沿行插值，根据数据分布选择。

2. SciPy interp1d手动一维插值 #

若不依赖Pandas，可使用SciPy的一维插值工具逐行/列处理：

from scipy.interpolate import interp1d

interp_array = a.copy()
rows, cols = a.shape

# 逐行处理（仅循环63次，成本极低）
for i in range(rows):
    row_data = a[i, :]
    nan_mask = np.isnan(row_data)
    if np.all(nan_mask):
        continue  # 整行全NaN可跳过或自定义填充逻辑
    
    # 获取非NaN的索引与对应值
    valid_idx = np.where(~nan_mask)[0]
    valid_vals = row_data[valid_idx]
    
    # 创建插值函数并填充NaN
    f = interp1d(valid_idx, valid_vals, kind='linear', fill_value="extrapolate")
    interp_array[i, nan_mask] = f(np.where(nan_mask)[0])

• 核心优势：基于纯数值计算库，可灵活自定义边缘填充规则，适合对依赖库有严格限制的场景。

3. 批量预处理优化（减少重复计算） #

通过提前预处理非NaN索引，进一步提升效率：

from scipy.interpolate import interp1d

interp_array = a.copy()
rows, cols = a.shape
col_indices = np.arange(cols)[np.newaxis, :]

# 预处理每行的非NaN索引与值
valid_mask = ~np.isnan(a)
valid_idxs = [col_indices[i, valid_mask[i, :]] for i in range(rows)]
valid_vals = [a[i, valid_mask[i, :]] for i in range(rows)]

# 批量插值填充
for i in range(rows):
    f = interp1d(valid_idxs[i], valid_vals[i], kind='linear', fill_value="extrapolate")
    interp_array[i, np.isnan(a[i, :])] = f(np.where(np.isnan(a[i, :]))[0])

为什么griddata不适用？ #

griddata针对不规则网格的二维插值，需要对所有NaN点进行全局空间计算，时间复杂度为O(N log N)（N为非NaN点数量）。对于近3亿元素的数组，非NaN点可能达数亿级别，会导致内存溢出或计算超时。而一维插值时间复杂度为O(M)（M为单维度元素数），仅需处理63次一维计算，效率提升几个数量级。

内容的提问来源于stack exchange，提问作者r.carnahan