Initial commit: Three-party SIR model Shiny app

- Main app.Rmd with interactive disease model simulation
- R/sir_model.R with three-group SIR model functions
- R/plot_utils.R with plotting helper functions
- Removed hardcoded API key for security

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

R/plot_utils.R

@@ -0,0 +1,28 @@
+library(cowplot)
+library(ggplot2)
+
+#' Get legend from a ggplot object
+#'
+#' @param plot ggplot object
+#' @param legend Optional legend identifier
+#' @return Legend grob
+get_legend_35 <- function(plot, legend = NULL) {
+    legends <- get_plot_component(plot, "guide-box", return_all = TRUE)
+
+    not_empty <- !vapply(legends, function(x) inherits(x, "zeroGrob"), logical(1))
+    legends <- legends[not_empty]
+
+    if (!is.null(legend)) {
+        pattern <- paste0("guide-box-", legend)
+        idx <- grep(pattern, names(legends))
+        if (length(idx) > 0) {
+            return(legends[[idx[1]]])
+        }
+    }
+
+    if (length(legends) > 0) {
+        return(legends[[1]])
+    } else {
+        return(NULL)
+    }
+}

R/sir_model.R

@@ -0,0 +1,991 @@
+
+library(deSolve)
+library(tidyverse)
+library(reshape2)
+
+# two group SIR model (a,b and c); each compartment is split into "Protected (P)" and "Unprotected (U)"
+# groups have their own average contact rates 
+# contacts rates for groups can vary over time in response to number of deaths
+# the way this is setup currently, it is simplest to equate protective behavior with mask usage
+# both groups have a background rate of adopting protective behavior which can differ between groups
+# protective behavior is protective of both transmission and infection
+# fraction of the population in each group can vary
+# homophily (degree to which members of a group mix within group versus outside their group) can vary; 
+
+## IN THE THREE GROUP MODEL, I am thinking of
+##   a -> democrats
+##   b -> republicans
+##   c -> independents
+##
+
+
+## given:
+##   dbar_a - avg degree of group a
+##   dbar_b - avg degree of group b
+##   dbar_c - avg degree of group c
+##   N_a - size of group a
+##   N_b - size of group b
+##   N_c - size of group c
+##   h_a - homophily parameter for group a
+##   h_b - homophily parameter for group b
+##
+## note that h_a, h_b corresponds more or less to
+## what Currarini et al call beta_a 
+## the share of group a's contacts that are with other members
+## of group a are given by m_a^(1/beta_a)
+## h_a is in [1, \infty) and the bigger h_a gets, the
+## more homophily there is
+##
+## this implies that if we have a target value of h_a, then the value
+## of Currarini et al's beta_a implied by that is
+##      h_a = m_a^(1/beta_a)
+##  <=> beta_a = log(m_a) / log(h_a)
+##
+## homophily in group c is completely determined by
+## h_a, and h_b, so it cannot be varied independently
+
+
+#for three groups contact matrix
+avg_contact_matrix_3gp <- function(dbar_a, dbar_b, dbar_c,
+                                   N_a, N_b, N_c,
+                                   beta_a, beta_b) {
+  #h_a, h_b) {
+  
+  # fraction of total degree in group a and group b 
+  m_a <- (N_a * dbar_a) / ((N_a * dbar_a) + (N_b * dbar_b) + (N_c * dbar_c))
+  m_b <- (N_b * dbar_b) / ((N_a * dbar_a) + (N_b * dbar_b) + (N_c * dbar_c))
+  m_c <- (N_c * dbar_c) / ((N_a * dbar_a) + (N_b * dbar_b) + (N_c * dbar_c))
+  
+  # this is the value of currarini et al's beta parameter
+  # implied by h_a
+  # we don't do anything with this directly, but it might be useful for
+  # debugging / substantive analysis down the line
+  #implied_beta_a <- log(m_a) / log(h_a)
+  #implied_beta_b <- log(m_b) / log(h_b)
+  
+  q_aa <- m_a^(1/beta_a) 
+  q_bb <- m_b^(1/beta_b) 
+  
+  #q_aa <- h_a
+  #q_bb <- h_b
+  
+  
+  ## for a's contacts w/ other groups, assume they are distributed
+  ## in proportion to the other groups' share of edges
+  ## (NB: this ignores homophily w/in other groups; could potentially improve this)
+  ## TODO - SHOULD REVISIT THIS
+  q_ab <- (1-q_aa) * (m_b / (m_b + m_c)) 
+  q_ac <- (1-q_aa) * (m_c / (m_b + m_c)) 
+  
+  # implied by q_ab and q_ac
+  q_ba <- q_ab*(dbar_a*N_a)/(N_b*dbar_b)
+  # and then q_bc has to make them add up to 1
+  q_bc <- (1 - (q_ba + q_bb))
+  
+  # implied by q_ac
+  q_ca <- q_ac*dbar_a*N_a/(N_c*dbar_c)
+  
+  # ... and then I think q_cb will be implied by that
+  q_cb <- q_bc * (dbar_b*N_b)/(dbar_c*N_c)
+  
+  # q_cc is residual
+  q_cc <- 1 - (q_ca + q_cb)
+  
+  # implied beta_c
+  beta_c_implied <- log(m_c) / log(q_bb)
+  
+  # q matrix
+  q_beta <- t(matrix(c(q_aa, q_ab, q_ac, 
+                       q_ba, q_bb, q_bc,
+                       q_ca, q_cb, q_cc),
+                     nrow=3)) 
+  ## this matrix has the avg number of contacts between groups
+  ## it should be symmetric
+  c_mat <- q_beta
+  c_mat[1,] <- c_mat[1,] * dbar_a 
+  c_mat[2,] <- c_mat[2,] * dbar_b
+  c_mat[3,] <- c_mat[3,] * dbar_c
+  
+  return(c_mat)
+  
+}
+
+
+sir_three_group_pu <- function(## parameters related to popn
+  N0 = 10000000, # population size
+  frac_a = 0.33, # fraction of population in group A
+  frac_b = 0.33, # fraction of population in group B
+  ## parameters related to contacts
+  cmax = NA, # initial/max average number of contacts per day (specify only if its the same for both groups, otherwise NA)
+  cmax_a = 8, cmax_b = 8, cmax_c=8, # initial/max average number of contacts per day by group 
+  cmin = NA, # min average contact
+  cmin_a = 3, cmin_b = 3, cmin_c=3,
+  beta_a=1, # homophily parameter for group a (beta_a = 1 means unbiased mixing; bigger values mean homophily) 
+  beta_b=1, # homophily parameter for group b (beta_b = 1 means unbiased mixing; bigger values mean homophily) 
+  #h_a=.5, # proportion of group A's total contact with members of their own group 
+  #h_b=.5, # proportion of group A's total contact with members of their own group 
+  ##
+  zeta = NA, # responsiveness of contact to deaths
+  zeta_a = 0.01, zeta_b = 0, zeta_c = 0.005,
+  trans_p = 0.05, # probability of transmission given contact (or susceptibility to infection given contact)
+  rho=1/10,# 1 / infectious period  or recovery rate
+  mu = NA,# probability of dying following infection
+  mu_a = 0.01, mu_b = 0.01, mu_c = .01, # can let it differ between groups to crudely account for difference in age composition between groups
+  kappa=0.9, # scaling factor for probability of transmission given contact resulting from protective behavior; kappa=1 means no protection, kappa = 0 means perfect protection
+  phi = NA, # waning of protective behavior
+  phi_a = 0, phi_b=0, phi_c=0,
+  I0_a=1, I0_b=1, I0_c=1, #intial infected in each group
+  time = 500, # time steps for simulation
+  pi = NA, # background rate of adopting protective behavior
+  pi_a = 0.05, pi_b = 0.05, pi_c = 0.05,
+  ell = 1, # time window for considering deaths that influence adoption of protective behavior
+  vacc = NA, # vaccination rate (goes from S to R)
+  # in Roubenoff et al. we assume about 2 million daily doses are distributed which is ~ 0.6 % of the population per day
+  vacc_a = 0.006, vacc_b = 0.006, vacc_c = 0.006,
+  vstart = 365, # start of vaccination
+  gamma = 1/182.5, # wanning immunity
+  
+  get_params=FALSE) {
+  # set following parameters to be the same for both groups unless specified otherwise 
+  
+  if(!is.na(mu)){  
+    mu_a<-mu  
+    mu_b<-mu
+    mu_c<-mu
+  }
+  
+  if(!is.na(zeta)){  
+    zeta_a<-zeta  
+    zeta_b<-zeta
+    zeta_c<-zeta
+  }
+  
+  if(!is.na(pi)){  
+    pi_a<-pi 
+    pi_b<-pi
+    pi_c<-pi
+  }
+  
+  if(!is.na(phi)){  
+    phi_a<-phi 
+    phi_b<-phi
+    phi_c<-phi
+  } 
+  
+  
+  if(!is.na(vacc)){  
+    vacc_a<-vacc
+    vacc_b<-vacc
+    vacc_c<-vacc
+  }
+  
+  if(!is.na(cmin)){  
+    cmin_a<-cmin
+    cmin_b<-cmin
+    cmin_c<-cmin
+  }
+  
+  
+  if(!is.na(cmax)){  
+    cmax_a<-cmax
+    cmax_b<-cmax
+    cmax_c<-cmax
+  }
+  
+  N0_a = N0 * frac_a 
+  N0_b = N0 * frac_b
+  N0_c = N0 * (1 - frac_a - frac_b)
+  
+  
+  
+  params<-c("trans_p" = trans_p, # probability of transmission given contact
+            "beta_a"= beta_a, # homophily parameter for group a (beta_a = 1 means unbiased mixing; beta_a > 1 means homophily)
+            "beta_b"= beta_b, # homophily parameter for group b (beta_b = 1 means unbiased mixing; beta_b > 1 means homophily)
+            #"h_a"= h_a, # # proportion of group A's total contact with members of their own group (bounded by population size and total contacts in group B)
+            #"h_b"= h_b, # # proportion of group A's total contact with members of their own group (bounded by population size and total contacts in group B)
+            "cmax_a" = cmax_a, "cmax_b" = cmax_b, "cmax_c" = cmax_c, # upper and lower bounds for contacts
+            "zeta_a" = zeta_a, "zeta_b" = zeta_b, "zeta_c" = zeta_c, # responsiveness of contact to deaths
+            "pi_a" = pi_a, "pi_b" = pi_b, "pi_c" = pi_c, # background rate of adopting protective behavior
+            "kappa" = kappa, # scaling factor for probability of transmission given contact resulting from protective behavior; kappa=1 means no protection, kappa = 0 means perfect protection
+            "phi_a"=phi_a, "phi_b" = phi_b, "phi_c" = phi_c, # waning rate of protective behavior
+            "rho"=rho, # 1 / infectious period
+            "mu_a"=mu_a, "mu_b" = mu_b, "mu_c" = mu_c, # probability of dying following infection
+            "time"=time, # time steps in simulation
+            "I0_a"=I0_a, "I0_b"=I0_b, "I0_c"=I0_c, # starting number infected in each group
+            "N0_a" = N0_a, "N0_b"= N0_b, "N0_c"=N0_c, # population size in each group at time zero
+            "ell"=ell, # time window for considering deaths that influence adoption of protective behavior
+            "vacc_a"=vacc_a, "vacc_b"=vacc_b, "vacc_c"=vacc_c, # vaccination rate
+            "vstart" = vstart, # start of vaccination
+            "gamma" = gamma #waning immunity
+  ) # responsiveness to proportion of protected individuals for adopting protective behavior
+  
+  state<-c( 
+    
+    #group a
+    ca <- cmax_a,
+    SUa<-(N0_a - I0_a),
+    SPa<-0,
+    IUa<-I0_a,
+    #IUaa<-0,
+    #IUab <- 0,
+    #IUac <- 0,
+    IPa<-0,
+    RUa<-0,
+    RPa<-0,
+    DUa<-0,
+    DPa<-0,
+    Casesa <- I0_a,
+
+    #group b
+    cb <- cmax_b,
+    SUb<-N0_b - I0_b,
+    SPb<-0,
+    IUb<-I0_b,
+    IPb<-0,
+    RUb<-0,
+    RPb<-0,
+    DUb<-0,
+    DPb<-0,
+    Casesb <- I0_b,
+
+    #group c 
+    cc <- cmax_c,
+    SUc<-N0_c - I0_c,
+    SPc<-0,
+    IUc<-I0_c,
+    IPc<-0,
+    RUc<-0,
+    RPc<-0,
+    DUc<-0,
+    DPc<-0,
+    Casesc <- I0_c
+  )
+  
+  names(state)<- c("ca","SUa", "SPa", "IUa", "IPa", "RUa", "RPa", "DUa", "DPa", "Casesa", 
+                   #"IUaa", "IUab", "IUac",
+                   "cb","SUb", "SPb", "IUb", "IPb", "RUb", "RPb", "DUb", "DPb", "Casesb",
+                   "cc","SUc", "SPc", "IUc", "IPc", "RUc", "RPc", "DUc", "DPc", "Casesc")
+  
+  sir_up<-function(t, state, parameter){
+    
+    with(as.list(c(parameter, state)), {
+      
+      
+      if(t<ell){
+        lag_ell<-rep(0, length(state))
+      }
+      else{
+        lag_ell<-lagvalue(t-ell) #provide access to past (lagged) values of all state variables; below we make sure we only look at D
+      }
+      
+      
+      if(t<vstart){
+        vacc_a <- 0
+        vacc_b <- 0
+        vacc_c <- 0
+      }
+      else{
+        vacc_a <- vacc_a
+        vacc_b <- vacc_b
+        vacc_c <- vacc_c
+      }
+      
+      # for the transitions there are seven processes that are currently modeled:
+      # 1: transmission process (infection, recovery/death)
+      # 2: changes in contact rate for group A in response to deaths
+      # 3: adoption of protective behavior 
+      # 5: waning of protective behavior
+      # 6: vaccination
+      # 7: waning of immunity
+      
+      
+      
+      # recalculate population size (will change due to deaths)
+      N_a = SUa + IUa + RUa + SPa + IPa + RPa 
+      N_b = SUb + IUb + RUb + SPb + IPb + RPb
+      N_c = SUc + IUc + RUc + SPc + IPc + RPc
+      
+      
+      
+      # contact matrix
+      c_mat <- avg_contact_matrix_3gp(dbar_a = ca,
+                                      dbar_b = cb,
+                                      dbar_c = cc,
+                                      N_a = N_a,
+                                      N_b = N_b,
+                                      N_c = N_c,
+                                      beta_a = beta_a,
+                                      beta_b = beta_b)
+      #h_a = h_a,
+      #h_b = h_b)
+      
+      ## TODO - should we throw an error if any implied contacts are negative? 
+      
+      
+      ################################
+      ## transitions for group a
+      ################################
+      
+      # Force of infection
+      lambda_a =  trans_p*(c_mat[1,1]*(IUa/N_a + kappa*IPa/N_a) + c_mat[1,2]*(IUb/N_b + kappa*IPb/N_b) + c_mat[1,3]*(IUc/N_c + kappa*IPc/N_c))
+      #below are force of infection caused by specific groups
+      lambda_aa = trans_p*(c_mat[1,1]*(IUa/N_a + kappa*IPa/N_a))
+      lambda_ab = trans_p*(c_mat[1,2]*(IUb/N_b + kappa*IPb/N_b))
+      lambda_ac = trans_p*(c_mat[1,3]*(IUc/N_c + kappa*IPc/N_c))
+      
+      # change in contacts in response to deaths occuring over a time window defined by ell (constrained to maximum/minimum observed values)
+      
+      deaths_t <- rho*mu_a*IUa + rho*mu_a*IPa + rho*mu_b*IUb + rho*mu_b*IPb + rho*mu_c*IUc + rho*mu_c*IPc
+      
+      
+      deaths_tminusell <- rho*mu_a*(lag_ell[which(names(state) == "IUa")]) + 
+        rho*mu_a*(lag_ell[which(names(state) == "IPa")]) + 
+        rho*mu_b*(lag_ell[which(names(state) == "IUb")]) + 
+        rho*mu_b*(lag_ell[which(names(state) == "IPb")]) +
+        rho*mu_c*(lag_ell[which(names(state) == "IUc")]) + 
+        rho*mu_c*(lag_ell[which(names(state) == "IPc")])
+      
+      cum_deaths_ell <- DUa - lag_ell[which(names(state) == "DUa")] +
+        DPa - lag_ell[which(names(state) == "DPa")] +
+        DUb - lag_ell[which(names(state) == "DUb")] +
+        DPb - lag_ell[which(names(state) == "DPb")] +
+        DUc - lag_ell[which(names(state) == "DUc")] +
+        DPc - lag_ell[which(names(state) == "DPc")]
+      
+      dca <-  (-zeta_a)*(cmax_a - cmin_a)*(deaths_t - deaths_tminusell)*exp(-zeta_a*(cum_deaths_ell))
+      
+      dSUa <- -SUa*lambda_a - #transmission process
+        pi_a * SUa + # adoption of protective behavior
+        phi_a * SPa + #waning of protective behavior
+        gamma * RUa - #waning immunity
+        vacc_a * SUa
+      
+      dSPa <- -SPa*kappa*lambda_a +
+        pi_a * SUa - 
+        phi_a * SPa  +
+        gamma * RPa -
+        vacc_a * SPa
+      
+      dIUa <- SUa*lambda_a  -
+        pi_a*IUa + 
+        phi_a * IPa -
+        (rho) * IUa 
+      
+      #dIUaa <- SUa*lambda_aa - pi_a*IUa + phi_a*IPa - (rho)*IUa
+      
+      #dIUab <- SUa*lambda_ab - pi_a*IUa + phi_a*IPa - (rho)*IUa 
+      
+      #dIUac <- SUa*lambda_ab - pi_a*IUa + phi_a*IPa - (rho)*IUa 
+      
+      dIPa <- SPa*kappa*lambda_a +
+        pi_a * IUa - 
+        phi_a * IPa -
+        (rho) * IPa 
+      
+      dRUa <- rho * (1-mu_a) * IUa  -
+        pi_a * RUa + 
+        phi_a * RPa  - 
+        gamma * RUa +
+        vacc_a * SUa
+      
+      dRPa <- rho * (1-mu_a) * IPa  +
+        pi_a * RUa - 
+        phi_a * RPa  - 
+        gamma * RPa +
+        vacc_a * SPa
+      
+      dDUa <- rho*mu_a*IUa
+      
+      dDPa <- rho*mu_a*IPa
+      
+      dCasesa <- SUa*lambda_a + SPa*kappa*lambda_a 
+      
+      ################################
+      ## transitions for group b
+      ################################
+      
+      # Force of infection
+      lambda_b = trans_p*(c_mat[2,1]*(IUa/N_a + kappa*IPa/N_a) + c_mat[2,2]*(IUb/N_b + kappa*IPb/N_b) + c_mat[2,3]*(IUc/N_c + kappa*IPc/N_c) )
+      
+      
+      # change in contacts in response to deaths (constrained to maximum/minimum observed values)
+      dcb <-  (-zeta_b)*(cmax_b - cmin_b)*(deaths_t - deaths_tminusell)*exp(-zeta_b*(cum_deaths_ell))
+      
+      
+      dSUb <- -SUb*lambda_b -
+        pi_b * SUb + 
+        phi_b * SPb +
+        gamma * RUb -
+        vacc_b * SUb
+      
+      dSPb <- -SPb*kappa*lambda_b  +
+        pi_b * SUb - 
+        phi_b * SPb  +
+        gamma * RPb -
+        vacc_b * SPb
+      
+      dIUb <- SUb*lambda_b -
+        pi_b * IUb + 
+        phi_b * IPb -
+        (rho) * IUb 
+      
+      dIPb <- SPb*kappa*lambda_b +
+        pi_b * IUb - 
+        phi_b * IPb -
+        (rho) * IPb 
+      
+      dRUb <- (1-mu_b) * rho * IUb -
+        pi_b * RUb +
+        phi_b * RPb  -
+        gamma * RUb +
+        vacc_b * SUb
+      
+      dRPb <- (1-mu_b) * rho * IPb +
+        pi_b * RUb - 
+        phi_b * RPb  -
+        gamma * RPb +
+        vacc_b * SPb
+      
+      dDUb <- rho*mu_b*IUb
+      
+      dDPb <- rho*mu_b*IPb
+      
+      dCasesb <- SUb*lambda_b + SPb*kappa*lambda_b 
+      
+      ################################
+      ## transitions for group c 
+      ################################
+      
+      # Force of infection
+      lambda_c = trans_p*(c_mat[3,1]*(IUa/N_a + kappa*IPa/N_a) + c_mat[3,2]*(IUb/N_b + kappa*IPb/N_b) + c_mat[3,3]*(IUc/N_c + kappa*IPc/N_c) )
+      
+      
+      # change in contacts in response to deaths (constrained to maximum/minimum observed values)
+      dcc <-  (-zeta_c)*(cmax_c - cmin_c)*(deaths_t - deaths_tminusell)*exp(-zeta_c*(cum_deaths_ell))
+      
+      
+      dSUc <- -SUc*lambda_c -
+        pi_c * SUc + 
+        phi_c * SPc  +
+        gamma * RUc -
+        vacc_c * SUc
+      
+      dSPc <- -SPc*kappa*lambda_c  +
+        pi_c * SUc - 
+        phi_c * SPc  +
+        gamma * RPc -
+        vacc_c * SPc
+      
+      dIUc <- SUc*lambda_c -
+        pi_c * IUc + 
+        phi_c * IPc -
+        (rho) * IUc 
+      
+      dIPc <- SPc*kappa*lambda_c +
+        pi_c * IUc - 
+        phi_c * IPc -
+        (rho) * IPc 
+      
+      dRUc <- (1-mu_c) * rho * IUc -
+        pi_c * RUc +
+        phi_c * RPc  -
+        gamma * RUc +
+        vacc_c * SUc
+      
+      dRPc <- (1-mu_c) * rho * IPc +
+        pi_c * RUc - 
+        phi_c * RPc  -
+        gamma * RPc +
+        vacc_c * SPc
+      
+      dDUc <- rho*mu_c*IUc
+      
+      dDPc <- rho*mu_c*IPc
+      
+      dCasesc <- SUc*lambda_c + SPc*kappa*lambda_c 
+      
+      
+      
+      return(list(c(dca, dSUa, dSPa, dIUa, dIPa, dRUa, dRPa, dDUa, dDPa, dCasesa,
+                    #dIUaa, dIUab, dIUac,
+                    dcb, dSUb, dSPb, dIUb, dIPb, dRUb, dRPb, dDUb, dDPb, dCasesb,
+                    dcc, dSUc, dSPc, dIUc, dIPc, dRUc, dRPc, dDUc, dDPc, dCasesc
+      )))
+    })
+  }
+  
+  times<-seq(from=0, to=time, by=1)
+  
+  as.data.frame(dede(state, times, sir_up, params))->sim
+  
+  if(get_params){
+    return(list(sim=sim, params=params))
+  }
+  
+  else{
+    return(sim)
+  }
+  
+  
+  
+}
+
+function_in_string <- "You are in expert in these functions being run in a shiny app:
+
+avg_contact_matrix_3gp <- function(dbar_a, dbar_b, dbar_c,
+                                   N_a, N_b, N_c,
+                                   beta_a, beta_b) {
+  #h_a, h_b) {
+  
+  # fraction of total degree in group a and group b 
+  m_a <- (N_a * dbar_a) / ((N_a * dbar_a) + (N_b * dbar_b) + (N_c * dbar_c))
+  m_b <- (N_b * dbar_b) / ((N_a * dbar_a) + (N_b * dbar_b) + (N_c * dbar_c))
+  m_c <- (N_c * dbar_c) / ((N_a * dbar_a) + (N_b * dbar_b) + (N_c * dbar_c))
+  
+  q_aa <- m_a^(1/beta_a) 
+  q_bb <- m_b^(1/beta_b) 
+  
+  #q_aa <- h_a
+  #q_bb <- h_b
+  
+  
+  ## for a's contacts w/ other groups, assume they are distributed
+  ## in proportion to the other groups' share of edges
+
+  q_ab <- (1-q_aa) * (m_b / (m_b + m_c)) 
+  q_ac <- (1-q_aa) * (m_c / (m_b + m_c)) 
+  
+  # implied by q_ab and q_ac
+  q_ba <- q_ab*(dbar_a*N_a)/(N_b*dbar_b)
+  # and then q_bc has to make them add up to 1
+  q_bc <- (1 - (q_ba + q_bb))
+  
+  q_ca <- q_ac*dbar_a*N_a/(N_c*dbar_c)
+  
+  q_cb <- q_bc * (dbar_b*N_b)/(dbar_c*N_c)
+  
+  # q_cc is residual
+  q_cc <- 1 - (q_ca + q_cb)
+  
+  # implied beta_c
+  beta_c_implied <- log(m_c) / log(q_bb)
+  
+  # q matrix
+  q_beta <- t(matrix(c(q_aa, q_ab, q_ac, 
+                       q_ba, q_bb, q_bc,
+                       q_ca, q_cb, q_cc),
+                     nrow=3)) 
+  ## this matrix has the avg number of contacts between groups
+
+  c_mat <- q_beta
+  c_mat[1,] <- c_mat[1,] * dbar_a 
+  c_mat[2,] <- c_mat[2,] * dbar_b
+  c_mat[3,] <- c_mat[3,] * dbar_c
+  
+  return(c_mat)
+  
+}
+
+
+sir_three_group_pu <- function(## parameters related to popn
+  N0 = 10000000, # population size
+  frac_a = 0.33, # fraction of population in group A
+  frac_b = 0.33, # fraction of population in group B
+  ## parameters related to contacts
+  cmax = NA, # initial/max average number of contacts per day (specify only if its the same for both groups, otherwise NA)
+  cmax_a = 8, cmax_b = 8, cmax_c=8, # initial/max average number of contacts per day by group 
+  cmin = NA, # min average contact
+  cmin_a = 3, cmin_b = 3, cmin_c=3,
+  beta_a=1, # homophily parameter for group a (beta_a = 1 means unbiased mixing; bigger values mean homophily) 
+  beta_b=1, # homophily parameter for group b (beta_b = 1 means unbiased mixing; bigger values mean homophily) 
+  #h_a=.5, # proportion of group A's total contact with members of their own group 
+  #h_b=.5, # proportion of group A's total contact with members of their own group 
+  ##
+  zeta = NA, # responsiveness of contact to deaths
+  zeta_a = 0.01, zeta_b = 0, zeta_c = 0.005,
+  trans_p = 0.05, # probability of transmission given contact (or susceptibility to infection given contact)
+  rho=1/10,# 1 / infectious period  or recovery rate
+  mu = NA,# probability of dying following infection
+  mu_a = 0.01, mu_b = 0.01, mu_c = .01, # can let it differ between groups to crudely account for difference in age composition between groups
+  kappa=0.9, # scaling factor for probability of transmission given contact resulting from protective behavior; kappa=1 means no protection, kappa = 0 means perfect protection
+  phi = NA, # waning of protective behavior
+  phi_a = 0, phi_b=0, phi_c=0,
+  I0_a=1, I0_b=1, I0_c=1, #intial infected in each group
+  time = 500, # time steps for simulation
+  pi = NA, # background rate of adopting protective behavior
+  pi_a = 0.05, pi_b = 0.05, pi_c = 0.05,
+  ell = 1, # time window for considering deaths that influence adoption of protective behavior
+  vacc = NA, # vaccination rate (goes from S to R)
+  # in Roubenoff et al. we assume about 2 million daily doses are distributed which is ~ 0.6 % of the population per day
+  vacc_a = 0.006, vacc_b = 0.006, vacc_c = 0.006,
+  vstart = 365, # start of vaccination
+  gamma = 1/182.5, # wanning immunity
+  
+  get_params=FALSE) {
+  # set following parameters to be the same for both groups unless specified otherwise 
+  
+  if(!is.na(mu)){  
+    mu_a<-mu  
+    mu_b<-mu
+    mu_c<-mu
+  }
+  
+  if(!is.na(zeta)){  
+    zeta_a<-zeta  
+    zeta_b<-zeta
+    zeta_c<-zeta
+  }
+  
+  if(!is.na(pi)){  
+    pi_a<-pi 
+    pi_b<-pi
+    pi_c<-pi
+  }
+  
+  if(!is.na(phi)){  
+    phi_a<-phi 
+    phi_b<-phi
+    phi_c<-phi
+  } 
+  
+  
+  if(!is.na(vacc)){  
+    vacc_a<-vacc
+    vacc_b<-vacc
+    vacc_c<-vacc
+  }
+  
+  if(!is.na(cmin)){  
+    cmin_a<-cmin
+    cmin_b<-cmin
+    cmin_c<-cmin
+  }
+  
+  
+  if(!is.na(cmax)){  
+    cmax_a<-cmax
+    cmax_b<-cmax
+    cmax_c<-cmax
+  }
+  
+  N0_a = N0 * frac_a 
+  N0_b = N0 * frac_b
+  N0_c = N0 * (1 - frac_a - frac_b)
+  
+  
+  
+  params<-c(‘trans_p’ = trans_p, # probability of transmission given contact
+            ‘beta_a’= beta_a, # homophily parameter for group a (beta_a = 1 means unbiased mixing; beta_a > 1 means homophily)
+            ‘beta_b’= beta_b, # homophily parameter for group b (beta_b = 1 means unbiased mixing; beta_b > 1 means homophily)
+            #’h_a’= h_a, # # proportion of group A's total contact with members of their own group (bounded by population size and total contacts in group B)
+            #’h_b’= h_b, # # proportion of group A's total contact with members of their own group (bounded by population size and total contacts in group B)
+            ‘cmax_a’ = cmax_a, ‘cmax_b’ = cmax_b, ‘cmax_c’ = cmax_c, # upper and lower bounds for contacts
+            ‘zeta_a’ = zeta_a, ‘zeta_b’ = zeta_b, ‘zeta_c’ = zeta_c, # responsiveness of contact to deaths
+            ‘pi_a’ = pi_a, ‘pi_b’ = pi_b, ‘pi_c’ = pi_c, # background rate of adopting protective behavior
+            ‘kappa’ = kappa, # scaling factor for probability of transmission given contact resulting from protective behavior; kappa=1 means no protection, kappa = 0 means perfect protection
+            ‘phi_a’=phi_a, ‘phi_b’ = phi_b, ‘phi_c’ = phi_c, # waning rate of protective behavior
+            ‘rho’=rho, # 1 / infectious period
+            ‘mu_a’=mu_a, ‘mu_b’ = mu_b, ‘mu_c’ = mu_c, # probability of dying following infection
+            ‘time’=time, # time steps in simulation
+            ‘I0_a’=I0_a, ‘I0_b’=I0_b, ‘I0_c’=I0_c, # starting number infected in each group
+            ‘N0_a’ = N0_a, ‘N0_b’= N0_b, ‘N0_c’=N0_c, # population size in each group at time zero
+            ‘ell’=ell, # time window for considering deaths that influence adoption of protective behavior
+            ‘vacc_a’=vacc_a, ‘vacc_b’=vacc_b, ‘vacc_c’=vacc_c, # vaccination rate
+            ‘vstart’ = vstart, # start of vaccination
+            ‘gamma’ = gamma #waning immunity
+  ) # responsiveness to proportion of protected individuals for adopting protective behavior
+  
+  state<-c( 
+    
+    #group a
+    ca <- cmax_a,
+    SUa<-(N0_a - I0_a),
+    SPa<-0,
+    IUa<-I0_a,
+    IPa<-0,
+    RUa<-0,
+    RPa<-0,
+    DUa<-0,
+    DPa<-0,
+    
+    #group b
+    cb <- cmax_b,
+    SUb<-N0_b - I0_b,
+    SPb<-0,
+    IUb<-I0_b,
+    IPb<-0,
+    RUb<-0,
+    RPb<-0,
+    DUb<-0,
+    DPb<-0,
+    
+    #group c 
+    cc <- cmax_c,
+    SUc<-N0_c - I0_c,
+    SPc<-0,
+    IUc<-I0_c,
+    IPc<-0,
+    RUc<-0,
+    RPc<-0,
+    DUc<-0,
+    DPc<-0
+  )
+  
+  names(state)<- c(‘ca’,’SUa’, ‘SPa’, ‘IUa’, ‘IPa’, ‘RUa’, ‘RPa’, ‘DUa’, ‘DPa’,
+                   ‘cb’,’SUb’, ‘SPb’, ‘IUb’, ‘IPb’, ‘RUb’, ‘RPb’, ‘DUb’, ‘DPb’,
+                   ‘cc’,’SUc’, ‘SPc’, ‘IUc’, ‘IPc’, ‘RUc’, ‘RPc’, ‘DUc’, ‘DPc’)
+  
+  sir_up<-function(t, state, parameter){
+    
+    with(as.list(c(parameter, state)), {
+      
+      
+      if(t<ell){
+        lag_ell<-rep(0, length(state))
+      }
+      else{
+        lag_ell<-lagvalue(t-ell) #provide access to past (lagged) values of all state variables; below we make sure we only look at D
+      }
+      
+      
+      if(t<vstart){
+        vacc_a <- 0
+        vacc_b <- 0
+        vacc_c <- 0
+      }
+      else{
+        vacc_a <- vacc_a
+        vacc_b <- vacc_b
+        vacc_c <- vacc_c
+      }
+      
+      # for the transitions there are seven processes that are currently modeled:
+      # 1: transmission process (infection, recovery/death)
+      # 2: changes in contact rate for group A in response to deaths
+      # 3: adoption of protective behavior 
+      # 5: waning of protective behavior
+      # 6: vaccination
+      # 7: waning of immunity
+      
+      
+      
+      # recalculate population size (will change due to deaths)
+      N_a = SUa + IUa + RUa + SPa + IPa + RPa 
+      N_b = SUb + IUb + RUb + SPb + IPb + RPb
+      N_c = SUc + IUc + RUc + SPc + IPc + RPc
+      
+      
+      
+      # contact matrix
+      c_mat <- avg_contact_matrix_3gp(dbar_a = ca,
+                                      dbar_b = cb,
+                                      dbar_c = cc,
+                                      N_a = N_a,
+                                      N_b = N_b,
+                                      N_c = N_c,
+                                      beta_a = beta_a,
+                                      beta_b = beta_b)
+      #h_a = h_a,
+      #h_b = h_b)
+      
+      ## TODO - should we throw an error if any implied contacts are negative? 
+      
+      
+      ################################
+      ## transitions for group a
+      ################################
+      
+      # Force of infection
+      lambda_a =  trans_p*(c_mat[1,1]*(IUa/N_a + kappa*IPa/N_a) + c_mat[1,2]*(IUb/N_b + kappa*IPb/N_b) + c_mat[1,3]*(IUc/N_c + kappa*IPc/N_c))
+      
+      
+      # change in contacts in response to deaths occuring over a time window defined by ell (constrained to maximum/minimum observed values)
+      
+      deaths_t <- rho*mu_a*IUa + rho*mu_a*IPa + rho*mu_b*IUb + rho*mu_b*IPb + rho*mu_c*IUc + rho*mu_c*IPc
+      
+      
+      deaths_tminusell <- rho*mu_a*(lag_ell[which(names(state) == ‘IUa’)]) + 
+        rho*mu_a*(lag_ell[which(names(state) == ‘IPa’)]) + 
+        rho*mu_b*(lag_ell[which(names(state) == ‘IUb’)]) + 
+        rho*mu_b*(lag_ell[which(names(state) == ‘IPb’)]) +
+        rho*mu_c*(lag_ell[which(names(state) == ‘IUc’)]) + 
+        rho*mu_c*(lag_ell[which(names(state) == ‘IPc’)])
+      
+      cum_deaths_ell <- DUa - lag_ell[which(names(state) == ‘DUa’)] +
+        DPa - lag_ell[which(names(state) == ‘DPa’)] +
+        DUb - lag_ell[which(names(state) == ‘DUb’)] +
+        DPb - lag_ell[which(names(state) == ‘DPb’)] +
+        DUc - lag_ell[which(names(state) == ‘DUc’)] +
+        DPc - lag_ell[which(names(state) == ‘DPc’)]
+      
+      dca <-  (-zeta_a)*(cmax_a - cmin_a)*(deaths_t - deaths_tminusell)*exp(-zeta_a*(cum_deaths_ell))
+      
+      dSUa <- -SUa*lambda_a - #transmission process
+        pi_a * SUa + # adoption of protective behavior
+        phi_a * SPa + #waning of protective behavior
+        gamma * RUa - #waning immunity
+        vacc_a * SUa
+      
+      dSPa <- -SPa*kappa*lambda_a +
+        pi_a * SUa - 
+        phi_a * SPa  +
+        gamma * RPa -
+        vacc_a * SPa
+      
+      dIUa <- SUa*lambda_a  -
+        pi_a*IUa + 
+        phi_a * IPa -
+        (rho) * IUa 
+      
+      dIPa <- SPa*kappa*lambda_a +
+        pi_a * IUa - 
+        phi_a * IPa -
+        (rho) * IPa 
+      
+      dRUa <- rho * (1-mu_a) * IUa  -
+        pi_a * RUa + 
+        phi_a * RPa  - 
+        gamma * RUa +
+        vacc_a * SUa
+      
+      dRPa <- rho * (1-mu_a) * IPa  +
+        pi_a * RUa - 
+        phi_a * RPa  - 
+        gamma * RPa +
+        vacc_a * SPa
+      
+      dDUa <- rho*mu_a*IUa
+      
+      dDPa <- rho*mu_a*IPa
+      
+      
+      ################################
+      ## transitions for group b
+      ################################
+      
+      # Force of infection
+      lambda_b = trans_p*(c_mat[2,1]*(IUa/N_a + kappa*IPa/N_a) + c_mat[2,2]*(IUb/N_b + kappa*IPb/N_b) + c_mat[2,3]*(IUc/N_c + kappa*IPc/N_c) )
+      
+      
+      # change in contacts in response to deaths (constrained to maximum/minimum observed values)
+      dcb <-  (-zeta_b)*(cmax_b - cmin_b)*(deaths_t - deaths_tminusell)*exp(-zeta_b*(cum_deaths_ell))
+      
+      
+      dSUb <- -SUb*lambda_b -
+        pi_b * SUb + 
+        phi_b * SPb +
+        gamma * RUb -
+        vacc_b * SUb
+      
+      dSPb <- -SPb*kappa*lambda_b  +
+        pi_b * SUb - 
+        phi_b * SPb  +
+        gamma * RPb -
+        vacc_b * SPb
+      
+      dIUb <- SUb*lambda_b -
+        pi_b * IUb + 
+        phi_b * IPb -
+        (rho) * IUb 
+      
+      dIPb <- SPb*kappa*lambda_b +
+        pi_b * IUb - 
+        phi_b * IPb -
+        (rho) * IPb 
+      
+      dRUb <- (1-mu_b) * rho * IUb -
+        pi_b * RUb +
+        phi_b * RPb  -
+        gamma * RUb +
+        vacc_b * SUb
+      
+      dRPb <- (1-mu_b) * rho * IPb +
+        pi_b * RUb - 
+        phi_b * RPb  -
+        gamma * RPb +
+        vacc_b * SPb
+      
+      dDUb <- rho*mu_b*IUb
+      
+      dDPb <- rho*mu_b*IPb
+      
+      
+      ################################
+      ## transitions for group c 
+      ################################
+      
+      # Force of infection
+      lambda_c = trans_p*(c_mat[3,1]*(IUa/N_a + kappa*IPa/N_a) + c_mat[3,2]*(IUb/N_b + kappa*IPb/N_b) + c_mat[3,3]*(IUc/N_c + kappa*IPc/N_c) )
+      
+      
+      # change in contacts in response to deaths (constrained to maximum/minimum observed values)
+      dcc <-  (-zeta_c)*(cmax_c - cmin_c)*(deaths_t - deaths_tminusell)*exp(-zeta_c*(cum_deaths_ell))
+      
+      
+      dSUc <- -SUc*lambda_c -
+        pi_c * SUc + 
+        phi_c * SPc  +
+        gamma * RUc -
+        vacc_c * SUc
+      
+      dSPc <- -SPc*kappa*lambda_c  +
+        pi_c * SUc - 
+        phi_c * SPc  +
+        gamma * RPc -
+        vacc_c * SPc
+      
+      dIUc <- SUc*lambda_c -
+        pi_c * IUc + 
+        phi_c * IPc -
+        (rho) * IUc 
+      
+      dIPc <- SPc*kappa*lambda_c +
+        pi_c * IUc - 
+        phi_c * IPc -
+        (rho) * IPc 
+      
+      dRUc <- (1-mu_c) * rho * IUc -
+        pi_c * RUc +
+        phi_c * RPc  -
+        gamma * RUc +
+        vacc_c * SUc
+      
+      dRPc <- (1-mu_c) * rho * IPc +
+        pi_c * RUc - 
+        phi_c * RPc  -
+        gamma * RPc +
+        vacc_c * SPc
+      
+      dDUc <- rho*mu_c*IUc
+      
+      dDPc <- rho*mu_c*IPc
+      
+      
+      
+      
+      return(list(c(dca, dSUa, dSPa, dIUa, dIPa, dRUa, dRPa, dDUa, dDPa,
+                    dcb, dSUb, dSPb, dIUb, dIPb, dRUb, dRPb, dDUb, dDPb,
+                    dcc, dSUc, dSPc, dIUc, dIPc, dRUc, dRPc, dDUc, dDPc
+      )))
+    })
+  }
+  
+  times<-seq(from=0, to=time, by=1)
+  
+  as.data.frame(dede(state, times, sir_up, params))->sim
+  
+  if(get_params){
+    return(list(sim=sim, params=params))
+  }
+  
+  else{
+    return(sim)
+  }
+  
+  
+  
+}
+"
+
+function_in_string <- gsub("\\s+", " ", function_in_string)
+function_in_string <- trimws(function_in_string)