#######################################
## Introduction 
#######################################
## Compartmental models and differential equation R tutorial
##  Part I/II
## E2M2 - 2020
## written by: Jessica Metcalf, Amy Wesolowski, Cara Brook
## modified by: Katie Gostic

## In this tutorial, we will solve a discrete-time model of simple population growth two ways:

## 1. Iteratively, using the change equation N_t+1 = N_t * lambda
## 2. Directly, using the state equation N_t = lambda^t * N_0

## We will also compare the discrete-time solution to the continuous time solution 
## In continuous time, we will only solve the equation one way, using the
## state equation: N(t) = N0*exp(rt)

##############################################################################
## Contents: 
# Modeling Population Growth
# Investigating continuous vs discrete time
# Structured population models
# Lotka-Volterra models of predator-prey dynamics
# Susceptible-Infected-Recovered models
##############################################################################


##############################################################################
## Modeling Population Growth
##############################################################################

## The simplest formulation for population growth in discrete time is N_t+1=lambda N_t
## where N_t is the size of the population at time t. The World Bank provides time-series 
## of population size for all countries in the world, including for Madagascar; this is in 
## the file "WorldBankPop.csv". Keep in mind that this is not data from direct observations 
## - censuses have not occurred in all the years for which data is available! 

## First, bring in this data by identifying the right path to the folder you are working in
## (use the command 'getwd' to find out where you are, and then use 'setwd' to get R to look 
## in the right place) and have a look at its structure in R (just the first 10 columns): 

##Load your data from the world bank and peek at the first few rows

pop.data <- read.csv("WorldBankPop.csv") ## read in the population data
head(pop.data) ## look at the first few rows of the data


## Make a sub-vector called 'mada.pop' that is the row of your dataset which 
## corresponds to Madagascar.
## Note that this vector contains a population estimate for each year from 1960 to 2015
mada.pop <- pop.data[pop.data$Country.Name=="Madagascar",]
mada.pop <- as.numeric(mada.pop[3:(length(mada.pop))])

## Make a second sub-vector called 'france.pop' that corresponds to France.
france.pop <- pop.data[pop.data$Country.Name=="France",]
france.pop <- as.numeric(france.pop[3:(length(france.pop))])

## Now, we can calculate the growth rate, N_t+1/N_t, by taking a vector that goes from the
## 2nd until the last value in our population vectors (this is years 1961 until 2016, and 
## thus we call it N_t+1). We will divide this by the the vector that goes from the first 
## until the before-last column, which will correspond to 1960 to 2015. Relative to the 
## first vector, this is N_t - i.e., the first number in the N_t+1 vector is 1961; and the 
## first number in the N_t vector is 1960, and so on.  

## First make a vector of Nt and Nt+1 for Mada
Nt.mada <- mada.pop[1:(length(mada.pop)-1)]
Nt1.mada <- mada.pop[2:length(mada.pop)]

## And do the same for France
Nt.france <- france.pop[1:(length(france.pop)-1)]
Nt1.france <- france.pop[2:length(france.pop)]

## Now estimate population growth rate for Madagascar and France as lambda = Nt+1/Nt
pop.growth.mada <- Nt1.mada/Nt.mada
pop.growth.france <- Nt1.france/Nt.france

## Is the population growing or declining?
## How do we interpret each entry in these vectors?

## Now plot (a) the population sizes for Madagascar and France from 1960-2016,
## side-by-side with the (b) population growth rates
par(mfrow=c(1,2))
plot(x=seq(1960,2016,1), y=mada.pop,type="l", xlab="",  #plot mada pop data
     ylab="Population size (World Bank estimate)", 
     ylim=range(c(mada.pop,france.pop),na.rm=TRUE))
points(x=seq(1960,2016,1), y=france.pop, type = "l", lty=2) #add france data
legend("topleft",legend=c("Madagascar","France"), lty=1:2, bty="n") #add legend

plot(x=seq(1960,2015,1), y=pop.growth.mada,type="l", xlab="",  #plot mada pop growth rate
     ylab="lambda = N(t+1)/N(t)", 
     ylim=range(c(pop.growth.mada,pop.growth.france),na.rm=TRUE))
points(x=seq(1960,2015,1), y=pop.growth.france, type = "l", lty=2) #add france


## The growth rates for Madagascar and France never drop below 1 (although France comes close).
## What does this tell us?


## On your own:
#### 1) Change the colour of the plots and the legend; 
#### 2) Experiment with plotting out different countries (Germany? Tanzania?)
#### 3) [Bonus] See if you can find which country in the entire data-base 
####    has the fastest growth, and when that occurred. You might want to use 
####    the 'which' command, with the argument 'arr.ind=TRUE' since this will 
####    index the matrix, i.e., tell you the row (which corresponds to the country) 
####    and the column (which corresponds to the year). Note that you will need 
####    to apply this to a matrix of growth rates!



##############################################################################
## Investigating continuous vs discrete time
##############################################################################

## Is the model above structured like a continuous or a discrete time model?

## Let's make a discrete time population model. We will use a for-loop to project a 
## population 100 time-steps, starting with population size of 10, and with a population 
## growth rate equal to 1.1

## Set up your (a) vector "N" of future population sizes, (b) fill in your initial population size,
## and (c) define "lambda", the population growth rate.
N <- rep(NA,100)
N[1] <- 10
lambda=1.1


## Now make a discrete time model using a for-loop to iterate this population 100 years
## into the future.

for(t in 2:100){
  N[t] <- lambda*N[t-1]
}


## Does anyone remember the equation to shorten this discrete time model? 
## If we do this, do we get the same result?
Ndirect = N[1]*lambda^(0:99) 

## Plot the results together.
par(mfrow=c(1,1))
plot(x=seq(1,100,1), y=N, pch=15, col="blue", ylab="N", xlab="time")
points(x=seq(1,100,1), y=Ndirect, pch=19, col="red", cex=0.5)
legend(x = 10, y = 100000, legend = c('Change equation: N(t+1)=N(t)*lambda', 'State equation: N(t)=N0*lambda^t'), col = c('blue', 'red'), pch = c(15, 19))

## The solutions are identical!


## ---------------------------------------------------------------------
##   CONTINUOUS TIME
## ---------------------------------------------------------------------

## What if we want to make a continuous time model? How do we write that?

## Recall that the equation for continuous-time growth of a simple population is:
##  N(t) = N0*exp(rt)
## Because we can solve this equation directly, the easiest way to model continuous time growth is to solve directly:
N0 = 10
r = log(lambda) 
t = 0:100
N_exact = N0*exp(r*t) 
N_exact


## Add the continuous time model to the plot
points(t, N_exact, col = 'orange', type = 'b')
legend(x = 10, 
       y = 100000, 
       legend = c('Discrete change equation: N(t+1)=N(t)*lambda', 'Discrete state equation: N(t)=N0*lambda^t', 'Continuous state equation: N(t)=N0*exp(rt)'), 
       col = c('blue', 'red', 'orange'), pch = c(15, 19, 1), lty = c(NA, NA, 1))


## The discrete and continous time models don't give solutions that are exactly the same.
## This is ok -- they are different models! But the solutions are close.



