R函数支持标量输入但无法处理向量输入的问题排查

问题背景 #

编写的对数似然函数ll在传入标量（如ll(cox.fit$coefficients)）时运行正常，但传入向量（如sq = seq(0.9, 1.1, length = 10000); ll(sq)）时触发大量number of items to replace is not a multiple of replacement length警告，且返回结果并非向量每个元素对应的似然值；但通过手动循环逐个调用向量元素时运行正常。

函数代码：

ll = function(b){
  o = numeric(length(d))
  l = numeric(length(d))
  
  for (i in 1:length(d)) {
    for (j in 1:nrow(a)) {
      if (a[j,1]>=d[i]){
        if (a[j,3]==1){
          o[i] = o[i] + 1
          
        }
        else {
          o[i] = o[i] + exp(b)
          
        }
      }
      if (a[which(a[,1]==d[i]),3] == 1) {
        l[i] = 1/o[i]
    }
      else {
        l[i] = exp(b)/o[i]
    }
    }
  }
  return(sum(log(l)))
}

问题原因 #

1. 输入类型不匹配：原函数设计为处理标量输入，当传入向量b时，exp(b)会生成与b长度相同的向量，但代码中试图将这个向量赋值给单个数值o[i]（o[i]是长度为1的标量），导致赋值时长度不匹配，触发警告。
2. 返回值逻辑错误：函数最后执行sum(log(l))返回单个标量，而非对应输入向量每个元素的似然值集合，因此即使没有警告，返回结果也不符合预期。

解决方法 #

方法1：利用原生向量化函数批量调用原函数 #

无需修改原ll函数，直接用sapply或vapply对输入向量的每个元素逐个调用ll，替代手动循环：

sq <- seq(0.9, 1.1, length = 10000)
# 使用sapply
k <- sapply(sq, ll)
# 更高效的vapply（指定返回值类型）
k <- vapply(sq, ll, FUN.VALUE = numeric(1))

方法2：重写函数，原生支持向量输入 #

通过向量化优化原函数逻辑，同时支持标量和向量输入，提升计算效率：

ll_vec <- function(b_vec) {
  # 预处理：对每个d中的值，提前找到a中a[,1]>=d[i]的行索引，避免重复计算
  d_matches <- lapply(d, function(x) which(a[, 1] >= x))
  
  # 对每个b值计算对数似然
  sapply(b_vec, function(b) {
    o <- numeric(length(d))
    l <- numeric(length(d))
    
    for (i in 1:length(d)) {
      idx <- d_matches[[i]]
      # 向量化计算o[i]：统计Group=1的数量 + exp(b)*Group=2的数量
      group1_cnt <- sum(a[idx, 3] == 1)
      group2_cnt <- sum(a[idx, 3] != 1)
      o[i] <- group1_cnt + group2_cnt * exp(b)
      
      # 获取当前d[i]对应的Group值，计算l[i]
      current_group <- a[which(a[,1] == d[i]), 3]
      l[i] <- ifelse(current_group == 1, 1/o[i], exp(b)/o[i])
    }
    sum(log(l))
  })
}

调用示例：

sq <- seq(0.9, 1.1, length = 10000)
k <- ll_vec(sq)

额外优化建议 #

原函数中的双层循环可以进一步向量化，完全去掉内层循环，大幅提升计算速度：

ll_optimized <- function(b_vec) {
  # 预处理d对应的Group值
  d_group <- sapply(d, function(x) a[which(a[,1] == x), 3])
  
  # 对每个b计算似然
  sapply(b_vec, function(b) {
    # 向量化计算每个d[i]对应的o[i]
    o <- sapply(d, function(x) {
      idx <- which(a[,1] >= x)
      sum(a[idx,3] == 1) + sum(a[idx,3] != 1)*exp(b)
    })
    # 计算l并求和对数
    l <- ifelse(d_group == 1, 1/o, exp(b)/o)
    sum(log(l))
  })
}

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