Apply a function in data.table grouped by a column

user_1393348 0
Viewed 82

I would like to apply a function in data.table, for different subsets of the data. This example will hopefully illustrate what I would like to achieve:

library(data.table)

# generate data
set.seed(123)
(dt = data.table(id = 1:20, 
                grp = sample(letters[1:3], size = 20, replace = TRUE), 
                R = sample(255, size = 20),
                G = sample(255, size = 20),
                B = sample(255, size = 20)))
#>     id grp   R   G   B
#>  1:  1   c 137   7 141
#>  2:  2   c 221 137 210
#>  3:  3   c  99 169  97
#>  4:  4   b  72  74 249
#>  5:  5   c  26  23  91
#>  6:  6   b   7 155 153
#>  7:  7   b 170 188  38
#>  8:  8   b 255  53  21
#>  9:  9   c 211 135 207
#> 10: 10   a 164 248  41
#> 11: 11   b  78 250 175
#> 12: 12   b  81 224  90
#> 13: 13   a  43 166  60
#> 14: 14   b 103 217 223
#> 15: 15   c 117  34  16
#> 16: 16   a  76 221 116
#> 17: 17   c 143  69  94
#> 18: 18   c  32  72   6
#> 19: 19   a 234  76 235
#> 20: 20   a 109  63 200

Suppose I want to apply the following function, applied to 3 columns ("R", "G", "B") for each group in "grp". So it takes 3 vectors of length n and returns 1 vector of length n.

fun = function(x1, x2, x3) {
    normalize = function(x) (x - min(x)) / diff(range(x))
    sqrt(normalize(x1)^2 + normalize(x2)^2 + normalize(x3)^2)
}

# mapping the column names of dt to the argument names of fun
vars = c(x1 = "R", x2 = "G", x3 = "B")

The following code produces the output I am looking for, but I am looking for a more efficient solution.

# solution, but very ugly and inefficient
dtgs = lapply(letters[1:3], function(g) {
    dtg = dt[grp==g,]
    dtg[, value := do.call(fun, unname(as.list(dtg[, vars, with = FALSE])))]
})
rbindlist(dtgs)
#>     id grp   R   G   B     value
#>  1: 10   a 164 248  41 1.1837788
#>  2: 13   a  43 166  60 0.5653052
#>  3: 16   a  76 221 116 0.9532667
#>  4: 19   a 234  76 235 1.4159583
#>  5: 20   a 109  63 200 0.8894540
#>  6:  4   b  72  74 249 1.0392584
#>  7:  6   b   7 155 153 0.7766996
#>  8:  7   b 170 188  38 0.9524469
#>  9:  8   b 255  53  21 1.0000000
#> 10: 11   b  78 250 175 1.2402336
#> 11: 12   b  81 224  90 0.9664781
#> 12: 14   b 103 217 223 1.2758577
#> 13:  1   c 137   7 141 0.8729010
#> 14:  2   c 221 137 210 1.6260248
#> 15:  3   c  99 169  97 1.1572081
#> 16:  5   c  26  23  91 0.4282122
#> 17:  9   c 211 135 207 1.5796092
#> 18: 15   c 117  34  16 0.4979543
#> 19: 17   c 143  69  94 0.8321982
#> 20: 18   c  32  72   6 0.4024126
data.table r
3 Answers

Try the code below

dt[, value := do.call(fun,unname(.SD)),grp,.SDcols = c("R","G","B")]

which gives

    id grp   R   G   B     value
 1:  1   c 137   7 141 0.8729010
 2:  2   c 221 137 210 1.6260248
 3:  3   c  99 169  97 1.1572081
 4:  4   b  72  74 249 1.0392584
 5:  5   c  26  23  91 0.4282122
 6:  6   b   7 155 153 0.7766996
 7:  7   b 170 188  38 0.9524469
 8:  8   b 255  53  21 1.0000000
 9:  9   c 211 135 207 1.5796092
10: 10   a 164 248  41 1.1837788
11: 11   b  78 250 175 1.2402336
12: 12   b  81 224  90 0.9664781
13: 13   a  43 166  60 0.5653052
14: 14   b 103 217 223 1.2758577
15: 15   c 117  34  16 0.4979543
16: 16   a  76 221 116 0.9532667
17: 17   c 143  69  94 0.8321982
18: 18   c  32  72   6 0.4024126
19: 19   a 234  76 235 1.4159583
20: 20   a 109  63 200 0.8894540
user_121587570 
dt[, value := fun(R, G, B), by = grp]

you can somewhat improve normalize :

normalize <- function(x){ 
  
  xx <- range(x)
  (x - xx[[1]]) / (x[[2]] - x[[1]])
}
user_121484020

Here is a direct (imo) option:

dt[, paste0(vars, "_n") := lapply(.SD, normalize), grp, .SDcols=vars]
dt[, ans := sqrt(Reduce(`+`, lapply(.SD, `^`, 2))), grp, .SDcols=paste0(vars, "_n")]
setorder(dt, grp, id)[]
user_19894800
Related