基于Poisson回归的计数型变量NaN插补代码故障排查

我有一个全为计数型变量（如日均水果摄入量、睡眠时长等整数数据）的Pandas DataFrame，用模拟数据集测试Poisson回归插补缺失值的流程：

1. 拆分数据为完整行和含缺失值的行
2. 用完整数据为每列训练Poisson回归模型
3. 用模型预测并填充缺失值

但执行代码后，原数据df和插补后的数据df2的NaN总数完全一致，缺失值没有被填充，请求排查代码错误。

模拟数据集代码 #

import numpy as np
import pandas as pd

# Set seed for reproducibility
np.random.seed(42)

# Define the ranges and number of unique values for each variable
range_0_24 = 25
range_0_30 = 31
range_0_4 = 5

total_rows = 20000
rows_without_nans = 400
rows_with_nans = total_rows - rows_without_nans

num_columns = 14

# Generate data for the rows without NaNs
data_without_nans = {
    f'var_{i}': np.random.randint(0, range_0_24, rows_without_nans)
    if i < 5 else np.random.randint(0, range_0_30, rows_without_nans)
    if i < 10 else np.random.randint(0, range_0_4, rows_without_nans)
    for i in range(num_columns)
}

# Create DataFrame for rows without NaNs
df_without_nans = pd.DataFrame(data_without_nans)

# Generate data for the rows with at least one NaN
data_with_nans = {
    f'var_{i}': np.random.randint(0, range_0_24, rows_with_nans)
    if i < 5 else np.random.randint(0, range_0_30, rows_with_nans)
    if i < 10 else np.random.randint(0, range_0_4, rows_with_nans)
    for i in range(num_columns)
}

# Introduce NaNs randomly in the DataFrame
for i in range(num_columns):
    rows_to_introduce_nans = np.random.choice(df_without_nans.index, rows_with_nans, replace=True)
    df_without_nans.loc[rows_to_introduce_nans, f'var_{i}'] = np.nan

# Concatenate the dataframes to create the final dataset
df = pd.concat([df_without_nans, pd.DataFrame(data_with_nans)])

# Shuffle the rows in the final dataset
df = df.sample(frac=1).reset_index(drop=True)

插补代码 #

import statsmodels.api as sm

# Split data into complete data and data with missing values
complete_data = df.dropna(subset=df.columns.values.tolist())
missing_data = df[df.isnull().any(axis=1)]

# Generate a copy of the df where I'll be imputing the values
df2 = df.copy()

for col in df.columns.tolist():
    
    X = complete_data.drop(columns=[col])
    
    # Add a constant for the intercept
    X = sm.add_constant(X)  
    y = complete_data[col]
    
    # Remove rows with NaN or infinite values from the data
    mask = (~np.isinf(y)) & (~np.isnan(y))
    X = X[mask]
    y = y[mask]
    
    poisson_model = sm.Poisson(y, X).fit()
    
    X_missing = missing_data.drop(columns=[col])
    X_missing = sm.add_constant(X_missing)

    predicted_values = poisson_model.predict(X_missing)
    
    # Round the predicted values
    predicted_values = np.round(predicted_values)
    
    # Set negative predictions to zero
    predicted_values[predicted_values < 0] = 0

    # Fill missing values in the original df for the specific column
    df2.loc[missing_data.index, col] = predicted_values.values

错误排查与修正 #

1. 模拟数据集的逻辑错误 #

原模拟代码的NaN引入逻辑完全不合理：

• df_without_nans仅400行，却要为每列随机选择19600个行索引（rows_with_nans=19600），且允许重复选择，最终导致df_without_nans的几乎所有行的所有列都被设为NaN，完全不符合"部分行含缺失"的模拟需求。

2. 插补代码的核心错误 #

（1）缺失行定位错误

原代码中missing_data是任意一列存在缺失的行，但插补时需要针对当前列col的缺失行，而非所有有缺失的行。直接给所有missing_data的行赋值，会覆盖原本有值的行，同时没精准定位到当前列的缺失位置。

（2）预测输入含NaN，导致预测结果为NaN

missing_data中的行本身包含其他列的缺失值，构建X_missing时这些NaN会被带入，statsmodels的Poisson模型预测时遇到输入NaN会直接返回NaN，最终填充的还是NaN，所以df2的NaN总数没有变化。

（3）冗余的完整数据处理

原代码中对complete_data的y做了NaN/inf过滤，但complete_data已经是dropna()后的结果，y不可能有NaN，这部分代码完全多余。

修正后的完整代码 #

修正后的模拟数据集代码

import numpy as np
import pandas as pd

# 设置随机种子保证可复现
np.random.seed(42)

# 定义变量范围和行数
range_0_24 = 25
range_0_30 = 31
range_0_4 = 5

total_rows = 20000
# 调整完整行比例，避免极端情况
rows_without_nans = 16000
rows_with_nans = total_rows - rows_without_nans

num_columns = 14

# 生成无缺失数据
data_without_nans = {
    f'var_{i}': np.random.randint(0, range_0_24, rows_without_nans)
    if i < 5 else np.random.randint(0, range_0_30, rows_without_nans)
    if i < 10 else np.random.randint(0, range_0_4, rows_without_nans)
    for i in range(num_columns)
}
df_without_nans = pd.DataFrame(data_without_nans)

# 生成含缺失的数据：先生成完整数据，再随机设置NaN
data_with_nans = {
    f'var_{i}': np.random.randint(0, range_0_24, rows_with_nans)
    if i < 5 else np.random.randint(0, range_0_30, rows_with_nans)
    if i < 10 else np.random.randint(0, range_0_4, rows_with_nans)
    for i in range(num_columns)
}
df_with_nans = pd.DataFrame(data_with_nans)

# 为含缺失的DataFrame随机设置NaN：每行随机选择1-3列设为NaN
for idx in df_with_nans.index:
    num_missing = np.random.randint(1, 4)
    cols_to_miss = np.random.choice(df_with_nans.columns, num_missing, replace=False)
    df_with_nans.loc[idx, cols_to_miss] = np.nan

# 合并并打乱数据
df = pd.concat([df_without_nans, df_with_nans])
df = df.sample(frac=1).reset_index(drop=True)

修正后的插补代码

import statsmodels.api as sm

# 拆分完整数据（无任何缺失的行）
complete_data = df.dropna()

# 复制原数据用于插补
df2 = df.copy()

for col in df.columns:
    # 1. 准备训练数据：用完整数据的其他列预测当前列
    X_train = complete_data.drop(columns=[col])
    X_train = sm.add_constant(X_train)  # 添加截距项
    y_train = complete_data[col]
    
    # 训练Poisson回归模型，关闭拟合信息输出
    poisson_model = sm.Poisson(y_train, X_train).fit(disp=0)
    
    # 2. 定位当前列的缺失行索引
    col_missing_idx = df2[df2[col].isna()].index
    if len(col_missing_idx) == 0:
        continue  # 该列无缺失，跳过
    
    # 3. 处理预测输入的NaN：用完整数据的均值填充其他列的缺失
    X_missing = df2.loc[col_missing_idx].drop(columns=[col])
    X_missing = X_missing.fillna(complete_data.mean())
    X_missing = sm.add_constant(X_missing)
    
    # 4. 预测并处理结果
    predicted_values = poisson_model.predict(X_missing)
    predicted_values = np.round(predicted_values).astype(int)
    predicted_values[predicted_values < 0] = 0
    
    # 5. 填充当前列的缺失值
    df2.loc[col_missing_idx, col] = predicted_values

# 检查NaN总数变化
print(f"原数据NaN总数：{df.isna().sum().sum()}")
print(f"插补后数据NaN总数：{df2.isna().sum().sum()}")

关键修正说明 #

1. 模拟数据修正：改为对指定行数的DataFrame随机设置部分列的NaN，更贴近真实场景。
2. 缺失行精准定位：针对每一列单独找到其缺失行的索引，避免无效赋值。
3. 处理预测输入的NaN：用完整数据的均值填充预测输入中的缺失值，确保模型能生成有效预测结果。
4. 简化冗余代码：移除对完整数据的NaN过滤步骤，因为complete_data已经是无缺失的行。

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