Python中Statsmodels OLS回归如何保留结果NaN并匹配原长度

当使用statsmodels的OLS回归并通过missing='drop'丢弃含NaN的样本后，得到的残差、拟合值等结果长度会短于原始数据。如果需要生成与原数据长度一致的结果数组，在含NaN的位置保留NaN，可以按以下步骤操作：

实现代码 #

import numpy as np
import statsmodels.api as sm

yvars = np.array([1.0, 6.0, 3.0, 2.0, 8.0, 4.0, 5.0, 2.0, np.nan, 3.0])
xvars = np.array(
    [
        [1.0, 8.0],
        [8.0, np.nan],
        [np.nan, 3.0],
        [3.0, 6.0],
        [5.0, 3.0],
        [2.0, 7.0],
        [1.0, 3.0],
        [2.0, 2.0],
        [7.0, 9.0],
        [3.0, 1.0],
    ]
)

# 执行OLS回归，自动丢弃含NaN的样本
res = sm.OLS(yvars, sm.add_constant(xvars), missing='drop').fit()

# 创建与原始y变量长度一致的全NaN数组
full_resid = np.full_like(yvars, np.nan)

# 筛选出所有y和x均不含NaN的有效样本索引
valid_mask = ~np.isnan(yvars) & ~np.isnan(xvars).any(axis=1)

# 将计算得到的残差填充到对应有效位置
full_resid[valid_mask] = res.resid

# 输出结果
print(full_resid)

输出结果 #

array([-0.71907958,        nan,        nan,  1.78811122,  1.18983701,
        2.63854267, -1.45254075, -1.54362416,        nan, -1.9012464 ])

关键步骤说明 #

• 创建全NaN数组：用np.full_like生成和原始y变量形状、数据类型完全一致的数组，确保长度匹配。
• 生成有效样本掩码：通过np.isnan分别检查y变量和x变量的NaN情况，用逻辑运算筛选出所有无NaN的样本位置，这和statsmodelsmissing='drop'的筛选逻辑完全一致。
• 填充残差：将回归得到的残差赋值到全NaN数组的有效位置，最终得到符合需求的结果。

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