# Simulate a Poisson distribution of new infections over time
#	code written by Mark Handcock and Martina Morris

# Set the seed so we can reproduce it
set.seed(1)
#
n <- 70         # the number of time steps
delta.t <- 0.01 # step size. Total time on n*delta.t
lambda <- 5     # rate per unit time
p <- rep(0,n)   # to store the prevalance
i <- rep(0,n)   # to store the incidence
p[1] <- 1       # initial prevalence
#
# set up a plot
plot(x=delta.t*(1:n),y=1:n, ylim=c(1,20), type="n", xlab="time", ylab="number of infections")
# step through time
for(k in 1:(n-1)){
	i[k] <- rpois(n=1,lambda=lambda*p[k]*delta.t)
	p[k+1] <- p[k]+i[k]
#if(p[k+1] < 15){ aaa <- locator(1) }
locator(1) # comment out to eliminate pause after each step; otherwise click on plot to continue
 points(x=delta.t*k,y=p[k+1],pch=19,cex=0.9,col=2)
}
#lines(x=1:n,y=p)
locator(2) # comment out to eliminate pause after each step; otherwise click on plot to continue
# plot the exponential deterministic model
lines(x=delta.t*(1:n), y=exp(lambda*delta.t*(0:(n-1))), col=2)
text(x=0.17,y=5,"exponential deterministic mean", col=2, cex=.8)
#
# Set up  simulation
locator(1)  # comment out to eliminate pause after each step; otherwise click on plot to continue
nsim <- 10     # The number of simulations
mp <- rep(0,n)  # to store the mean prevalence at each time
plot(x=delta.t*(1:n),y=1:n, ylim=c(1,20), type="n", xlab="time", ylab="number of infections")
for(j in 1:nsim){
	for(k in 1:(n-1)){
		i[k] <- rpois(n=1,lambda=lambda*p[k]*delta.t)
		p[k+1] <- p[k]+i[k]
	}
	mp <- mp + p / nsim
	lines(x=delta.t*(1:n),y=p,col=5, lwd=2)
	locator(1)
}
locator(1)

lines(x=delta.t*(1:n), y=mp, lty=2, lwd=2, col=3) # plot the empirical mean
text(x=0.4,y=2,"empirical mean", cex=.8)
locator(1) # comment out to eliminate pause after each step; otherwise click on plot to continue

lines(x=delta.t*(1:n), y=exp(lambda*delta.t*(0:(n-1))), lwd=2, col=2) # plot the deterministic mean
text(x=0.17,y=5,"deterministic mean", col=2, cex=.8)


#  repeat the next section to see the stochastic nature of the empirical mean

nsim <- 100     # The number of simulations
mp <- rep(0,n)  # to store the mean prevalence at each time
plot(x=delta.t*(1:n),y=1:n, ylim=c(1,20), type="n", xlab="time", ylab="number of infections")
for(j in 1:nsim){
 for(k in 1:(n-1)){
  i[k] <- rpois(n=1,lambda=lambda*p[k]*delta.t)
  p[k+1] <- p[k]+i[k]
 }
mp <- mp + p / nsim
lines(x=delta.t*(1:n),y=p,col=5)
#locator(1)
#points(x=1:n,y=p,pch=19,cex=0.2)
}
#locator(1)
lines(x=delta.t*(1:n), y=mp, lty=2,col=3) # plot the empirical mean
text(x=0.4,y=2,"empirical mean", cex=.8)
#locator(1)
lines(x=delta.t*(1:n), y=exp(lambda*delta.t*(0:(n-1))),col=2) # plot the deterministic mean
text(x=0.17,y=5,"deterministic mean", col=2, cex=.8)