hetero_age_2022_07_10.Rmd

---
title: "Heterogeneity and Age Structure in SIR Models"
author: "Matthew Ferrari (adapted from Helen Wearing and Aaron King)"
date: '2022-06-22'
output:
  html_document: default
  pdf_document: default
editor_options:
  markdown:
    wrap: sentence
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
```


## A Model With 2 Classes
We'll start with the simplest mechanistic model of two classes we can think of, which has separate classes for two groups $a$ and $b$. These groups could represent different socioeconomic classes, for example.

```{r}
require(diagram)
elpos <- rbind(
               Sa=c(2,3),
               Sb=c(3,3),
               Ia=c(2,2),
               Ib=c(3,2),
               Ra=c(2,1),
               Rb=c(3,1)
               )
elpos[,1] <- (2*elpos[,1]-1)/8
elpos[,2] <- (2*elpos[,2]-1)/6

fromto <- rbind(
                SaIa=c(1,3),
                SbIb=c(2,4),
                IaRa=c(3,5),
                IbRb=c(4,6)
                )

op <- par(mar=c(1,1,1,1))
openplotmat(asp=0.9)
arrpos <- matrix(ncol=2,nrow=nrow(fromto))
for (i in seq_len(nrow(fromto))){
  arrpos[i,] <- straightarrow(
                              to=elpos[fromto[i,2],],
                              from=elpos[fromto[i,1],],
                              lwd=2,
                              arr.pos=0.65,
                              arr.length=0.5
                              )}
textrect(elpos["Sa",],0.07,0.07,lab="Sa",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Sb",],0.07,0.07,lab="Sb",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Ia",],0.07,0.07,lab="Ia",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Ib",],0.07,0.07,lab="Ib",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Ra",],0.07,0.07,lab="Ra",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Rb",],0.07,0.07,lab="Rb",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
text(arrpos[1,1]-0.05,2/3,quote(lambda[a]),cex=2)
text(arrpos[2,1]+0.05,2/3,quote(lambda[b]),cex=2)
text(arrpos[3,1]-0.05,1/3,quote(gamma[a]),cex=2)
text(arrpos[4,1]+0.05,1/3,quote(gamma[b]),cex=2)

par(op)
```

Which can be written in equations as,
$$
\begin{aligned}
    \frac{dS_a}{dt} &= -\lambda_a\,S_a \phantom{-\gamma\,I_b}\\
    \frac{dS_b}{dt} &= -\lambda_b\,S_b \phantom{-\gamma\,I_b}\\
    \frac{dI_a}{dt} &= \phantom{-}\lambda_a\,S_a -\gamma\,I_a\\
    \frac{dI_b}{dt} &= \phantom{-}\lambda_b\,S_b-\gamma\,I_b\\
    \frac{dR_a}{dt} &= \phantom{-\lambda_a\,S_b}+\gamma\,I_a\\
    \frac{dR_b}{dt} &= \phantom{-\lambda_a\,S_b}+\gamma\,I_b\\
  \end{aligned}
$$
The $\lambda$s denote the group-specific force of infections:
$$
\begin{aligned}
        \lambda_a &= \beta_{aa}\,I_a+\beta_{ab}\,I_b\\
        \lambda_b &= \beta_{ba}\,I_a+\beta_{bb}\,I_b
\end{aligned}
$$
In this model, each population can infect each other but the infection moves through the populations separately.
Let's simulate such a model.
To make things concrete, we'll assume that the transmission rates $\beta$ are greater within groups than between them.

```{r}
b1 <- 0.005
b2 <- 0.02
gamma <- 10
```

```{r}
# Here we set up the ODE model that matches the equations above
ba.model <- function (t, x, ...) {     
  s <- x[c("Sb","Sa")]                  # susceptibles
  i <- x[c("Ib","Ia")]                  # infecteds
  r <- x[c("Rb","Ra")]                  # recovereds
  n <- s+i+r                            # total pop
  lambda.b <- (b1+b2)*i[1]+b1*i[2]      # group B force of infection
  lambda.a <- b1*i[1]+(b1+b2)*i[2]      # group A force of infection
  list(
       c(                               # these are the rates from the 6 equations above
         -lambda.b*s[1],
         -lambda.a*s[2],
          lambda.b*s[1]-gamma*i[1],
          lambda.a*s[2]-gamma*i[2],
                        gamma*i[1],
                        gamma*i[2]
         )
       )
}
```

```{r}
require(deSolve)

## initial conditions
yinit <- c(Sb=2000,Sa=1000,Ib=1,Ia=1,Rb=0,Ra=0) # set starting conditions

sol_ba <- ode(                          # run the ode solver on the model above
           y=yinit,
           times=seq(0,2,by=0.001),
           func=ba.model
           )
```

```{r}
# this code will plot the resulting time series
plot(sol_ba, mfcol=c(2,3))
dim(sol_ba)
head(sol_ba)
plot(sol_ba,log='y', mfcol=c(2,3))
```

```{r}
plot(sol_ba, mfcol=c(2, 3))
par(mfcol=c(1,1))
```

The results of the above are plotted below:

```{r}
# plot the proportion of individuals in each state for the two groups
time <- sol_ba[,1]                         # time
y <- sol_ba[,-1]                           # all other variables
n <- apply(y,1,sum)                     # population size
prop <- y/n                             # fractions
subsampled.prop <- prop[seq(1,length(time),by=10),]
subsampled.time <- time[seq(1,length(time),by=10)]
barplot(
        t(subsampled.prop),
        names.arg=subsampled.time,
        xlab='time',main='Group structure',
        space=0,
        col=c(
          rgb(0.5,1,0.5),
          rgb(0,1,0),
          rgb(1,0.5,0.5),
          rgb(1,0,0),
          rgb(0.5,0.5,1),
          rgb(0,0,1)
          ),
        legend = colnames(prop),
        args.legend=list(bg="white")
        )
```

## A Model With 2 Age Classes

Note that age is a special kind of heterogeneity in an epidemic model because individuals necessarily move from one class (younger) to another class (older) in a directional fashion that is independent of the infection and recovery process.

<!-- Today's first lecture showed how force of infection can vary with age. -->

<!-- What sort of mechanisms might give rise to these effects? -->

<!-- Here we'll see to what extent we can infer these mechanisms on the basis of age-specific incidence and seroprevalence data. -->

We'll start by introducing age into the model above.
So now $a$ becomes juveniles and $b$ becomes adults.
And, independent of the disease process, juveniles (of any category) age into adults.
Additionally, new juveniles are added through births (always first susceptible) and old individuals are lost to death.

```{r}
require(diagram)
elpos <- rbind(
               B=c(1,3),
               Sj=c(2,3),
               Sa=c(3,3),
               Ij=c(2,2),
               Ia=c(3,2),
               Rj=c(2,1),
               Ra=c(3,1),
               Ds=c(4,3),
               Di=c(4,2),
               Dr=c(4,1)
               )
elpos[,1] <- (2*elpos[,1]-1)/8
elpos[,2] <- (2*elpos[,2]-1)/6

fromto <- rbind(
                BSj=c(1,2),
                RaD=c(7,10),
                SjSa=c(2,3),
                IjIa=c(4,5),
                RjRa=c(6,7),
                SjIj=c(2,4),
                SaIa=c(3,5),
                IjRj=c(4,6),
                IaRa=c(5,7),
                SaD=c(3,8),
                IaD=c(5,9)  
                )

op <- par(mar=c(1,1,1,1))
openplotmat(asp=0.9)
arrpos <- matrix(ncol=2,nrow=nrow(fromto))
for (i in seq_len(nrow(fromto)))
  arrpos[i,] <- straightarrow(
                              to=elpos[fromto[i,2],],
                              from=elpos[fromto[i,1],],
                              lwd=2,
                              arr.pos=0.65,
                              arr.length=0.5
                              )
textrect(elpos["B",],0.07,0.07,lab="B",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Sj",],0.07,0.07,lab="Sj",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Sa",],0.07,0.07,lab="Sa",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Ij",],0.07,0.07,lab="Ij",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Ia",],0.07,0.07,lab="Ia",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Rj",],0.07,0.07,lab="Rj",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Ra",],0.07,0.07,lab="Ra",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Ds",],0.07,0.07,lab="D",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Di",],0.07,0.07,lab="D",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
textrect(elpos["Dr",],0.07,0.07,lab="D",box.col=gray(0.7),shadow.col=gray(0.4),shadow.size=0.01,cex=2)
text(1/2,arrpos[3,2]+0.05,quote(alpha),cex=2)
text(1/2,arrpos[4,2]+0.05,quote(alpha),cex=2)
text(1/2,arrpos[5,2]+0.05,quote(alpha),cex=2)
text(arrpos[6,1]-0.05,2/3,quote(lambda[J]),cex=2)
text(arrpos[7,1]+0.05,2/3,quote(lambda[A]),cex=2)
text(arrpos[8,1]-0.05,1/3,quote(gamma),cex=2)
text(arrpos[9,1]+0.05,1/3,quote(gamma),cex=2)


##text(3/4,arrpos[2,2]+0.05,quote(mu),cex=2)
par(op)
```

<!-- $$ -->

<!-- \begin{aligned} -->

<!--   \frac{dS_J}{dt} &= -\lambda_J\,S_J \phantom{-\gamma\,I_A}\\ -->

<!--   \frac{dS_A}{dt} &= -\lambda_A\,S_A \phantom{-\gamma\,I_A}\\ -->

<!--   \frac{dI_J}{dt} &= \phantom{-}\lambda_J\,S_J -\gamma\,I_J\\ -->

<!--   \frac{dI_A}{dt} &= \phantom{-}\lambda_A\,S_A-\gamma\,I_A\\ -->

<!--   \frac{dR_J}{dt} &= \phantom{-\lambda_J\,S_A}+\gamma\,I_J\\ -->

<!--   \frac{dR_A}{dt} &= \phantom{-\lambda_J\,S_A}+\gamma\,I_A\\ -->

<!-- \end{aligned} -->

<!-- $$ The $\lambda$s denote the age-specific force of infections: $$ -->

<!-- \begin{aligned} -->

<!--     \lambda_J &= \beta_{JJ}\,I_J+\beta_{JA}\,I_A\\ -->

<!--     \lambda_A &= \beta_{AJ}\,I_J+\beta_{AA}\,I_A -->

<!--   \end{aligned} -->

<!-- $$  -->

We can do this very simply using the same ingredients that go into the basic SIR model.
In that model, the waiting times in the S and I classes are exponential.
Let's assume the same thing about the aging process.
We'll also add in births into the juvenile susceptible class and deaths from the adult classes.

$$
  \begin{aligned}
    \frac{dS_J}{dt} &= -\lambda_J\,S_J \phantom{-\gamma\,I_A}  +B \phantom{-\mu\,S_A}-\alpha\,S_J\\
    \frac{dS_A}{dt} &= -\lambda_A\,S_A \phantom{-\gamma\,I_A +B} -\mu\,S_A+\alpha\,S_J\\
    \frac{dI_J}{dt} &= \phantom{-}\lambda_J\,S_J -\gamma\,I_J\phantom{+B-\mu\,S_A}-\alpha\,I_J\\
    \frac{dI_A}{dt} &= \phantom{-}\lambda_A\,S_A-\gamma\,I_A \phantom{+B} -\mu\,I_A+\alpha\,I_J\\
    \frac{dR_J}{dt} &= \phantom{-\lambda_J\,S_A}+\gamma\,I_J\phantom{+B-\mu\,S_A}-\alpha\,R_J\\
    \frac{dR_A}{dt} &= \phantom{-\lambda_J\,S_A}+\gamma\,I_A \phantom{+B}-\mu\,R_A+\alpha\,R_J\\
  \end{aligned}
$$

Now, let's simulate this model, under the same assumptions about transmission rates as above.

```{r}
# define the parameters
b1_demog <- 0.002
b2_demog <- 0.002
gamma_demog <- 10
births_demog <- 100
da_demog <- c(20,60)                          # alpha = 1/da
```

```{r}
ja.demog.model <- function (t, x, ...) {
  s <- x[c("Sj","Sa")]                  # susceptibles
  i <- x[c("Ij","Ia")]                  # infecteds
  r <- x[c("Rj","Ra")]                  # recovereds
  n <- s+i+r                            # total pop
  lambda.j <- (b1_demog+b2_demog)*i[1]+b1_demog*i[2]      # juv. force of infection 
  lambda.a <- b1_demog*i[1]+(b1_demog+b2_demog)*i[2]      # adult. force of infection
  alpha <- 1/da_demog
  list(
       c(
         -lambda.j*s[1]           -alpha[1]*s[1]+births_demog,
         -lambda.a*s[2]           +alpha[1]*s[1]-alpha[2]*s[2],
          lambda.j*s[1]-gamma_demog*i[1]-alpha[1]*i[1],
          lambda.a*s[2]-gamma_demog*i[2]+alpha[1]*i[1]-alpha[2]*i[2],
                        gamma_demog*i[1]-alpha[1]*r[1],
                        gamma_demog*i[2]+alpha[1]*r[1]-alpha[2]*r[2]
         )
       )
}
```

Note that in this function, $\mu=$ `alpha[2]`, i.e., death, is just like another age class.

```{r}
require(deSolve)

## initial conditions
yinit <- c(Sj=2000,Sa=3000,Ij=0,Ia=1,Rj=0,Ra=1000)

sol_demog <- ode(
           y=yinit,
           times=seq(0,200,by=0.1),
           func=ja.demog.model
           )

```

### Exercise 1: Use this code to plot the number of susceptible, infected, and recovered individuals over time.
```{r}
plot(sol_demog, mfcol=c(2, 3))
dim(sol_demog)
head(sol_demog)
plot(sol_demog,log='y',mfcol=c(2, 3))
```

```{r}
plot(sol_demog, mfcol=c(2, 3))
par(mfcol=c(1,1))
```

Note that now that births are replenishing susceptibles infection persists. The results of the above are plotted here :

```{r}
#plot the relative proportion of individuals in each category
par(mfrow=c(1,2))
time <- sol_demog[,1]                         # time
y <- sol_demog[,-1]                           # all other variables
n <- apply(y,1,sum)                     # population size
prop <- y/n                             # fractions
subsampled.prop <- prop[seq(1,length(time),by=10),]
subsampled.time <- time[seq(1,length(time),by=10)]
barplot(
        t(subsampled.prop),
        names.arg=subsampled.time,
        xlab='time',main='Age structure',
        space=0,
        col=c(
          rgb(0.5,1,0.5),
          rgb(0,1,0),
          rgb(1,0.5,0.5),
          rgb(1,0,0),
          rgb(0.5,0.5,1),
          rgb(0,0,1)
          ),
        legend = colnames(prop),
        args.legend=list(bg="white"),
        border = NA
        )
equil <- drop(tail(sol_demog,1))[-1]
n <- equil[c("Sj","Sa")]+equil[c("Ij","Ia")]+equil[c("Rj","Ra")]
seroprev <- equil[c("Rj","Ra")]/n            # get the proportion in the R class at equilibrium
names(seroprev) <- c("J","A")                

#plot the proportion of each age class 
# that has been previously infected (is in the R class) 
# and would thus be expected to have antibodies (i.e. seropositive)
barplot(height=seroprev,width=da_demog,ylab="seroprevalence")  
```

```{r}
par(mfrow=c(1,1))
```


To compute $R_0$, we need to know the stable age distribution (the relative proportion in the juvenile and adult age classes) of the population, which we can find by solving for the disease-free equilibrium: $S_J^*=B/\alpha$ and $S_A^*=B/\mu$.

With the stable age distribution, we can calculate $R_0$ by constructing the next generation matrix. Details on this method are described in the last section of this worksheet and are not required for completing the exercises in the worksheet. The code below outlines how the next generation matrix is constructed using the $\alpha$ (aging from juvenile to adult), $\mu$ death), $n$ (total births), $\gamma$ (recovery), $da$ (width of age groups in years), and $\beta$ (transmission) parameters. 


```{r}

# this code can be re-run to set parameter values
alpha <- 1/da_demog[1]
mu <- 1/da_demog[2]
n <- births_demog/c(alpha,mu)
beta_demog <- matrix(c(b1_demog+b2,b1_demog,b1_demog,b1_demog+b2),nrow=2,ncol=2)

# this command craetes the next generation matrix
ngm <- matrix(
              c(
                n[1]*beta_demog[1,1]/(gamma_demog+alpha)+
                  alpha/(gamma_demog+mu)*n[1]*beta_demog[1,2]/(gamma_demog+mu),
                n[2]*beta_demog[2,1]/(gamma_demog+alpha)+
                  alpha/(gamma_demog+mu)*n[2]*beta_demog[2,2]/(gamma_demog+mu),
                n[1]*beta_demog[1,2]/(gamma_demog+mu),
                n[2]*beta_demog[2,2]/(gamma_demog+mu)
                ),
              nrow=2,
              ncol=2
              )

# this command outputs the R0 value
max(Re(eigen(ngm,only.values=T)$values))
```



## Getting more realistic: adding more age classes

In the models above, the aging process follows an exponential distribution, which means that whether an individual is 1\~year old or 10 years old, the chance of them becoming an adult is the same!
To improve on this, we can assume that the time a juvenile must wait before becoming an adult follows a gamma distribution.
This is equivalent to saying that the waiting time is a sum of some number of exponential distributions.
This suggests that we can achieve such a distribution by adding age classes to the model, so that becoming an adult means passing through some number of stages.
We'll use 30 age classes, and since they don't have to be of equal duration, we'll assume that they're not.
Specifically, we'll have 20 1-yr age classes to take us up to adulthood and break adults into 10 age classes of 5\~yr duration each. The last age class covers age 66-80.

Now, when we had just two age classes, we could write out each of the equations easily enough, but now that we're going to have 30, we'll need to be more systematic.
In particular, we'll need to think of $\beta$ as a matrix of transmission rates.
Let's see how to define such a matrix in **R**.
So that we don't change too many things all at once, let's keep the same contact structure as in the juvenile-adult model.

```{r}

b1_ages <- 0.02  #transmission rate
b2_ages <- 0.01  #transmission rate
gamma_ages <- 10
births_ages <- 100


ages <- c(seq(1,20,by=1),seq(25,65,by=5),80) # upper end of age classes
da_ages <- diff(c(0,ages))                        # widths of age classes

# set up a matrix of contact rates between classes -- more contact within juveniles and adults than between
beta_ages <- matrix(nrow=30,ncol=30)    
beta_ages[1:20,1:20] <- b1_ages # transmission rate for juveniles 
beta_ages[21:30,21:30] <- b2_ages # transmission rate for adults
beta_ages[1:20,21:30] <- b2_ages/2 # lower transmission rate between juveniles and adults
beta_ages[21:30,1:20] <- b2_ages/2 # lower transmission rate between juveniles and adults

#WAIFW stands for Who Aquires Infection From Whom
filled.contour(ages, ages, beta_ages,                
               plot.title=title(
                 main="WAIFW matrix",      
                 xlab="age",ylab="age"))

```

We'll assume that, at the time of introduction, all children are susceptible, as are adults over 45, but that individuals aged 20--45 have seen the pathogen before and are immune.
The vector `yinit` expresses these initial conditions.

```{r}
yinit_ages <- c(
           S=c(rep(100,20),rep(0,5),rep(200,5)),
           I=c(rep(0,25),1,rep(0,4)),
           R=c(rep(0,20),rep(1000,5),rep(0,5))
           )
```

Note that we're starting out with 1 infected individual in the 26th age class.

The codes that follow will be a bit easier to follow if we introduce some indexes that will allow us to pick out certain bits of the `yinit` vector.

```{r}
sindex <- 1:30
iindex <- 31:60
rindex <- 61:90
juvies <- 1:20
adults <- 21:30
```

Now, to capture the aging process, it's convenient to define another matrix to hold the rates of movement between age classes.

```{r}
aging <- diag(-1/da_ages)
aging[row(aging)-col(aging)==1] <- 1/head(da_ages,-1)
```

Have a look at the aging matrix, for example by doing:

```{r}
# move fast through the 1-year age classes - negatives are moves out, positives are moves in
aging[1:5,1:5] 

# don't age between these classes -- e.g. can't age from 1 to 6
aging[1:5,6:10] 

# move slowly between the wider age classes
aging[25:30,25:30] 

#plot the aging matrix
filled.contour(ages, ages, aging,                
               plot.title=title(
                 main="Aging matrix",
                 xlab="age",ylab="age"))
```

### Exercise 2: What can you say about its structure? How are the different age groups in contact with each other?

Now we can put the pieces together to write a simulator for the age-structured SIR dynamics.

```{r}
ja.multistage.model <- function (t, x, ...) {
  s <- x[sindex]                  # susceptibles
  i <- x[iindex]                  # infecteds
  r <- x[rindex]                  # recovereds
  lambda <- beta_ages%*%i              # force of infection
  dsdt <- -lambda*s+aging%*%s
  didt <-  lambda*s+aging%*%i-gamma_ages*i
  drdt <-           aging%*%r+gamma_ages*i 
  dsdt[1] <- dsdt[1]+births_ages
  list(
       c(
         dsdt,
         didt,
         drdt
         )
       )
  
}
```

We can plug this into `ode` just as we did the simpler models to simulate an epidemic.
We'll then plot the epidemic curve.

```{r}

sol_ms <- ode(
           y=yinit_ages,
           times=seq(0,100,by=0.1),
           func=ja.multistage.model,
           parms = c(beta_ages = beta_ages, aging=aging, births_ages = births_ages)
           )
time_ms <- sol_ms[,1]
infects_ms <- sol_ms[,1+iindex]
plot(time_ms,apply(infects_ms,1,sum),type='l')
lines(time_ms,apply(infects_ms[,juvies],1,sum),col='red')
lines(time_ms,apply(infects_ms[,adults],1,sum),col='blue')
```

Let's mimic a situation where we have cross-sectional seroprevalence data (e.g. measures of antibodies that tell you someone is in the R class).
In using such data, we'd typically assume that the system was at equilibrium.

### Exercise 3: What does the equilibrium age-specific seroprevalence look like in this example?
Use the code below to display the age-specific seroprevalence (i.e., the seroprevalence for each age group at equilibrium)

```{r}
equil_ms <- drop(tail(sol_ms,1))[-1]
n_ms <- equil_ms[sindex]+equil_ms[iindex]+equil_ms[rindex]
seroprev_ms <- equil_ms[rindex]/n_ms
names(seroprev_ms) <- ages
barplot(height=seroprev_ms,width=da_ages)
```

Let's also compute $R_0$.
To do so, we'll need the stable age distribution.
We can get that by simulating an infection-free population, which we get by setting the initial I to all 0s:

```{r}
yinit.sonly <- c(
                 S=c(rep(250,30)),
                 I=c(rep(0,30)),
                 R=c(rep(0,30))
                 )
sol_sonly <- ode(
           y=yinit.sonly,
           times=seq(0,300,by=1),
           func=ja.multistage.model,
           parms = c(beta=beta, aging=aging, births=births_demog)
           )
time_sonly <- sol_sonly[,1]
pop <- apply(sol_sonly[,-1],1,sum)
plot(time_sonly,pop,type='l')
```

Alternatively, we can get the stable age distribution by finding the population structure that balances the birth, aging, and death processes.
At equilibrium, we have the matrix equation $$
 \begin{pmatrix}
    -\alpha_1 & 0 & 0 & \cdots & 0\\
    \alpha_1 & -\alpha_2 & 0 & \cdots & 0\\
    0 & \alpha_2 & -\alpha_3 & \cdots & 0\\
    \vdots &  & \ddots & \ddots & \vdots \\
    0 & \cdots & & \alpha_{29} & -\alpha_{30}\\
  \end{pmatrix} . \begin{pmatrix}
    n_1 \\ n_2 \\ n_3 \\ \vdots \\ n_{30}
  \end{pmatrix} + \begin{pmatrix}
    B \\ 0 \\ 0 \\ \vdots \\ 0
  \end{pmatrix}=
  \begin{pmatrix}
    0 \\ 0 \\ 0 \\ \vdots \\ 0
  \end{pmatrix}
$$ To solve this equation in **R**, we can do

```{r}
## get stable age distribution
n <- solve(aging,-c(births_ages,rep(0,29)))
```

The following lines then compute $R_0$ using the next generation matrix method. More details are available in the "Bonus" section at the end of the document.
This calculation comes from a recipe described in detail by **Diekmann & Heesterbeek, 2000** and **Hurford *et. al*, 2010**.

```{r}
F <- diag(n)%*%beta_ages+aging-diag(diag(aging))
V <- diag(gamma_ages-diag(aging))
max(Re(eigen(solve(V,F),only.values=T)$values))

```

```{r,}
seroprev1 <- seroprev
```

### Exercise 4: 
#### a. Change the juvenile and adult contact rates (b1_ages and b2_ages) to reflect different transmission within groups. Make the juvenile contact rate 0.02 to reflect higher contact among kids (e.g. in schools).
#### b. Use `image` or `filled.contour` to plot the $\beta$ matrix.
#### c. Compute $R_0$ for your assumptions.
#### d. Simulate and plot the age-structured SIR dynamics under your assumptions and record how the age-specific seroprevalence has changed.

```{r}
b1_r <- 0.007
b2_r <- 0.02
b3_r <- 0.03
beta <- matrix(data=b1_r,nrow=30,ncol=30)
beta[1:20,1:20] <- b2_r
beta[6:16,6:16] <- b3_r
beta_ages <- beta
filled.contour(ages, ages, beta_ages,                
               plot.title=title(
                 main="WAIFW matrix",
                 xlab="age",ylab="age"))



sol_r <- ode(
           y=yinit_ages,
           times=seq(0,400,by=0.1),
           func=ja.multistage.model,
           parms = c(beta_ages=beta_ages, aging=aging, births_ages=births_ages)
           )

plot(sol_r[,1],apply(sol_r[,1+iindex],1,sum),type='l')
lines(sol_r[,1],apply(sol_r[,1+iindex[juvies]],1,sum),col='red')
lines(sol_r[,1],apply(sol_r[,1+iindex[adults]],1,sum),col='blue')

equil_r <- drop(tail(sol_r,1))[-1]
n_r <- equil_r[sindex]+equil_r[iindex]+equil_r[rindex]
seroprev_r <- equil_r[rindex]/n_r
names(seroprev_r) <- ages
barplot(height=seroprev_r,width=da_ages)

## get stable age distribution
n_r <- solve(aging,-c(births_ages,rep(0,29)))
## get R0
F <- diag(n_r)%*%beta_ages+aging-diag(diag(aging))
V <- diag(gamma_ages-diag(aging))
max(Re(eigen(solve(V,F),only.values=T)$values))
# R0 = 6.53
infects_r <- sol_r[dim(sol_r)[1],1+iindex]
sum(ages * infects_r/sum(infects_r))
sum(infects_r[15:23])

```

```{r}
seroprev2 <- seroprev
```

```{r}
require(ggplot2)
require(reshape)
sol_r <- as.data.frame(sol_r)
sol_rm <- melt(sol_r,id="time")
sol_rm$class <- factor(substr(as.character(sol_rm$variable),start=1,stop=1))
sol_rm$age <- ages[as.integer(substr(as.character(sol_rm$variable),2,stop=5))]
ggplot(data=sol_rm)+geom_line(aes(x=time,y=value,col=age,group=age))+facet_grid(class~.,scales="free")
```


## R0 and the mean age of infection

For simplicity, let's return to the earlier models with a simple age-class mixing matrix.
But this time, we'll calculate $R_0$, the mean age of infection, and the number of cases between 15-35 years as we increase the rate of contact. Recall from the rubella and CRS example that the risk of severe disease outcomes depends on the risk of infection in reproductive age women. Recall also that increasing vaccination reduces $R_E$ -- so here we'll evalate at several values of $R_0$ as a proxy for the impact of vaccination. We'll then calculate how the mean age of infection changes, and specifically how the absolute number of cases among individuals between the ages of 15-35 (as a proxy for reproductive age women) changes. 
To do so, we'll make a loop and evaluate the code for each of 10 increasing levels of mixing (which whill change R0)

```{r}
 # a vector of scaling factors, we'll reduce the contact rate 
 # from the original code by each of the values in this vector
scale <- seq(.2,1,length=10) 
R0 <- numeric() # somewhere to store results for the mean age
mean_age <- numeric() # somewhere to store results for the mean age
sum_cases <- numeric() # somewhere to store results for the mean age

# this loop will run for as many different levels of contact that we specify 
# for the scale variable 
# this is the same code as above, but now we've included a multiplier for the 
# contact matrix
for(ii in 1:length(scale)){ 
  b1 <- 0.007
  b2 <- 0.02
  b3 <- 0.03
  beta <- matrix(data=b1,nrow=30,ncol=30)
  beta[1:20,1:20] <- b2
  beta[6:16,6:16] <- b3
  beta_ages <- beta * scale[ii]
  #filled.contour(beta)

  
  sol_R0 <- ode(
    y=yinit_ages,
    times=seq(0,400,by=0.1),
    func=ja.multistage.model,
    parms = c(beta_ages=beta_ages, aging=aging, births_ages=births_ages)
  )
  
  
  equil_R0 <- drop(tail(sol_R0,1))[-1]
  n_R0 <- equil_R0[sindex]+equil_R0[iindex]+equil_R0[rindex]
  seroprev_R0 <- equil_R0[rindex]/n_R0
  names(seroprev_R0) <- ages
  # barplot(height=seroprev,width=da)
  
  ## get stable age distribution
  n_R0 <- solve(aging,-c(births_ages,rep(0,29)))
  ## get R0
  F_R0 <- diag(n_R0)%*%beta_ages+aging-diag(diag(aging))
  V_R0 <- diag(gamma_ages-diag(aging))
  R0[ii] <- max(Re(eigen(solve(V_R0,F_R0),only.values=T)$values))
  # R0 = 6.53
  infects_R0 <- sol_R0[dim(sol_R0)[1],1+iindex]
  mean_age[ii] <- sum(ages * infects_R0/sum(infects_R0))
  sum_cases[ii] <- sum(infects_R0[15:23])
}


```

Now we can make a table of the results and plot mean age and the sum of cases between 15-35 years of age as a function of $R_0$.

```{r}
df <- data.frame(R0 = R0, "mean age" = mean_age, "cases (15-35y)" = sum_cases)

knitr::kable(df)

par(mfrow=c(1,2))
plot(R0,mean_age,xlab="R0",ylab="mean age of infection")
plot(R0,sum_cases,xlab="R0",ylab="total cases between 15-35y")


```
```{r }
par(mfrow=c(1,1))
```

### Exercise 5: Try the same as you fix $R_0$ but change the birth rate (as if new infants were vaccinated), instead of changing the contact rate. 

For example, if the birth rate was 100 before, try using a sequence of birth rates ranging from 100 (i.e. no new infants are vaccinated) to 75 (i.e., 25% of new infants are vaccinated). You can use the technique we used above, where we ran the simulation in a loop multiple times, for varying levels of contact, modifying it to run on varying levels of births (e.g., between 75 and 100). 

```{r}
births_seq <- seq(75,100,length=10)  # a vector of scaling factors, we'll reduce the contact rate from the original code by each of the values in this vector
R0_b <- numeric() # somewhere to store results for the mean age
mean_age_b <- numeric() # somewhere to store results for the mean age
sum_cases_b <- numeric() # somewhere to store results for the mean age


  
for(ii in 1:10){
#this is the same code as above, but now we're modifying births
  b1 <- 0.007
  b2 <- 0.02
  b3 <- 0.03
  beta <- matrix(data=b1,nrow=30,ncol=30)
  beta[1:20,1:20] <- b2
  beta[6:16,6:16] <- b3
  beta_ages <- beta
  births_ages <- births_seq[ii]


  #filled.contour(beta)

  sol_b <- ode(
    y=yinit_ages,
    times=seq(0,400,by=0.1),
    func=ja.multistage.model,
    parms = c(beta_ages=beta_ages, aging=aging, births_ages=birth_ages)
  )

  equil_b <- drop(tail(sol_b,1))[-1]
  n_b <- equil_b[sindex]+equil_b[iindex]+equil_b[rindex]
  seroprev_b <- equil_b[rindex]/n_b
  names(seroprev_b) <- ages
  # barplot(height=seroprev,width=da)

  ## get stable age distribution
  n_b <- solve(aging,-c(births_ages,rep(0,29)))
  ## get R0
  F <- diag(n_b)%*%beta_ages+aging-diag(diag(aging))
  V <- diag(gamma_ages-diag(aging))
  R0_b[ii] <- max(Re(eigen(solve(V,F),only.values=T)$values))
  # R0 = 6.53
  infects_b <- sol_b[dim(sol_b)[1],1+iindex]
  mean_age_b[ii] <- sum(ages * infects_b/sum(infects_b))
  sum_cases_b[ii] <- sum(infects_b[15:23])
}

```

```{r}
df <- data.frame(births = births_seq, R0 = R0_b, "mean age" = mean_age_b, 
                 "cases (15-35y)" = sum_cases_b)

knitr::kable(df)

par(mfrow=c(1,2))
plot(births_seq,mean_age_b,xlab="births",ylab="mean age of infection")
plot(births_seq,sum_cases_b,xlab="births",ylab="total cases between 15-35y")
#plot(births_seq,R0_b,xlab="births",ylab="R0")


```



## What do real contact networks look like?

The POLYMOD study **Mossong, 2008** was a journal-based look into the contact network in contemporary European society.
Let's have a look what these data tell us about the contact structure.

```{r}
moss <- read.csv(
                 url("http://www.math.mcmaster.ca/~bolker/eeid/data/mossong.csv"),
                 as.is=TRUE
                 )
age.categories <- moss$contactor[1:30]
moss$contactor <- ordered(moss$contactor,levels=age.categories)
moss$contactee <- ordered(moss$contactee,levels=age.categories)
```

Since contacts are symmetric, we'll need to estimate the symmetric contact matrix.

```{r}
x1 <- with(
           moss,
           tapply(contact.rate,list(contactor,contactee),unique)
           )
xsym <- (x1+t(x1))/2

```

```{r}
filled.contour(ages,ages,log10(xsym))
filled.contour(
               ages,ages,log10(xsym),
               plot.title=title(
                 main=quote(log[10](contact~rate)),
                 xlab="age",ylab="age")
               )
barplot(height=apply(x1,1,sum))
barplot(height=apply(x1,2,sum))
```

While this matrix tells us how many contacts are made per year by an individual of each age, it doesn't tell us anything about the probability that a contact results in communication of infection.
Let's assume that each contact has a constant probability $q$ of resulting in a transmission event.

```{r}
q <- 3e-5
beta_ages <- q*xsym
filled.contour(ages,ages,log10(beta_ages),                
               plot.title=title(
                 main="WAIFW matrix based on POLYMOD data",
                 xlab="age",ylab="age"))

```

Now let's simulate the introduction of such a pathogen into a population characterized by this contact structure.

```{r}
sol_p <- ode(
           y=yinit_ages,
           times=seq(0,200,by=0.5),
           func=ja.multistage.model,
           parms = c(beta_ages=beta_ages, aging=aging, births=births_ages)
           )

time <- sol_p[,1]
infects_p <- sol_p[,1+iindex]
plot(time,apply(infects_p,1,sum),type='l')
lines(time,apply(infects_p[,juvies],1,sum),col='red')
lines(time,apply(infects_p[,adults],1,sum),col='blue')

```

As before, we can also look at the equilibrium seroprevalence

```{r}
equil_p <- drop(tail(sol_p,1))[-1]
n_p <- equil_p[sindex]+equil_p[iindex]+equil_p[rindex]
seroprev_p <- equil_p[rindex]/n_p
names(seroprev_p) <- ages
barplot(height=seroprev_p,width=da_ages)
```

and compute the $R_0$ for this infection.

```{r}
n_p <- solve(aging,-c(births_ages,rep(0,29)))
F_p <- diag(n_p)%*%beta_ages+aging-diag(diag(aging))
V_p <- diag(gamma_ages-diag(aging))
max(Re(eigen(solve(V_p,F_p),only.values=T)$values))

```
How does this R0 value compare to the R0 value obtained from Exercise 4?

```{r}
par(mfrow=c(1,1))
```

## Bonus: Calculating R0 Using a Next Generation Matrix

The next generation matrix is a matrix that specifies how many new age-specific infections are generated by a typical infected individual of each age class (in a fully susceptible population).
For example, let's consider an infected adult and ask how many new juvenile infections it generates: this is the product of the number of susceptible juveniles (from the stable age distribution), the per capita transmission rate from adults to juveniles and the average duration of infection, i.e. $S_J^* \times \beta_{JA} \times 1/ (\gamma+\mu)$.
This forms one element of our next generation matrix.
The other elements look very similar, except there are extra terms when we consider an infected juvenile because there is a (very small) chance they may age during the infectious period and therefore cause new infections as an adult: $$
 \mathrm{NGM} = \left(
  \begin{matrix}
    \frac{S_J^* \beta_{JJ} }{(\gamma+\alpha)} +\frac{\alpha}{(\gamma+\mu)}\frac{S_J^*\beta_{JA}}{(\gamma+\mu)} &  
    \frac{S_J^* \beta_{JA}}{(\gamma+\mu)} \\
    \frac{S_A^* \beta_{AJ} }{(\gamma+\alpha)} +\frac{\alpha}{(\gamma+\mu)}\frac{S_A^*\beta_{AA}}{(\gamma+\mu)} &  \frac{S_A^* \beta_{AA}}{(\gamma+\mu)} 
  \end{matrix}
\right)
$$

$R_0$ can then be computed as the dominant eigenvalue (i.e., the one with the largest real part) of this matrix. Let's take an example from a model with 2 age classes, from above. First, let's define the components of the next generation matrix:

```{r}
da_ngm <- c(20,60) # this classifies the two age groups (0-20, 21-60) 
b1_ngm = 0.005
b2_ngm = 0.006
alpha_ngm <- 1/da_ngm[1]
mu_ngm <- 1/da_ngm[2]
n_ngm <- births_demog/c(alpha_ngm,mu_ngm)
beta_ngm <- matrix(c(b1_ngm+b2_ngm,b1_ngm,b1_ngm,b1_ngm+b2_ngm),nrow=2,ncol=2) 
gamma_ngm = 10
```



The Next Generation Matrix can be calculated in **R** as:

```{r}
ngm <- matrix(
              c(
                n_ngm[1]*beta_ngm[1,1]/(gamma_ngm+alpha_ngm)+
                  alpha_ngm/(gamma_ngm+mu_ngm)*n[1]*beta_ngm[1,2]/(gamma_ngm+mu_ngm),
                n_ngm[2]*beta_ngm[2,1]/(gamma_ngm+alpha_ngm)+
                  alpha_ngm/(gamma_ngm+mu_ngm)*n[2]*beta_ngm[2,2]/(gamma_ngm+mu_ngm),
                n_ngm[1]*beta_ngm[1,2]/(gamma_ngm+mu_ngm),
                n_ngm[2]*beta_ngm[2,2]/(gamma_ngm+mu_ngm)
                ),
              nrow=2,
              ncol=2
              )
eigen(ngm)
eigen(ngm,only.values=TRUE)
max(Re(eigen(ngm,only.values=T)$values))

```