R/SIRsimulation.R

In fintzij/augSIR: Estimation of stochastic SIR epidemic parameters via Bayesian data augmentation.

# SIRsim simulates an SIR epidemic

infec_rate <- function(a, b, St, It){
  (b*It + a)*St
}

recov_rate <- function(mu, It){
  mu*It
}

SIRsim <- function(popsize, initdist, b, mu, a=0, tmax, censusInterval, sampprob, trim = TRUE, returnX = FALSE) {
    
    # create a matrix to store the individual level trajectories. The first column is the time of infection or recovery, the second column contains subject ids, the 
    # third column records 1 for an infection, -1 for recovery, 0 for no event. The matrix is ordered according to time, then event (only relevant for time=0). 
    
    if(returnX == TRUE){
        X <- as.matrix(data.frame(time=rep(0,popsize*2), id=rep(1:popsize,each=2), event=rep(0,2*popsize)))
        
    }
    
    # create matrix to store the simulation results. Observed will contain the binomial samples, time is the time of recovery or infection, truth is the number of infecteds
    # at each event time. 
    
    SIRres <- data.frame(time = seq(0, tmax, by = censusInterval),
                         Observed = 0, 
                         Truth = 0)
    
    # sample initial number of infected and susceptible individuals according to multinomial distribution
    initcounts <- rmultinom(1, size = popsize, prob = initdist)
    susceptiblenow <- initcounts[1]; infectednow <- initcounts[2]
    
    # if the initial number of infecteds is zero, redraw since there is no need to consider an impossible epidemic
    if(infectednow == 0){
        while(infectednow == 0){
            initcounts <- rmultinom(1, size = popsize, prob = initdist)
            susceptiblenow <- initcounts[1]; infectednow <- initcounts[2]       
        }
    }
    
    # record the initial number of infecteds
    SIRres$Truth <- infectednow
    
    # if returnX == TRUE, initialize the vectors of IDs for infecteds and susceptibles
    if(returnX == TRUE){
        X[seq(1,(2*infectednow - 1),by=2), 3] <- 1 #record an infection, infection time for these cases is zero
        
        # initialize vector of susceptible IDs
        if(infectednow != popsize){
            which.susc <- (infectednow+1):popsize
        } else {
            which.susc <- 0
        }
        
        # initialize vector of IDs for infecteds
        which.inf <- 1:infectednow
        
    }
        
    # infection and recovery rates
    infec.rate <- infec_rate(a=a, b=b, St=susceptiblenow, It=infectednow) 
    recov.rate <- recov_rate(mu=mu, It=infectednow)
    
    keep.going <- TRUE; timenow <- 0
    
    
    while(keep.going == TRUE){
        
        # the hazard is the sum of the rate of a new infection and the rate or a new recovery
        hazard <- infec.rate + recov.rate
        
        # sample the time of the next infection or recovery
        tau <- rexp(1, rate = hazard)
                
        # conditional on an event happening, sample whether it is an infection or a recovery
        event <- sample.int(n = 2, size = 1, prob = c(infec.rate, recov.rate))
        
        
        if (event == 1) { #infection happens
            
            # update time
            timenow <- timenow + tau
            
            # update the current time, count of infecteds (by adding 1), and count of susceptibles (by subtracting 1)
            infectednow <- infectednow + 1; susceptiblenow <- susceptiblenow - 1
            
            # update results matrix
            SIRres$Truth[SIRres$time >= timenow] <- infectednow
            
            # update the infection and recovery rates
            infec.rate <- infec_rate(a=a, b=b, St=susceptiblenow, It=infectednow); recov.rate <- recov_rate(mu=mu, It=infectednow)
            
            # if returnX == TRUE, update X
            if(returnX == TRUE){
                # assign the infection to the next susceptible individual
                X[which(X[,2] == which.susc[1])[1],1] <- timenow # record the infection time
                X[which(X[,2] == which.susc[1])[1],3] <- 1 # record the infection code
                
                # remove the infected individual from the list of susceptibles and add him to the list of infecteds
                which.inf <- c(which.inf, which.susc[1]); which.susc <- which.susc[-1]
            }
            
            
        } else if(event == 2){ # recovery happens
            
            # update time
            timenow <- timenow + tau
            
            # update time and count of infecteds (subtracting 1). No change to the number of susceptibles.
            infectednow <- infectednow - 1
            
            # update results matrix
            SIRres$Truth[SIRres$time >= timenow] <- infectednow
                        
            # update infection and recovery rates
            infec.rate <- infec_rate(a=a, b=b, St=susceptiblenow, It=infectednow); recov.rate <- recov_rate(mu=mu, It=infectednow)
            
            # if returnX == TRUE, select an infected individual to recover and update the vector of infecteds
            if(returnX == TRUE){
                # choose an individual to recover (discrete uniform over the infected individuals, by the memoryless property of exponentials and the fact that all 
                # individuals recover at the same rate)
                whorecovers <- ifelse(length(which.inf) > 1, sample(which.inf, 1), which.inf)
                
                # record the time and recovery status for the infected individual
                X[which(X[,2] == whorecovers)[2],1] <- timenow #record the recovery time 
                X[which(X[,2] == whorecovers)[2],3] <- -1 #record the recovery if it occurs before tmax
                
                #remove recovered individual from list of infecteds
                which.inf <- which.inf[which.inf != whorecovers]
            }            

        }
        
        # no more infected individuals so the epidemic has been fully observed and we stop
        if(infectednow == 0) keep.going <- FALSE
    }
    
    # if returnX == TRUE, we may need to censor some observations in the X matrix
    if(returnX == TRUE){
        # censor individuals outside of observation window
        X[X[,1] > tmax, c(1,3)] <- 0
        X <- X[order(X[,1],X[,3]),]
    }
    
    
    #generate binomial samples
    SIRres$Observed <- rbinom(n = nrow(SIRres), size = SIRres$Truth, prob = sampprob)
        
    # If first observed count is zero, resample that observation
#     if(SIRres$Observed[1] == 0){
#         while(SIRres$Observed[1] == 0){
#             SIRres$Observed[1] <- rbinom(n = nrow(SIRres), size = SIRres$Truth[1], prob = sampprob)
#         }
#     }
        

    # Get rid of the matrix after the infection has died out (no more infecteds) if trim == TRUE
    if(trim == TRUE){
        if(any(SIRres[,3] == 0)){
            ind <- which(SIRres[,3] == 0)[1]
            SIRres <- SIRres[1:ind,]
        }
    }

    
    if(returnX == FALSE){
        return(SIRres)
    } else {
        return(list(results = SIRres, trajectory = X))
    }
}

# sim_one_SIR simulates a single SIR trajectory for one individual
sim_one_SIR <- function(Xcount, obstimes, b, m, initdist, tmax, returnpath = FALSE){
    Xt <- rep(1, length(obstimes))
    eventtimes <- c(Xcount[,1], tmax); numsick <- Xcount[,2]
    
    initstate <- sample.int(n=3, size=1, prob=initdist)
    
    if(initstate == 2){
        tau <- rexp(1, m)
        path <- c(0,tau)
        Xt[obstimes <= tau] <- 2
        Xt[obstimes > tau] <- 3
        
    } else if(initstate == 1){
        path <- c(0,0)
        cur.time <- 0; ind <- 1
        keep.going <- TRUE
        rate <- b*numsick[ind]
        
        while(keep.going == TRUE){
            
            if(cur.time < tmax){
                
                tau <- rexp(1, rate)
                if((cur.time + tau) < eventtimes[ind+1]){
                    path[1] <- cur.time + tau
                    cur.time <- cur.time + tau
                    keep.going <- FALSE
                    
                    path[2] <- cur.time + rexp(1, m)
                    
                    Xt[obstimes <= path[1]] <- 1
                    Xt[(obstimes > path[1]) & (obstimes <= path[2])] <- 2
                    Xt[obstimes > path[2]] <- 3
                    
                    
                } else if((cur.time + tau) >= eventtimes[ind + 1]){
                    cur.time <- eventtimes[ind + 1]
                    ind <- ind+1
                    rate <- b*numsick[ind]
                    
                    if(numsick[ind] == 0){
                        keep.going <- FALSE
                        path <- c(0,0)
                        Xt[1:length(Xt)] <- 1
                    }
                }
                
            } else if(cur.time >= tmax){
                keep.going <- FALSE
                
                path <- c(0,0)
                
                Xt[1:length(Xt)] <- 1
                
            }            
            
        }
        
    }
    
    if(returnpath == FALSE) {
        return(Xt)
    } else{
        return(path)
    }
}


# Gillespie on extended state space ---------------------------------------

SIRsim_extended <- function(popsize, initdist, b, m, obstimes, samp_prob, trim = TRUE, returnX = TRUE){
    
    X <- as.matrix(data.frame(time=rep(0,popsize*2), id=rep(1:popsize,each=2), event=rep(0,2*popsize)))
    
    which_initinfec <- which(runif(popsize) < initdist[2])
    # someone must be initially infected. 
    if(length(which_initinfec) == 0) {
        while(length(which_initinfec) == 0){
            which_initinfec <- which(runif(popsize) < initdist[2])
        }
    }
    
    whichsusc <- 1:popsize; whichinf <- which_initinfec
    whichsusc <- whichsusc[-which_initinfec]
    
    recov_times <- rep(0, length(whichinf))
    
    for(k in 1:length(which_initinfec)){
        subj_id = which_initinfec[k]
        X[X[,"id"] == subj_id, "event"][1] <- 1
        X[X[,"id"] == subj_id, "event"][2] <- -1
        X[X[,"id"] == subj_id, "time"][2] <- recov_times[k] <- rexp(1, rate = m)
    }
    
    time <- 0; tmax = max(obstimes)
    
    while(length(whichinf) != 0){
        
        if(length(whichsusc) == 0){
            infec_times <- Inf
        } else {
            infec_times <- time + rexp(length(whichsusc), rate = b * length(whichinf))
        } 

        if(min(infec_times) < min(recov_times)){ #infection occurs
            time <- min(infec_times)
            subj <- whichsusc[which.min(infec_times)]
            X[X[,"id"] == subj, "time"][1] <- time
            X[X[,"id"] == subj, "event"][1] <- 1
            
            # draw recovery
            X[X[,"id"] == subj, "time"][2] <- time + rexp(1, rate = m)
            X[X[,"id"] == subj, "event"][2] <- -1
            
            recov_times <- c(recov_times, X[X[,"id"] == subj, "time"][2])
            
            whichinf <- c(whichinf, subj)
            whichsusc <- whichsusc[whichsusc != subj]
            
        } else { #recovery occurs
            time <- min(recov_times)
            whichinf <- whichinf[-which.min(recov_times)]
            
            recov_times <- recov_times[recov_times > time]
        }
        
    }
    
    X[X[,"time"] > tmax, c("time","event")] <- 0
    X <- X[order(X[,"time"],X[,"event"]),]  
    
    Xcount <- build_countmat(X, popsize)
    
    # build observation matrix
    SIRres <- data.frame(time = obstimes,
                         Observed = 0, 
                         Truth = 0)
    
    # get true counts of infecteds
    for(j in 1:length(obstimes)){
        SIRres[j,"Truth"] <- Xcount[sum(Xcount[,"time"] <= obstimes[j]), "numsick"]
    }
    
    # get binomial counts of infecteds
    SIRres[,"Observed"] <- rbinom(length(obstimes), SIRres[,"Truth"], samp_prob)
    
    if(trim == TRUE){
        if(any(SIRres[,3] == 0)){
            ind <- which(SIRres[,3] == 0)[1]
            SIRres <- SIRres[1:ind,]
        }
    }
    
    if(returnX == FALSE){
        return(SIRres)
    } else {
        return(list(results = SIRres, trajectory = X))
    }
}

# sim_one_tpms simulates infection statuses at a sequence of observation times directly from transition probability matrices
sim_one_tpms <- function(tpms, initdist) {
    states <- rep(0, dim(tpms)[3] + 1)
    
    states[1] <- sample.int(n = 3, size = 1, prob = initdist)
    
    for(s in 1:dim(tpms)[3]) {
        states[s+1] <- sample.int(n = 3, size = 1, prob = tpms[states[s],,s])
        
    }
    
    return(states)
}