Python 3.10：移除NumPy数组NaN值的正确方法（适配polyfit）

问题背景 #

需要在使用np.polyfit对X-Y数据进行多项式拟合前，移除Y轴中含NaN值对应的X、Y数据对。

错误原因分析 #

你之前的尝试报错，核心问题有两个：

• 直接用Python列表做布尔索引：Python原生列表不支持布尔数组索引，只有NumPy数组支持这种操作，因此会抛出TypeError: only integer scalar arrays can be converted to a scalar index。
• 错误判断NaN：dataY != np.nan永远返回True，因为NaN的特性是不等于任何值（包括自身），必须用np.isnan()或np.isfinite()来检测。

正确实现方法 #

步骤如下：

1. 将Python列表转换为NumPy数组，这样才能使用布尔索引筛选数据。
2. 生成有效数据的索引：筛选出Y值不为NaN的位置（若X也可能含NaN/inf，可同时检查X的有限性）。
3. 用索引过滤X、Y数组，再传入np.polyfit进行拟合。

代码示例 #

import numpy as np

def main():
    # 原始数据（列表形式）
    dataX = [1, 2, 3, 4, 5]
    dataY = [1, np.nan, 5, np.nan, 1]
    
    # 转换为NumPy数组，支持布尔索引操作
    x = np.array(dataX)
    y = np.array(dataY)
    
    # 生成有效索引：筛选y中非NaN的位置
    valid_idx = ~np.isnan(y)
    # 若x也可能存在NaN/inf，可改用以下方式
    # valid_idx = np.isfinite(x) & np.isfinite(y)
    
    # 筛选有效数据并执行多项式拟合
    poly = np.polyfit(x[valid_idx], y[valid_idx], 2)
    print(poly)

if __name__ == "__main__":
    main()

代码说明 #

• np.array()：将原生列表转为NumPy数组，为布尔索引操作提供支持。
• ~np.isnan(y)：通过取反操作，得到Y数组中非NaN值的位置索引（布尔数组）。
• x[valid_idx]、y[valid_idx]：筛选出对应位置的有效数据，传入np.polyfit即可完成正常拟合。

内容的提问来源于stack exchange，提问作者Runsva