如果满足条件,R 在 DF rowwise 上应用函数

R apply function on DF rowwise if condition is met

假设以下数据框和函数:

cond    var_x    var_y
cond1       2        3
cond1       3        0.1
cond2       2        2
cond2       1        0.2
my_fun1 <- function(x) {
  act_vec <- rnorm(500, x, x/4)
  RT = 0
  for (i in 1:length(act_vec)) {
      my_act = act_vec[i]
      my_rt = my_act
      RT = RT + my_rt
  }
  RT = RT/500
  return(RT)
} 
my_fun2 <- function(x,y) {
  act_vec1 <- rnorm(500, x, x/4)
  act_vec2 <- rnorm(500, y, y/4)
  RT = 0
  for (i in 1:length(act_vec1)) {
      my_act1 = act_vec1[i]
      my_act2 = act_vec1[i]*act_vec2[i]
      my_rt = min(my_act1, my_act2)
      RT = RT + my_rt
  }
  RT = RT/500
  return(RT)
} 

如果 DF$cond == 'cond1' 则以下代码将函数 my_fun1 应用于 DF 中的每一行,否则 my_fun2.

my_test_vec = c()
for (i in (1:nrow(DF))) {
  if (DF$cond[i] == 'cond1') {
    my_test = my_fun1(DF$var_x[i])
  } else {
    my_test = my_fun2(DF$var_x[i], DF$var_y[i])
  }
  my_test_vec = c(my_test_vec, my_test)
}

但是,这个 for 循环的计算效率很低。因此,我想以一种一次性应用所有操作的方式来实现它。到目前为止,我的方法不起作用,看起来像这样:

DF$results <- ifelse(DF$cond == 'cond1', 
                           my_fun1(DF$var_x), 
                           my_fun2(DF$var_x, DF$var_y))

有什么建议吗?

您应该提高对 R 向量化的了解。以下是您的函数的向量化版本。首先,代码在函数中应用向量运算。其次,整个功能也是矢量化的; Vectorize() 确保它们可以使用向量作为输入,就像在 ifelse().

中的情况一样
my_fun1A <- Vectorize(function(x) {
  act_vec <- rnorm(500, x, x/4)
  RT <- sum(act_vec) / 500
  return(RT)
}) 

my_fun2A <- Vectorize(function(x, y) {
  act_vec1 <- rnorm(500, x, x/4)
  act_vec2 <- rnorm(500, y, y/4)
  RT <- sum(pmin(act_vec1, act_vec1*act_vec2)) / 500
  return(RT)
}) 

结果

set.seed(42)
resA <- ifelse(DF$cond == 'cond1', my_fun1A(DF$var_x), my_fun2A(DF$var_x, DF$var_y))
resA
# [1] 1.9849769 2.9837980 1.9772901 0.2028555

与原版对比

Vectorize使用内部的 for 循环编辑您的原始版本:

my_fun1v <- Vectorize(my_fun1)
my_fun2v <- Vectorize(my_fun2)
set.seed(42)
res <- ifelse(DF$cond == 'cond1', my_fun1v(DF$var_x), my_fun2v(DF$var_x, DF$var_y))
stopifnot(all.equal(res, resA))  ## same result

基准

fun <- function() ifelse(DF$cond == 'cond1', my_fun1v(DF$var_x), my_fun2v(DF$var_x, DF$var_y))
funA <- function() ifelse(DF$cond == 'cond1', my_fun1A(DF$var_x), my_fun2A(DF$var_x, DF$var_y))
set.seed(42)
DF <- DF[sample(1:nrow(DF), 1e4, replace=T), ]

set.seed(42)
microbenchmark::microbenchmark(fun(), funA(), times=3)
# Unit: seconds
# expr      min       lq     mean   median       uq      max neval cld
#  fun() 5.319736 5.352284 5.392068 5.384831 5.428234 5.471636     3   b
# funA() 1.793795 1.793863 1.835880 1.793931 1.856923 1.919914     3  a 

我们可以清楚地看到超过 60% 的改进。


数据:

DF <- read.table(header=T, text="cond    var_x    var_y
cond1       2        3
cond1       3        0.1
cond2       2        2
cond2       1        0.2")