如何加快分组数据帧中的跨列计算

Question

--- Crossposting from rstudio community forum for potential solutions outside of tidyverse .

The basic situation is that the calculations are independent between groups, but each group need to be fed some arguments calculated from itself. The trivial example is to find the index of the first element that's less than half of the column maxima. The only twist is that one column X need to use maxima calculated through the others A, B, C .

I have a solution using group_map (similar to do ) for my question on grouped calculation . But the performance does not appear to be optimal. It seems that summarise_at takes much longer when used with group_map (compared to timings w/o it)

library(tidyverse)

times <- 1e5
cols <- 4
df3 <- as.data.frame(x = matrix(rnorm(times * cols, mean = 5), ncol = cols)) %>% 
   rename(A = V1, B = V2, C = V3, X = V4)

df3 <- cbind(grp = rep(seq_len(1e3), each = 100), df3) %>% 
   group_by(grp)

system.time(
  df3 %>% 
    group_map(~
    { 
      all_max <- summarise_at(., vars(A:C), max) %>% mutate(X = rowMeans(.))
      map2_df(., all_max, ~match(TRUE, .x < 0.5 * .y))
    }
    )
)
#>    user  system elapsed 
#>    3.87    0.00    3.98

system.time(
  df3 %>% summarise_at(vars(A:C), max) %>% mutate(X = rowMeans(.))
)
#>    user  system elapsed 
#>    0.02    0.00    0.01

system.time(
  df3 %>% summarise_at(vars(A:X), ~match(TRUE, . < 0.5 * max(.)))
)  
#>    user  system elapsed 
#>    0.25    0.02    0.26

^{Created on 2019-04-05 by the reprex package (v0.2.1)}

Any idea to improve the performance? It seems that most functions are column based and I have not yet find a solution to do this simple task efficiently.

Answer 1

From what I can tell this accomplishes the same as your code in less than half a second on my machine:

library(data.table)
DT = as.data.table(matrix(rnorm(times * cols, mean = 5), times, cols))
setnames(DT, c('A', 'B', 'C', 'X'))
DT[ , grp := rep(seq_len(1e3), each = 100)]

setkey(DT, grp)

DT[DT[ , lapply(.SD, max), keyby = grp, .SDcols = !'X'
       ][ , X := Reduce(`+`, .SD)/ncol(.SD), .SDcols = !'grp'], {
  i.A; i.B; i.C; i.X
  lapply(names(.SD), function(j) 
    which.max(eval(as.name(j)) < .5 * eval(as.name(paste0('i.', j)))))
}, on = 'grp', by = .EACHI, .SDcols = !'grp']
#        grp V1 V2 V3 V4
#    1:    1  3 30  1  4
#    2:    2  6 15  4 10
#    3:    3  2  5  7  2
#    4:    4  8 16  5  8
#    5:    5 10  3  1  7
#   ---                 
#  996:  996 11  5  3 13
#  997:  997  3  3  3 11
#  998:  998 14 21  2 10
#  999:  999 18  2  1 41
# 1000: 1000  8  7  3  3

Essentially, you are creating a look-up table of the relevant caps and joining back.

You could separate this by writing:

lookup = 
  DT[ , lapply(.SD, max), keyby = grp, .SDcols = !'X'
     ][ , X := Reduce(`+`, .SD)/ncol(.SD), .SDcols = !'grp']
DT[lookup, on = 'grp', {
  i.A; i.B; i.C; i.X
  lapply(names(.SD), function(j) 
    which.max(eval(as.name(j)) < .5 * eval(as.name(paste0('i.', j)))))
}, by = .EACHI, .SDcols = !'grp']

Once it's separated, you also gain the flexibility of getting get (which in my experience is slower than eval(as.name()) ):

DT[lookup, on = 'grp', {
  lapply(names(.SD), function(j) 
    which.max(eval(as.name(j)) < .5 * get(paste0('i.', j))))
}, by = .EACHI, .SDcols = !'grp']
#        grp V1 V2 V3 V4
#    1:    1  1  5 26  3
#    2:    2  6  7  3  4
#    3:    3  2  6  1 13
#    4:    4  5  2 12  5
#    5:    5  9 12  2  4
#   ---                 
#  996:  996  1  3  4  1
#  997:  997  1  6  3 13
#  998:  998 10 13  9  8
#  999:  999  2  4 13  4
# 1000: 1000  7 30 19 14