Search code examples
rstatisticsmodeling

Reducing number of conditional statements in R?


I have this massive block of code that has different if statements and different for statement for these if statements. Is there a way that I will be able to reduce the number of if and for loops in this code to make it more condensed and efficient?

library(tidyverse)

occ_simulation <- function(nyears, lambda, alpha, beta){
  
  data_matrix <- matrix(, nrow = nyears, ncol = 6)
  
  for (z in 1:nyears){
    data_matrix[z][1] <- z
  }
  
  for (yr in 1:nyears){
    
    poisson_sim = rpois(1, lambda)
    
    for (number_of_events in poisson_sim){
      
      if (number_of_events == 1){
        
        beta_sim = rbeta(1, alpha, beta)
        data_matrix[yr, 2] <- beta_sim
        
      } else if (number_of_events == 2){
        
          for (i in 2:3){
            
            beta_sim = rbeta(1, alpha, beta)
            data_matrix[yr, i] <- beta_sim
            
          }
        
      } else if (number_of_events == 3){
          
          for (i in 2:4){
            
            beta_sim = rbeta(1, alpha, beta)
            data_matrix[yr, i] <- beta_sim
        
          }
        
      } else if (number_of_events == 4){
        
          for (i in 2:5){
            
            beta_sim = rbeta(1, alpha, beta)
            data_matrix[yr, i] <- beta_sim
        
          }
          
      } else{
        
          for (i in 2:6){
            
            beta_sim = rbeta(1, alpha, beta)
            data_matrix[yr, i] <- beta_sim
        }
    
      }
      
    }
    
  }
  
  sorted_matrix <- cbind(data_matrix[,1],t(apply(data_matrix[,2:6],1,function(x) sort(x))))

G <- sorted_matrix %>% as.data.frame %>%
  pivot_longer(-V1) %>%
  ggplot(aes(x=factor(V1),y=value,color=name,group=name))+
  geom_point()+
  labs(color='Column',x='Time (Years)', y ='Probability')+
  theme_bw()
return(G)
}

manual = occ_simulation(10, 10, 2, 20)
manual

What this code does is take an input of the timeframe you want a simulation, then for each year it does a poisson simulation to calculate the number of events that happened in that year. Then depending on the number of events that happened in that year (up to a cut off of 5). Then depending on the number of events that happened it runs a beta simulation for each event so you have a value for the 5 "events" that happened in the year. Then it just places this into a matrix so we have are left with a matrix of values for the number of years against the values for the number of events that have occurred.

      [,1]        [,2]       [,3]       [,4]       [,5]       [,6]
 [1,]    1 0.011466314 0.05644778 0.09415982 0.24654834 0.31950041
 [2,]    2 0.005988383 0.05334661 0.07467286 0.10501165 0.18135706
 [3,]    3 0.030178740 0.04501533 0.07642787 0.11537804 0.12156121
 [4,]    4 0.033208015 0.04391486 0.10769396 0.12492813 0.19938551
 [5,]    5 0.026646013 0.03085844 0.06188760 0.10742163 0.16181556
 [6,]    6 0.021271089 0.03010617 0.08581937 0.11578462 0.40768207
 [7,]    7 0.032381818 0.06602770 0.08423590 0.10511102 0.15810456
 [8,]    8 0.030124977 0.04863202 0.08484395 0.08790361 0.09787569
 [9,]    9 0.043273119 0.05148791 0.06314097 0.10162677 0.10308642
[10,]   10 0.045976676 0.06381831 0.07352065 0.08667746 0.11843700

Cheers!


Solution

  • This does everything you talk about in the text, but skips the sorting since you don't mention it and I'm not sure why you do it:

    occ_simulation2 = function(n_year, lambda, alpha, beta, max_event){
      beta_events = matrix(rbeta(n_year * max_event, shape1 = alpha, shape2 = beta), nrow = n_year)
      n_events_per_year = rpois(n_year, lambda = lambda)
      for(i in which(n_events_per_year < max_event)) {
        beta_events[i, (n_events_per_year[i] + 1):max_event] = NA
      }
      cbind(1:n_year, beta_events)
    }
    
    occ_simulation2(n_year = 10, lambda = 10, alpha = 2, beta = 20, max_event = 5)
    #       [,1]        [,2]        [,3]       [,4]       [,5]       [,6]
    #  [1,]    1 0.035977205 0.131498127 0.06717396 0.07395549 0.16532084
    #  [2,]    2 0.092411677 0.010091762 0.14054564 0.07560796 0.17471096
    #  [3,]    3 0.033394007 0.017621993 0.06469264 0.10337253 0.19579706
    #  [4,]    4 0.165301623 0.131111006 0.03285909 0.38068480 0.05473743
    #  [5,]    5 0.007820079 0.196192197 0.25096961 0.06851775 0.35516715
    #  [6,]    6 0.028011863 0.003841574 0.16924708 0.12446178 0.06525773
    #  [7,]    7 0.133261625 0.059417090 0.09608348 0.07471339 0.08303839
    #  [8,]    8 0.076545497 0.110469131 0.23364757 0.09250536 0.03295593
    #  [9,]    9 0.051244148 0.070419370 0.07127251 0.11847306 0.04112807
    # [10,]   10 0.104567386 0.188888704 0.02556781 0.10075848 0.02456839