Apply function to vectorised column values in data.table

Question

Please consider this

library(data.table)
mydt <- 
  data.table(id = 1:100,
           p1 = sample(seq(0,1,length.out=1000),100))
mydt$p2 <- 1 - mydt$p1

I want to apply a function using as the argument a vector from columns p1 and p2 .

myFun <- function(x) {
  sample(c(1,2), 1, prob = x)
}

This works,

mydt$outcome <- apply(mydt[,2:3], 1, myFun)

but I have a 25M rows, so I reach the memory limit.

I tried this, but it doesn't work.

mydt[,mydt := mapply(myFun, p1, p2)]

Answer 1

prob argument in sample requires a vector. And to apply myFun to each row, you can use by=1:nrow(mydt) or by=1:mydt[,.N]

mydt[, chosen := myFun(c(p1, p2)), by=1:nrow(mydt)]

---

Hat-tip to @Roland for his usage of rbinom . His vectorized version for this Bernoulli trial is much faster.

> system.time(mydt[, chosen := myFun(c(p1, p2)), by=1:nrow(mydt)])
   user  system elapsed 
   4.82    0.00    4.86 
> system.time(mydt[, outcome2 := rbinom(.N, 1, p2) + 1])
   user  system elapsed 
   0.05    0.02    0.06 

data used in timings:

library(data.table)
set.seed(0L)
m <- 1e6
mydt <- data.table(id = 1:m, p1 = runif(m))[, p2 := 1 - p1]
myFun <- function(x) sample(c(1,2), 1, prob = x)

accuracy check:

n <- 0L
while (n < 1e3) {
    set.seed(n)
    mydt[, chosen := myFun(c(p1, p2)), by=1:nrow(mydt)]

    set.seed(n)
    mydt[, outcome2 := rbinom(.N, 1, p2) + 1]

    if(!all.equal(mydt$chosen, mydt$outcome2)) stop("mismatch")
    n <- n + 1
}