如果满足条件，R 在 DF 上按行应用函数

Question

假设以下数据框和函数：

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)
} 

以下代码将函数my_fun1 应用于DF 中的每一行，如果是DF$cond == 'cond1'，否则是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))

有什么建议吗？

Answer 1

您应该提高对 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")