#####################
## Solving ODEs
## with lsoda in R
#####################
rm(list = ls())
require(odesolve)
#####################
## Logistic growth
#####################
ODEfun <- function(t, y, parms){ # `fun' = function, not fun ;)
     r <- parms[1] 
     K <- parms[2] 
     N <- y[1]
     dN <- r * N * (1 - N/K) 
     list(dN) 
} 

logistic <- lsoda(c(N=0.1), times= seq(0, 10, by= 0.1), func= ODEfun, parms= c(r= 0.9, K= 5))
head(logistic)
plot(logistic[,"time"], logistic[,"N"], ylab="N", xlab= "time", col= "navy")


#############################
### The SI model 
#############################
# assumes births & deaths equal and deaths from disease occur on timescale longer than we're interested in.
tInt = seq(0, 25, by = 1/2) 
pars = c(beta= 0.75) 
Initial = c(S = 4999, I = 1) 

SIfun = function(t, y, parms) { 
	S = y[1] 
	I = y[2]
	dS <- -parms[1] * S * (I/(S+I))
	dI <- parms[1] * S * (I/(S+I))
	ODEs = c(dS, dI) 
	list(ODEs)
} 

SIout <- lsoda(Initial, times= tInt, func= SIfun, parms= pars)
head(SIout)

par(mfrow = c(1,2)) 
plot(SIout[,"time"], SIout[,"S"], ylab = "Susceptible= blue", xlab = "time", col="navy")
points(SIout[,"time"], SIout[,"I"], col= "darkred")
plot(SIout[,"I"], SIout[,"S"], ylab = "Susceptible", xlab = "Infected") 


########################
### SIR model 
########################
tInt = seq(0, 150, by= 1/2) 
pars = c(beta= 0.9, gamma= 0.8) 
initial = c(S = 4999, I = 1, R = 0) 

SIRfun = function(t, y, parms) { 
	S = y[1] 
	I = y[2]
	R = y[3]
	dS <- -parms[1] * S * (I/(S+I))
	dI <- parms[1] * S * (I/(S+I)) - (parms[2] * I)
	dR <- parms[2] * I
	ODEs = c(dS, dI, dR) 
	list(ODEs)
} 

SIRout <- lsoda(initial, times= tInt, parms= pars, SIRfun)
head(SIRout)

par(mfrow = c(2, 2))
plot(SIRout[,"time"], SIRout[,"S"], ylab = "Susceptible", xlab = "time", col="navy")
plot(SIRout[,"time"], SIRout[,"I"], ylab= "Infected", xlab = "time", col= "darkred")
plot(SIRout[,"time"], SIRout[,"R"], ylab= "Recovered", xlab = "time", col= "darkgreen")
plot(SIRout[,"S"], SIRout[,"I"], ylab = "Infected", xlab = "Susceptible") 


######################
### Ottar's example 
### SEIR model; from EEID 2007
#################################
tInt = seq(0, 10, by = 1/120) 
parameters = c(mu = 1/50, N = 1, beta = 1000, sigma = 365/8, gamma = 365/5) 
Init = c(S = 0.06, E = 0, I = 0.001, R = 0.939) 

SEIRmod = function(t, y, parms) { 
	S = y[1] 
	E = y[2] 
	I = y[3] 
	R = y[4] 
	with(as.list(parms), { 
		dS = mu * (N - S) - beta * S * I/N 
		dE = beta * S * I/N - (mu + sigma) * E 
		dI = sigma * E - (mu + gamma) * I 
		dR = gamma * I - mu * R 
		ODEs = c(dS, dE, dI, dR) 
		list(ODEs) 
	}) 
} 

SEIRout <- lsoda(Init, times= tInt, parms= parameters, SEIRmod)
head(SEIRout)

par(mfrow = c(2, 2)) 
plot(SEIRout[,"time"], SEIRout[,"S"], ylab = "Susceptible", xlab = "time")
plot(SEIRout[,"time"], SEIRout[,"I"], ylab= "Infected", xlab = "time", col= 'darkred')
plot(SEIRout[,"time"], SEIRout[,"R"], ylab= "Recovered", xlab = "time", col= 'green') 
plot(SEIRout[,"S"], SEIRout[,"I"], ylab = "Infected", xlab = "Susceptible") 


##############################################
### Exercises Evolutionary response (a) & (b)
###############################################
tInt= 0:100
pars <- c(r= 0.2, aopt= 10, E= 0.03)

EvolFun = function(t, y, parms) { 
	r= parms[1]
	aopt= parms[2]
	E= parms[3]
	x = y[1] 
	a = y[2]
	dx <- x * (r - (a - aopt)^2)
	da <- E * (-2 * (a - aopt))
	ODEs <- c(dx, da)
	list(ODEs)
} 

Evol <- lsoda(c(x= 2, a= 9.15), times= tInt, parms= pars, EvolFun)
head(Evol)

par(mfrow=c(1,2))
plot(Evol[,"time"], Evol[,"x"], ylab = "Population Size", xlab = "time", col="navy")
plot(Evol[,"time"], Evol[,"a"], ylab= "Trait Value (mean activation temp.)", xlab = "time", col= "darkred")

