r - 在 mutate 部分返回几个向量的函数

Question

这是一个困扰我的问题。编写函数以在返回多个向量的变异部分中使用的最简单(也是最优雅)的方法是什么。

我举个例子。假设我有这样的功能。

f1 = function(x, n){
  y1 = rep(NA, length(x))
  y2 = rep(NA, length(x))
  y3 = rep(NA, length(x))
  y4 = rep(NA, length(x))
  for(i in (n+1):(length(x)-n)){
    idx = (i-n):(i+n)
    y1[i] = sin(mean(x[idx])/max(x[idx]))
    y2[i] = cos(mean(x[idx])/max(x[idx]))
    y3[i] = tan(mean(x[idx])/max(x[idx]))
    y4[i] = 1/tan(mean(x[idx])/max(x[idx]))
  }
  data.frame(
    y1 = y1,
    y2 = y2,
    y3 = y3,
    y4 = y4
  )
}

请不要分析它的数学意义，这只是一个例子。如您所见，此函数接受一个向量并返回四个相同长度的向量。如果我想在变异部分使用这个函数，该函数将被调用四次。不幸的是，输入向量很长，需要很长时间。

这是一个例子。

n = 10000
df = tibble(
  key = rep(c("a", "b", "c", "d"), n),
  val = rep(rnorm(n), 4)
)

f1test = function(df) df %>% 
  group_by(key) %>% 
  mutate(
    y1 = f1(val, 100) %>% pull(y1),
    y2 = f1(val, 100) %>% pull(y2),
    y3 = f1(val, 100) %>% pull(y3),
    y4 = f1(val, 100) %>% pull(y4)
  )
f1test(df)

在寻找解决方案时，我有一个稍微不同的想法，即一次返回所有四个向量，然后以某种方式将它们分开。所以我创建了第二个示例函数，它执行相同的计算，只是返回结果的方式不同。

f2 = function(x, n){
  ret = rep(NA, length(x))
  for(i in (n+1):(length(x)-n)){
    idx = (i-n):(i+n)
    ret[i] = paste(
      sin(mean(x[idx])/max(x[idx])),
      cos(mean(x[idx])/max(x[idx])),
      tan(mean(x[idx])/max(x[idx])),
      1/tan(mean(x[idx])/max(x[idx])), sep = ";")
  }
  ret
}

使用这样的函数可能看起来像这样:

f2test = function(df) df %>% 
  group_by(key) %>% 
  mutate(ret = f2(val, 100)) %>% 
  separate(ret, paste0("y", 1:4), sep=";", convert = TRUE)  
f2test(df)

您可以立即看到后者应该更快。确实对于 n = 1000，带 f2 的版本大约快 2 倍。

对于 n = 10000，它快四倍。

现在我的问题。有谁知道解决这个问题的更好(更优雅)的方法吗？

Answer 1

无论您如何构建函数，都无需计算四次

mutate(
    y1 = f1(val, 100) %>% pull(y1),
    y2 = f1(val, 100) %>% pull(y2),
    y3 = f1(val, 100) %>% pull(y3),
    y4 = f1(val, 100) %>% pull(y4)
  )

您可以改变(无需赋值)返回帧的内容:

f1 <- function(x) data.frame(mysin = sin(x), mycos = cos(x), mytan = tan(x))
mtcars %>%
  mutate(f1(cyl)) %>%
  head
#                    mpg cyl disp  hp drat    wt  qsec vs am gear carb   mysin   mycos  mytan
# Mazda RX4         21.0   6  160 110 3.90 2.620 16.46  0  1    4    4 -0.2794  0.9602 -0.291
# Mazda RX4 Wag     21.0   6  160 110 3.90 2.875 17.02  0  1    4    4 -0.2794  0.9602 -0.291
# Datsun 710        22.8   4  108  93 3.85 2.320 18.61  1  1    4    1 -0.7568 -0.6536  1.158
# Hornet 4 Drive    21.4   6  258 110 3.08 3.215 19.44  1  0    3    1 -0.2794  0.9602 -0.291
# Hornet Sportabout 18.7   8  360 175 3.15 3.440 17.02  0  0    3    2  0.9894 -0.1455 -6.800
# Valiant           18.1   6  225 105 2.76 3.460 20.22  1  0    3    1 -0.2794  0.9602 -0.291

如果你想在这个过程中分配名字，那么你可以这样做

mtcars %>%
  mutate(setNames(f1(cyl), c("A","B","C"))) %>%
  head
#                    mpg cyl disp  hp drat    wt  qsec vs am gear carb       A       B      C
# Mazda RX4         21.0   6  160 110 3.90 2.620 16.46  0  1    4    4 -0.2794  0.9602 -0.291
# Mazda RX4 Wag     21.0   6  160 110 3.90 2.875 17.02  0  1    4    4 -0.2794  0.9602 -0.291
# Datsun 710        22.8   4  108  93 3.85 2.320 18.61  1  1    4    1 -0.7568 -0.6536  1.158
# Hornet 4 Drive    21.4   6  258 110 3.08 3.215 19.44  1  0    3    1 -0.2794  0.9602 -0.291
# Hornet Sportabout 18.7   8  360 175 3.15 3.440 17.02  0  0    3    2  0.9894 -0.1455 -6.800
# Valiant           18.1   6  225 105 2.76 3.460 20.22  1  0    3    1 -0.2794  0.9602 -0.291

与您的 f1 相比，这显然是一种简化，我的意思是您可以 native 使用返回值中的所有列，而无需弄清楚如何将它们分开。

这可能会使你的其他功能

f1test = function(df) df %>% 
  group_by(key) %>% 
  mutate(f1(val, 100))

(有或没有 setNames)。