spreadOfInfection.R

library(dplyr)
library(tidyr)

choose_partner <- function(People){
  Males <- filter(People, gender == "M")
  Females <- filter(People, gender == "F")
  
  if (nrow(Males) <= nrow(Females)){
    Males$partner_id <- sample(Females$id, nrow(Males))
    People <- full_join(Males, Females,
                        by = c("partner_id" = "id")) %>%
      mutate(coupling = 1:n()) %>%
      arrange(coupling) %>%
      gather(group, gender, 
             gender.x, gender.y) %>%
      arrange(coupling) %>%
      mutate(hiv = ifelse(gender == "F", hiv.y, hiv.x)) %>%
      select(-hiv.x, -hiv.y, -group) %>%
      rename(male_id = id, 
             female_id = partner_id) %>%
      arrange(coupling, female_id, male_id) %>%
      group_by(coupling) %>%
      mutate(hiv = ifelse(is.na(male_id), hiv,
                          ifelse(hiv == "No" & any(hiv %in% "Yes"),
                                 sample(c("No", "Yes"), 1, prob = c(.9, .1)),
                                 hiv)))
  }
  else{
    Females$partner_id <- sample(Males$id, nrow(Females))
    People <- full_join(Females, Males,
                        by = c("partner_id" = "id")) %>%
      mutate(coupling = 1:n()) %>%
      arrange(coupling) %>%
      gather(group, gender, 
             gender.x, gender.y) %>%
      arrange(coupling) %>%
      mutate(hiv = ifelse(gender == "M", hiv.y, hiv.x)) %>%
      select(-hiv.x, -hiv.y, -group) %>%
      rename(female_id = id, 
             male_id = partner_id) %>%
      arrange(coupling, female_id, male_id) %>%
      group_by(coupling) %>%
      mutate(hiv = ifelse(is.na(female_id), hiv,
                          ifelse(hiv == "No" & any(hiv %in% "Yes"),
                                 sample(c("No", "Yes"), 1, prob = c(.9, .1)),
                                 hiv)))
  }
  People %>%
    gather(id_group, id, female_id, male_id) %>%
    filter(!(gender == "M" & id_group == "female_id") & 
             !(gender == "F" & id_group == "male_id")) %>%
    select(-coupling, -id_group)
}


simulate_hiv_rates <- function(nrep, People){
  hiv_rates <- data.frame(rep_num = 0:nrep,
                          n_hiv = rep(NA, nrep + 1),
                          p_hiv = rep(NA, nrep + 1))
  hiv_rates$n_hiv[1] <- sum(People$hiv == "Yes")
  hiv_rates$p_hiv[1] <- sum(People$hiv == "Yes") / nrow(People)
  
  for (i in 1:nrep){
    People <- choose_partner(People)
    hiv_rates$n_hiv[i + 1] <- sum(People$hiv == "Yes")
    hiv_rates$p_hiv[i + 1] <- sum(People$hiv == "Yes") / nrow(People)
  }
  hiv_rates
}


set.seed(104)
People <- data_frame(gender = sample(c("M", "F"), 5000, replace = TRUE),
               hiv = sample(c("Yes", "No"), 5000, replace = TRUE, prob = c(0.003, 0.997))) %>%
  mutate(id = 1:n())


set.seed(654)
(hiv_rates_10 <- simulate_hiv_rates(10, People))

set.seed(654)
(hiv_rates_100 <- simulate_hiv_rates(100, People))