# This is the R version of program 3.3 from page 79 of 
# "Modeling Infectious Disease in humans and animals" 
# by Keeling & Rohani.
 
# It is the SIR model with two different age-groups.
# Note we do not explicitly model the recovered classes.

# mu, S0 and I0 are all vectors.

# Set up equations as a function
SIS_age = function(time,state,pars){
  with(as.list(c(state,pars)),{
    nC=mu[2]/(lC+mu[2])
    nu=(lC+mu[2])*nC
    dSCdt = nu - (beta[1,1]*IC+beta[1,2]*IA)*SC - mu[1]*SC - lC*SC
    dICdt = (beta[1,1]*IC+beta[1,2]*IA)*SC - gamma*IC - mu[1]*IC - lC*IC
    dSAdt = lC*SC-(beta[2,1]*IC+beta[2,2]*IA)*SA - mu[2]*SA
    dIAdt = lC*IC+(beta[2,1]*IC+beta[2,2]*IA)*SA - gamma*IA - mu[2]*IA
    return(list(c(dSCdt,dICdt,dSAdt,dIAdt)))
  })
}

# Set parameters (all rates are per year)
beta=matrix(c(100, 10, 10, 20), 
            nrow = 2, ncol = 2, byrow = TRUE)    # Matrix of transmission rates
gamma=10                # Recovery rate
lC=0.066667             # Rate at which children mature and move into adult class
mu = c(0, 0.0166667)    # Death rate in childhood, adult group
pars = c(beta,gamma,lC,mu)

nC=mu[2]/(lC+mu[2])     # Proportion of population in childhood group
n=c(nC, 1-nC)
nu=(lC+mu[2])*nC        # Birth rate into childhood class

# Initial conditions
SC0 = 0.1             # Initial proportion that are susceptible and in childhood group
SA0 = 0.1             # Initial proportion that are susceptible and in adult group
IC0 = 0.0001          # Initial proportion that are infectious and in childhood group
IA0 = 0.0001          # Initial proportion that are infectious and in adult group
S0=c(SC0,SA0)
I0=c(IC0,IA0)
y0 = c(SC=SC0,IC=IC0,SA=SA0,IA=IA0)

# Time range
MaxTime=100
step=0.1
times = seq(0,MaxTime,step)

# Check all parameters are valid
for (i in 1:2) {
  if(I0[i]<0){
    print(paste("Initial level of infection in age group is less than zero, I0[",i,"] =",I0[i]))   
    stop()}
  if(S0[i]<0){
    print(paste("Initial level of susceptibles in age group is less than zero, S0[",i,"] =",S0[i]))   
    stop()}  
  if(mu[i]<0){
    print(paste("Death rate in age group is less than zero, mu[",i,"] =",mu[i]))   
    stop()}  
  if(I0[i]>n[i]){
    print(paste("Initial level of infection in age group greater than group size, I0[",i,"] =",I0[i], ", n[",i,"] =",n[i]))   
    stop()}  
  if(S0[i]>n[i]){
    print(paste("Initial level of susceptibles in age group is greater than group size, S0[",i,"] =",S0[i], ", n[",i,"] =",n[i]))   
    stop()}  
  if(S0[i]+I0[i]>n[i]){
    print(paste("Initial level of infecteds+susceptibles in age group is greater than group size S0+I0[",i,"] =",S0[i]+I0[i], ", n[",i,"] =",n[i]))   
    stop()}  
}
 
if(length(which(beta<0))>0){
  print(paste("Transmission rate is less than zero, beta =",beta[1,1],beta[1,2],beta[2,1],beta[2,2]))
  stop()} 
if(gamma<=0){
  print(paste("Recovery rate is less than or equal to zero, gamma =",gamma))  
  stop()}
if(MaxTime<=0){
  print(paste("Maximum run time is less than or equal to zero, MaxTime =",MaxTime))  
  stop()}
if(beta[1,2]!=beta[2,1]){
  print(paste("Transmission rate beta is not symmetric beta =",beta[1,1],beta[1,2],beta[2,1],beta[2,2]))  
  stop()}

# Solve differential equations
library(deSolve)
SIS_age.out = ode(y0,times,SIS_age,pars,rtol=1e-5)

# Plot results using ggplot2
library(ggplot2)
out = as.data.frame(SIS_age.out)

p1=ggplot() +
  geom_line(data=out,aes(y=IC,x= time,linetype="Y1"),
            color="red") +
  geom_line(data=out,aes(y=IA,x= time,linetype="Y2"),
            color="red") +
  xlab("Time") + ylab("Infectious") +
  scale_linetype_manual(name = "", labels=c("Child group","Adult group"),
                        values = c(1,2)) +  
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),limits=c(0,0.003)) +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.87,0.83),
        legend.box.background = element_rect(colour = "black"),
        legend.spacing.y = unit(0, "mm")) +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))

p2=ggplot() +
  geom_line(data=out,aes(y=SC,x= time,linetype="Y1"),
            color="green3") +
  geom_line(data=out,aes(y=SA,x= time,linetype="Y2"),
            color="green3") +
  xlab("Time") + ylab("Susceptibles") +
  scale_linetype_manual(name = "", labels=c("Child group","Adult group"),
                        values = c(1,2)) +  
  scale_y_continuous(expand=expansion(mult=c(0,0.05)),limits=c(0.08,0.2)) +
  scale_x_continuous(expand=expansion(mult=c(0,0))) +
  theme_classic() +
  theme(legend.title = element_blank(),legend.position = c(0.13,0.83),
        legend.box.background = element_rect(colour = "black"),
        legend.spacing.y = unit(0, "mm")) +
  theme(panel.border = element_rect(color = "black", 
                                    fill = NA, linewidth = 0.6))

library(patchwork)
p1/p2