### SEIR model
## You may already have your own version of this code from Workshop 1

## ----sir-diffeq, echo=FALSE-------------------------------------------------------------
library(deSolve)
seir_model = function (t, y, params)
{
  # label state variables from within the y vector
  # Now including an Exposed compartment
  S = y[1]        # susceptibles
  E = y[2]
  I = y[3]        # infectious
  R = y[4]        # recovered
  N = S + E + I + R
  
  # Pull parameter values from the params vector
  # Now including the sigma parameter, for the rate at which the exposed become infectious
  beta = params["beta"]
  sigma = params["sigma"]
  gamma = params["gamma"]
  
  
  # Define equations
  # now including dE, for rate of change in Exposed
  dS = -(beta*S*I)/N
  dE = (beta*S*I)/N - sigma*E
  dI = sigma*E - gamma*I
  dR = gamma * I
  
  # return list of gradients   
  # now including dE as an output
  
  res = c(dS, dE, dI, dR)
  list(res)
}

# This will produce output for 50 days
# in increments of 0.2 days
timepoints = seq (0, 50, by=0.2)



## ---------------------------------------------------------------------------------------
# Some fairly arbitrary values here
params_seir = c(
  beta = 1.25, 
  sigma = 1/2,
  gamma = 1/4
)



## ---------------------------------------------------------------------------------------
# Now including an initial value for the E compartment
initial_seir = c (S = 999, E=0, I = 1, R = 0)


## ---------------------------------------------------------------------------------------
output_seir = ode (y=initial_seir, times = timepoints, func = seir_model, parms= params_seir)
outdf_seir = as.data.frame(output_seir)

# display the first 6 rows
head(outdf_seir)

# This function doesn't need to be re-defined
# we just include 'E' in compartments
plot_sir = function(
  sir_df, 
  Compartments = c("S", "I", "R")
){
  # Collect the values in columns S, I and R (or those given as "Compartments")
  #  into one column named 'value', and create an additional column named 
  # "Compartment".
  output_tidy = pivot_longer(sir_df, cols = all_of(Compartments), names_to = "Compartment")
  # Plot the values, with a different coloured line for each compartment
  ggplot(data=output_tidy, aes(x=time, y=value, col=Compartment)) + 
    geom_line() + ylab("People")  
}

plot_sir(outdf_seir, Compartments = c("S", "E", "I", "R"))