# Continuous-time Markov chain simulation via Gillespie's stochastic simulation algorithm (SSA)
# STAT 753
# April 7, 2020


##########################################################
# Simulate a CTMC - uses the 'GillespieSSA' package in R
##########################################################

# Install GillespieSSA package
library("GillespieSSA")

### Demo of examples -- for some reason, this is not working...
demo(GillespieSSA)

# Set parameters/arguments for ssa function
# x0: initial state vector
# a: vector of propensity functions (aka transition rate function)
# nu: state change vector (change in number of individuals in each state caused by one reaction)
# parms: vector of model parameters
# tf: final time


############################################################################
### Simple example: decay at a constant rate (e.g. radioactive decay model)
### Single reaction X -> 0 at rate c
############################################################################

## Large initial population size (X=1000 or 10,000)
parms <- c(c=0.5)
x0 <- c(X=1000)
a <- c("c*X")
nu <- matrix(-1)
out <- ssa(x0,a,nu,parms,tf=10,simName="Decay model") # SSA direct method

# Use built in plot function:
ssa.plot(out) # This doesn't divide by population size

# Plot output
plot(out$data[,1],out$data[,2]/1000,col="red",cex=0.5,pch=19,xlab="Time",ylab="Frequency",main="Decay Model")


## Smaller initial population size (X=100)
x0 <- c(X=100)
out <- ssa(x0,a,nu,parms,tf=10) # SSA direct method

# Plot: compare to large population size above, plot on same plot
points(out$data[,1],out$data[,2]/100,col="green",cex=0.5,pch=19)


## Small initial population size (X=10)
x0 <- c(X=10)
out <- ssa(x0,a,nu,parms,tf=10) # SSA direct method
points(out$data[,1],out$data[,2]/10,col="blue",cex=0.5,pch=19)


#################################
## Example: Logistic growth
#################################

# This is a birth and death process 
# N is the population size, K is the carrying capacity of the environment
# b = birth rate
# d = death rate


# Here, birth rate is twice the death rate
parms <- c(b=2, d=1, K=1000)
x0 <- c(N=500)
a <- c("b*N", "(d+(b-d)*N/K)*N")
nu <- matrix(c(+1,-1),ncol=2)
out1 <- ssa(x0,a,nu,parms,tf=10, method=ssa.d(), maxWallTime=5, simName="Logistic Growth")
ssa.plot(out1)


# Here, death rate is slightly higher than birth rate
parms <- c(b=1, d=1.2, K=1000)
x0 <- c(N=500)
a <- c("b*N", "(d+(b-d)*N/K)*N")
nu <- matrix(c(+1,-1),ncol=2)
out2 <- ssa(x0,a,nu,parms,tf=10, method=ssa.d(), maxWallTime=5, simName="Logistic Growth")
ssa.plot(out2)


# Here, birth rate = death rate
parms <- c(b=1, d=1, K=1000)
x0 <- c(N=500)
a <- c("b*N", "(d+(b-d)*N/K)*N")
nu <- matrix(c(+1,-1),ncol=2)
out3 <- ssa(x0,a,nu,parms,tf=10, method=ssa.d(), maxWallTime=5, simName="Logistic Growth")
ssa.plot(out3)



#########################################
## Example: Lotka predator-prey model -- need to look at relevant parameters, this isn't working properly
#########################################

parms <- c(c1=10, c2=.01, c3=10)
x0 <- c(Y1=1000,Y2=1000)
a <- c("c1*Y1","c2*Y1*Y2","c3*Y2")
nu <- matrix(c(+1,-1,0,0,+1,-1),nrow=2,byrow=TRUE)
out <- ssa(x0,a,nu,parms,tf=100, method=ssa.etl(), simName="Lotka predator-prey model")
ssa.plot(out)