R语言中predict()结果行数与原数据不匹配的解决方法

问题背景 #

用户尝试基于多轮随访数据预测生存概率，编写了R代码但运行时出现行数不匹配的错误。以下是原始代码：

ID <- c("A", "B", "C", "D", "E", "F")
age <- c(54, 61, 65, 55, 60, 60)
sex <- c(0, 0, 1, 1, 1, 0)
Q1 <- c(0, 0, 0, 0, 0, 0)
Q2 <- c(0, 1, 0, 0, 0, 1) 
Q3 <- c(0, 1, 1, 0, 0, 1) 
Q4 <- c(0, 1, 1, 1, 0, 1) 
Q5 <- c(0, 1, 1, 1, 0, 1)
E1 <- c(2, 1, 0, 0, 0, 0)
E2 <- c(0, 1, 2, 0, 1, 0)
E3 <- c(0, 0, 1, 0, 1, 1)
E4 <- c(1, 0, 0, 0, 0, 0)
E5 <- c(0, 0, 0, 0, 2, 2)
Sint <- c(4, 3, 4, 1, 0, 2)
surv1 <- c(1, 1, 1, 1, 1, 1)
surv2 <- c(1, NA, 0, 1, 1, 1)
surv3 <- c(1, NA, 0, 1, 1, 1)
surv4 <- c(1, NA, 0, 1, 1, 0)
surv5 <- c(1, NA, 0, 1, NA, 0)
surv6 <- c(1, NA, 0, 1, NA, 0)


dta <- data.frame(ID, age, sex, Q1, Q2, Q3, Q4, Q5, E1, E2, E3, E4, E5, Sint,
                  surv1, surv2, surv3, surv4, surv5, surv6)

# Arrays
surv_wave <- c("surv1", "surv2", "surv3", "surv4", "surv5", "surv6")
predictors_fix <- c("age", "sex", "Sint")
predictors_Q <- c("Q1", "Q2", "Q3", "Q4", "Q5")
predictors_E <- c("E1", "E2", "E3", "E4", "E5")


# Survival denominator
# Probability variables that will be predicted
dta$wsd2 <- NA
dta$wsd3 <- NA
dta$wsd4 <- NA
dta$wsd5 <- NA
dta$wsd6 <- NA

# vector of variables that will be predicted
wsurv_den <- c("wsd2", "wsd3", "wsd4", "wsd5", "wsd6")

# iterate all waves
for(i in 2:6) {
  
  # subset people who survived in the previous wave
  Subset <- subset(dta, dta[[surv_wave[i-1]]] == 1)
  
  # logistic regression
  f <- as.formula(
    paste(surv_wave[i], 
          paste(c(predictors_fix, predictors_Q[i-1], predictors_E[i-1]),
                collapse=" + "),
          sep = "~"))
  Den_surv_s <- glm(f,  family = binomial(link = "logit"), na.action=na.exclude, 
                    data = Subset)
  
  # predict probabilities of survival based on logistic regression
  Den_surv_p_s <- predict(Den_surv_s, type = "response")
  
  # Add predicted values to original dataset
  dta[dta[[surv_wave[i-1]]] == 1,][[wsurv_den[i-1]]]<-Den_surv_p_s
  
}

运行后出现错误：

Error in `[[<-.data.frame`(`*tmp*`, wsurv_den[i - 1], value = c(`1` = 0.999999999941738,  : 
  replacement has 4 rows, data has 5

用户已明确错误原因：筛选上一轮存活样本时包含了存在NA的行，但glm的na.action=na.exclude会自动排除这些样本，导致预测结果的行数与赋值目标的行数不匹配。

解决方案 #

方法1：通过行索引匹配保留NA样本 #

核心思路是先给筛选出的存活样本（包括有NA的）创建预测列，仅给模型实际使用的样本赋值，最后将结果同步回原数据集：

# 迭代所有随访轮次
for(i in 2:6) {
  # 定义当前轮次的关键列名
  prev_surv_col <- surv_wave[i-1]
  curr_wsd_col <- wsurv_den[i-1]
  curr_surv_col <- surv_wave[i]
  
  # 筛选上一轮存活的所有样本（包含有NA的）
  Subset <- dta[dta[[prev_surv_col]] == 1, ]
  
  # 构建回归公式
  f <- as.formula(
    paste(curr_surv_col, 
          paste(c(predictors_fix, predictors_Q[i-1], predictors_E[i-1]),
                collapse=" + "),
          sep = "~"))
  
  # 拟合逻辑回归模型
  Den_surv_s <- glm(f, family = binomial(link = "logit"), na.action=na.exclude, 
                    data = Subset)
  
  # 找到模型实际使用的样本在Subset中的行索引
  used_rows <- rownames(Subset) %in% rownames(Den_surv_s$data)
  
  # 初始化Subset的预测列为NA，仅给模型使用的样本赋值
  Subset[[curr_wsd_col]] <- NA
  Subset[[curr_wsd_col]][used_rows] <- predict(Den_surv_s, type = "response")
  
  # 将预测值同步回原数据集
  dta[rownames(Subset), curr_wsd_col] <- Subset[[curr_wsd_col]]
}

方法2：提前清理含NA的样本 #

在拟合模型前，直接移除筛选样本中包含模型变量NA的行，确保预测结果行数与赋值目标完全匹配：

# 迭代所有随访轮次
for(i in 2:6) {
  # 定义当前轮次的关键列名
  prev_surv_col <- surv_wave[i-1]
  curr_wsd_col <- wsurv_den[i-1]
  curr_surv_col <- surv_wave[i]
  
  # 筛选上一轮存活的样本
  Subset <- dta[dta[[prev_surv_col]] == 1, ]
  
  # 获取当前模型用到的所有变量
  model_vars <- c(curr_surv_col, predictors_fix, predictors_Q[i-1], predictors_E[i-1])
  
  # 清理样本：移除包含NA的行
  Subset_clean <- na.omit(Subset[, model_vars])
  
  # 构建回归公式（排除因变量）
  f <- as.formula(paste(curr_surv_col, paste(setdiff(model_vars, curr_surv_col), collapse=" + "), sep="~"))
  
  # 拟合逻辑回归模型
  Den_surv_s <- glm(f, family = binomial(link = "logit"), data = Subset_clean)
  
  # 预测生存概率
  Den_surv_p_s <- predict(Den_surv_s, type = "response")
  
  # 将预测值赋值回原数据集
  dta[rownames(Subset_clean), curr_wsd_col] <- Den_surv_p_s
}

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