R multiple loops using purrr where output is a data table/tibble/matrix

Viewed 72
rc<-1 # the number of red balls
wc<-1 # the number of white balls
red<-rep("Red",rc)
white<-rep("White",wc)
jar<-c(red,white)
nb<-5 ### number of draws
N=10 # Number of experiments
new_data<-matrix(NA,nrow = N,ncol = nb+1)
count_red_new<-matrix(NA,nrow = N,ncol = nb)
red_count<-rep(NA,N)
#### nested loop
for(i in 1:N){
  for(j in 1:nb){
    new_data[i,j]<-sample(jar,1,replace = T)
    if(new_data[i,j]=="Red"){
      count_red_new[i,j]=1
    }else{
      count_red_new[i,j]=0
    }
    red_count[i]=sum(count_red_new[i,])
    new_data[i,nb+1]=red_count[i]
  }
}
colnames(new_data)<-c("Draw1","Draw2","Draw3","Draw4","Draw5","Red Ball Count")
new_data<-data.frame(new_data)
#new_data$Red.Ball.Count<-as.integer(new_data$Red.Ball.Count)
new_data

The above code is used to run an experiment with respect to pulling red and white balls in 5 draws. This experiment is carried out for 10 iterations. I have used a nested loop in between as above. How to use purrr for nested loop so that the output can be stored in a matrix or a dataframe?

2 Answers

Here's one approach with map. If you only have red and white balls, you can take advantage of rbinom() for draws. Assume draws of 1 are red balls.

library(tidyverse)

nb <- 1:5
n <- 10

nb %>% 
  set_names(paste0("draw", .)) %>% 
  map_df(~rbinom(n, 1, prob = .5)) %>% 
  rowwise() %>% 
  mutate(red_ball_count = sum(c_across(everything())))

# A tibble: 10 x 6
# Rowwise: 
   draw1 draw2 draw3 draw4 draw5 red_ball_count
   <int> <int> <int> <int> <int>          <int>
 1     1     1     1     0     0              3
 2     0     1     0     0     1              2
 3     0     0     0     1     1              2
 4     0     0     0     1     1              2
 5     1     0     0     0     0              1
 6     1     0     1     0     1              3
 7     1     0     1     1     1              4
 8     0     0     1     0     1              2
 9     1     1     0     1     0              3
10     0     1     0     1     1              3

Note that this technically performs 5 rounds of 10 draws instead of 10 rounds of 5 draws. If you really want to do 10 rounds of 5 draws, it's the same idea but you need to transpose with t().

FWIW here's a base R approach:

df <- data.frame(t(replicate(n, rbinom(nb, 1, prob = .5))))
colnames(df) <- gsub("X", "Draw", colnames(df))
df['red_ball_count'] <- rowSums(df)

Finally, given that order (usually) doesn't matter in Bernoulli trials, note that you can obtain red_ball_count with rbinom() alone, by specifying n=10 and size=5:

rbinom(n, nb, prob = .5)

You could reduce complexity and add expressiveness by using a functional approach.

This simulates 3 experiments where you randomly draw 5 balls with equal probability of red/white. It returns a list, with each item representing the sequence of random draws.

set.seed(0)

samples <- lapply(1:3, function(x)
    sample(c('red', 'white'), size = 5,
           prob = c(0.5, 0.5), replace = TRUE))

samples

[[1]]
[1] "red"   "white" "white" "red"   "red"  

[[2]]
[1] "white" "red"   "red"   "red"   "red"  

[[3]]
[1] "white" "white" "white" "red"   "white"

To summarize the number of red balls picked in each experiment, this returns a vector with one item (i.e. number of reds) per experiment.

sapply(samples, function(x) sum(x == 'red'))
[1] 3 4 1
 

If you still want everything in a data frame...

 results <- as.data.frame(Reduce(rbind, samples))
 
 results$n.reds <- Reduce("+", lapply(results,
                                      function(x) x == 'red'))
 
 results
        V1    V2    V3  V4    V5 n.reds
init   red white white red   red      3
X    white   red   red red   red      4
X.1  white white white red white      1
 
Related