---
title: "Generation time in stage-structured populations under fluctuating environments (Simulations)"
author:    |
    | Stefano Giaimo (giaimo@evolbio.mpg.de) & Arne Traulsen 
    | Evolutionary Theory Department
    | Max Planck Institute for Evolutionary Biology
    | August-Thienemann-Str. 2, 24306 Ploen, Germany
date: "05 July 2022"
output:
  pdf_document: default
  html_document: default
---

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

## Introduction

Here, we report the code to generate Figure 2 in the main text.
Code is written in R (version 4.1.2) and run on the platform x86_64-apple-darwin17.0 (64-bit).

## Figure 1

Clean the environment:

```{r clean2}
rm(list=ls())
```

Data for this figure come from Claessen, *AmNat* $2005(165), \textrm{ pp. }E27$–$E35$.

Number environmental states:

```{r envsfig2}
my.envs <- c(1,2) # 1 is a good year, 2 is bad year
my.envs
```

Set probabilities of a good year and a bad year:

```{r envprobsfig2}
pr.good <- 0.1
pr.bad <- 1-pr.good 
```

Number of stages:

```{r nrstagesfig2}
nr.stages <- 3 
```

Number of time steps in the simulations to estimate the stochastic elasticities:

```{r TlimitFig2}
Tlimit <- 1000
```

Survival matrices from Eq.11 in Claessen, *AmNat*  $2005(165), \textrm{ pp. }E27$-$E35$ as a list:

```{r SmatsFig2}
# these are the survival matrices
basesurv.mats <- list(good = matrix( c(  0,  0.0029,   0.042, # good years
                               0,  0.0012,   0.017,
                               0,  0.78,     0.57) ,
                           ncol=nr.stages,nrow=nr.stages,byrow = TRUE),
                   
                      bad = matrix( c(  0,       0,       0.0001, # bad years
                               0,       0.0018,  0.026,
                               0,       0.78,    0.57) ,
                           ncol=nr.stages,nrow=nr.stages,byrow = TRUE))
```

Fecundity matrices from Eq.11 in Claessen, *AmNat*  $2005(165), \textrm{ pp. }E27$-$E35$ as a list:

```{r FmatsFig2}
# these are the fecundity matrices
basefec.mats <- list(good = matrix( c(  14.7,  0,  0, # good years
                              1.49,  0,  0,
                              970,   0,  0) ,
                          ncol=nr.stages,nrow=nr.stages,byrow = TRUE),
                  
                     bad = matrix( c(  0.019,  0,  0, # bad years
                              0.89,   0,  0,
                              388,    0,  0) ,
                          ncol=nr.stages,nrow=nr.stages,byrow = TRUE))
```

Set average survival matrix:

```{r aveSFig2}
U.survmat <-  (pr.good * basesurv.mats[[1]]) + (pr.bad * basesurv.mats[[2]])
U.survmat
```

Set average fecundity matrix
```{r aveFFig2}
U.fecmat <-  (pr.good *  basefec.mats[[1]]) + (pr.bad * basefec.mats[[2]])
U.fecmat
```

Set average projection matrix
```{r aveFig2}
U.mat <-  U.survmat + U.fecmat
U.mat
```

Set parameter regulating variability ($\theta$ in main text)
```{r theta}
v.par <- seq(0,1,length.out = 41)
v.par
```

Generate all survival and fecundity matrices for each value of theta following Eqs. 6-7 Claessen, *AmNat*  $2005(165), \textrm{ pp. }E27$-$E35$ and put them in lists:

```{r SandFwiththeta}
allsurv.mats <- lapply(v.par, function(x) {
  l <- lapply(basesurv.mats, function(y) {U.survmat + x*(y - U.survmat)})
  return(l)
})
head(allsurv.mats)
tail(allsurv.mats)

allfec.mats <- lapply(v.par, function(x) {
  l <- lapply(basefec.mats, function(y) {U.fecmat + x*(y - U.fecmat)})
  return(l)
})
head(allfec.mats)
tail(allfec.mats)
```

Generate a sequence of Tlimit environments:

```{r seqenFig2}
seq.env <- sample(my.envs, size = Tlimit, prob = c(pr.good,pr.bad), replace = TRUE)
```
Set initial initial stage distribution, i.e. the stage distribution at time 0 as positive random vector of the appropriate dimension with components adding up to 1:

```{r sd0FIg2}
sd0 <- runif(nr.stages, min = 1,max = 10)
# then the vector is normalized so that its components add up to 1 and put in column form
sd0 <- t(t(sd0 / sum(sd0)))
sd0
```

Set the final reproductive value vector, i.e. the stage-specific reproductive value at time Tlimit, as a positive random vector of the appropriate dimension with components adding up to 1:

```{r rvFig2}
rvT <- runif(nr.stages, min = 1,max = 10)
# then the vector is normalized so that its components add up to 1 and put in column form
rvT <- (t(t(rvT / sum(rvT))))
rvT
```

A function to project the stage distribution (sd) over one time step using a given population projection matrix (ppm) and retrieve population growth over this time step:
 
```{r updatesdFig2}
update.sd <- function(sd,ppm) {
  new.sd <- ppm%*%sd
  new.Rt <- sum(new.sd)
  new.sd <- new.sd /new.Rt
  return(list(new.sd=new.sd, new.Rt=new.Rt))
}
```

A function to project backward in time the reproductive value (rv) over one time step using a given population projection matrix (ppm):

```{r updatervFig2}
update.rv <- function(rv,ppm) {
  new.rv <- t(t(rv)%*%ppm)
  new.rv <- new.rv /sum(new.rv)
  return(new.rv)
}
```

Find stochastic elasticity matrix for each value of $\theta$:

```{r eSfromtheta}
veS <- lapply(1:length(v.par), function(x) {
  
  # compose the projection matrix for the given value of v
  my.mats <- list(allsurv.mats[[x]][[1]] + allfec.mats[[x]][[1]],
                  allsurv.mats[[x]][[2]] + allfec.mats[[x]][[2]])
    
  # generate a sequence of population projection matrices
  # corresponding to the sequence of environments
  seq.mats <- sapply(seq.env, function(x) my.mats[[x]], simplify = FALSE)
  
  # initialize the sequence of stage distribution as a list where
  # the element j is the stage distribution at time j-1
  seq.sd <- list(sd0)
  
  # initialize the sequence of time specific populaton growth rates
  seq.Rt <- c()
  
  # generate the sequences of stage distributions and time-specific
  # growth rates,
  j <- 0
  while (j <= (Tlimit-1)) {
    # the stage distribution in the next time step and the realized growth
    # from the current to the next time step are computed from the current
    # stage distribution seq.sd[j+1] and population projection matrix
    my.sdRt <- update.sd( sd = seq.sd[[j+1]], ppm = seq.mats[[j+1]])
    # the sequences of stage distributions and growth rates are updated accordingly
    seq.sd <- append(seq.sd, list(my.sdRt[["new.sd"]]))
    seq.Rt <- c(seq.Rt, my.sdRt[["new.Rt"]])
    # the time index is updated
    j <- j+1
  }
  
  # generate the backward sequence of reproductive values
  seq.rv <- Reduce(f = update.rv, x = rev(seq.mats), init  = rvT, accumulate = TRUE)
  # the sequence is then put in forward time as a list
  seq.rv <- rev(seq.rv)
  
  # the sequence of relative reproductive values is computed
  seq.rrv <- mapply(function(x,y) x/c(t(x)%*%y), seq.rv, seq.sd, SIMPLIFY = FALSE)
  
  
  # compute the stochastic elasticities
  eS <- lapply(1:Tlimit, function(x) {
    term <- ((seq.rrv[[x+1]] %*% t(seq.sd[[x]])) * seq.mats[[x]]) / (seq.Rt[x])
    return(term)
  } )
  eS <- (Reduce(f = "+", eS))/Tlimit
  # returns the estimated stochastic elasticities
  if ( x < length(v.par)) {  return(eS) }
  # for the case of v=1, it also returns the sequence of stage distributions,
  # the time-specific growth rates and time-specific relative reproductive values
  if ( x == length(v.par)) {  return(list(eS=eS,
                                          seq.sd=seq.sd,
                                          seq.Rt=seq.Rt,
                                          seq.rrv=seq.rrv)) }
})
```

Store time-specific quantities (stage distribution, reproductive values...) from the simulation of the demographic process for $\theta=1$:

```{r theta1data}
v1process <- veS[[length(v.par)]]
v1process$eS <- NULL
```

These quantities will be used below to compute the generation time using Eq. A16 in the Appendix to main text.

Reports the stochastic elasticities for each value of $\theta$:

```{r eStheta}
veS[[length(v.par)]] <- veS[[length(v.par)]]$eS
veS
```

Check that stochastic elasticities for each matrix add up to 1 and that row and column sum of the elasticity matrices are approximately equal:

```{r checkeSFig2}
lapply(veS, function(x) sum(x))
lapply(veS, function(x) return(cbind(colSums(x),rowSums(x))))
```

Reproduce Fig 2C in Claessen, *AmNat*  $2005(165), \textrm{ pp. }E27$-$E35$:

```{r reproFig2CClaessen}
plot(v.par, sapply(veS, function(x) x[1,1]), pch= 16, ylim = c(0,0.7),
     xlab = "Variability (v)", ylab = "Stochastic elasticity")
points(v.par, sapply(veS, function(x) x[2,1]), pch= 1)
points(v.par, sapply(veS, function(x) x[3,1]), pch= 17)
points(v.par, sapply(veS, function(x) x[2,2]), pch= 6)
points(v.par, sapply(veS, function(x) x[3,2]), pch= 15)
points(v.par, sapply(veS, function(x) x[3,3]), pch= 0)
```

Compute the generation time from stochastic elasticities for each value of $\theta$ using Eq. 9 in the main text:

```{r stocgentimeEq14}
velasgentime <- sapply(veS, function(x) 1/sum(x[U.fecmat > 0]), simplify = TRUE)
velasgentime
```

Function to find the deterministic generation time from a given projection matrix using Eq. 3 in the main text:

```{r detTFig2}
detgentime <- function(ppm) {
  # the deterministic growth rate for the projection matrix as its leading eigenvalue
  lambda <- (abs(eigen(ppm,only.values = TRUE)$values))
  lambda <- max(lambda)
  # the stable stage distribution for the projection matrix  
  detSSD <- (eigen(ppm)$vectors)
  detSSD <- Re(detSSD[,which.max(abs(eigen(ppm, only.values = TRUE)$values))])
  # scaled so that its components add up to 1
  detSSD <- detSSD/sum(detSSD)
  # store reproductive value for the average projection matrix  
  detRV <- (eigen(t(ppm))$vectors)
  detRV <- Re(detRV[,which.max(abs(eigen(t(ppm), only.values = TRUE)$values))])
  # scaled so that its dot product with the stable stage distribution is equal to 1
  detRV <- detRV/(sum(detRV * detSSD))
  # determinstic elasticities
  eD <- sweep(ppm,2,detSSD,"*")
  eD <- sweep(eD,1,detRV,"*")
  eD <- eD/lambda
  # deterministic generation time
  detGT <- 1/sum(eD[U.fecmat > 0])
  return(detGT)
  }
```

Compute deterministic generation time for good years, bad years and average years for each $\theta$ value:

```{r GTtheta}
veladgentime <- sapply(1:length(v.par), function(x) {
  # projection matrix for the good years
  good.mat <- (allsurv.mats[[x]][[1]] + allfec.mats[[x]][[1]])
  # projection matrix for the bad years
  bad.mat <-  (allsurv.mats[[x]][[2]] + allfec.mats[[x]][[2]])
  # projection matrix for the average years
  ave.mat <- (pr.good * good.mat) + (pr.bad * bad.mat)
  # returns the generation time for the good years, bad years and average year
  return(c(detgentime(good.mat), detgentime(bad.mat), detgentime(ave.mat)))
}, simplify = TRUE)
rownames(veladgentime) <- c("good", "bad", "ave")
veladgentime
```


Generate a list containing the sequence of the fecundity matrices for $\theta=1$ corresponding to the sequence of environments:

```{r seqFFig2}
seq.Fmats <- sapply(seq.env, function(x) allfec.mats[[length(v.par)]][[x]],
                    simplify = FALSE)
```

Generate a list containing the sequence of the survival matrices  for $\theta=1$ corresponding to the sequence of environments:

```{r seqSFig2}
seq.Smats <- sapply(seq.env, function(x) allsurv.mats[[length(v.par)]][[x]],
                    simplify = FALSE)
```

Compute the denominator on the right hand side of Eq. A16 in the Appendix for all times points for $\theta=1$:

```{r denoFig2}
denosgentime <- mapply(function(x,y,w,z) c(t(y)%*%w%*%z)/x,
                       v1process$seq.Rt,
                       (v1process$seq.rrv)[-1],
                       seq.Fmats,
                       (v1process$seq.sd)[-(Tlimit+1)])
denosgentime <- Reduce("+", denosgentime, accumulate = TRUE)
```

Set the initial maximum age in the population
```{r initmaxageFIg2}
bigK <- 3
bigK
```

Set the initial, random age distribution within each stage in the population as a list with as many elements as the initial maximum age, each element is a vector with number of components equal to the number of stages. The $j$ component of the vector that is the $i$ element of the list reports the population fraction in stage $j$ and of age $i$ at time 0. This initial age-by-stage distribution is set at random with the constraint that the initial stage distribution should match that chosen before:

```{r agestagedist}
# First a random array with dimensions (nr.stages X bigK ) is created
agestdist <- replicate(n = bigK,
                       expr = sample(x = 1:100,
                                     size = nr.stages,
                                     replace = TRUE))
# then it is normalized so that its row sums are identical
# to the initial stage distribution
agestdist <- sweep(agestdist, MARGIN = 1,
                   STATS = sd0/rowSums(agestdist), FUN = "*")
# a little check that this is really so
rowSums(agestdist)
t(sd0)
```

Set a vector of all possible ages in the population:

```{r agesFig2}
ages <- seq_len(bigK)
ages
```

Generate the vector containing the numerator in Eq. A16 in the Appendix for all times points:

```{r numFig2}
numsgentime <- c()
j <- 0
while (j <= (Tlimit-1)) {
  # the columns of the age by stage distribution are weighted
  # by the age they refer to
  wagestdist <- sweep(agestdist, 2, STATS = ages, FUN = "*")
  # the inner sum in the numerator in Eq. 21 is computed
  insum <- sum((t((v1process$seq.rrv)[[j+2]])%*%(seq.Fmats[[j+1]])%*%(wagestdist)))/
    (v1process$seq.Rt)[j+1]
  numsgentime <- c(numsgentime, insum)
  # the age by stage distribution is updated by survival from each age class to the next
  agestdist <- apply(agestdist, 2, function(x) seq.Smats[[j+1]]%*%x)
  # and by adding the new age class 1
  agestdist <- cbind(seq.Fmats[[j+1]]%*%(v1process$seq.sd)[[j+1]], agestdist)
  agestdist <- agestdist/(v1process$seq.Rt)[j+1]
  # the vector of ages gets elongated
  ages <- c(ages, bigK+j+1)
  # the time index is updated
  j <- j+1
}
numsgentime <- Reduce("+", numsgentime, accumulate = TRUE)
head(numsgentime)
tail(numsgentime)
```


Check that the age vector is correct and the age-stage distribution adds up to 1:
```{r checkage}
head(ages, n = 2*bigK)
tail(ages, n = 2*bigK)
sum(agestdist)
```

Generate sequence of estimations of the generation time for $\theta=1$ at each time point using Eq. A16 in the Appendix to the main text and check that asymptotically, there is convergence to the generation time computed via Eq. 9 for $\theta=1$:

```{r stocgentimeFig2}

stochgentime <- numsgentime / denosgentime
head(stochgentime)
tail(stochgentime)
stochgentime[Tlimit]
tail(velasgentime, n=1)
```

Plot fig 2
```{r Fig2}
par(mfrow=c(1,2), mar = c(3,2,1,1))
my.ylim <- c(min(veladgentime[,length(v.par)]["good"],
                 veladgentime[,length(v.par)]["bad"])-1,
             max(veladgentime[,length(v.par)]["good"],
                 veladgentime[,length(v.par)]["bad"])+1)

fig2.mat <- t(rbind(veladgentime, velasgentime))
matplot(fig2.mat, xlab = "", xaxt = "n", ylab = "Generation time",
        ylim = my.ylim, type = "o", pch=16, lty = 1,
        col = c("deepskyblue","gold2", "pink", "red"), lwd=2)
axis(side = 1, at = c(1,11,21,31,41), labels = c(0,0.25,0.5,0.75,1))

time.window <- 1000
plot(NA, NA,xlim = c(0,time.window),  xlab = "Time", ylab = "Generation time",
     ylim = my.ylim, lwd=2)
xcor.goodyears <- which(seq.env[1:time.window] == 1)
ycor.goodyears <- c(0,0,rep(max(veladgentime[,length(v.par)]["good"],
                                veladgentime[,length(v.par)]["bad"])+1, 2))
invisible(sapply(xcor.goodyears, function(z) {
  polygon(x= c(z-0.5,z+0.5,z+0.5,z-0.5), y = ycor.goodyears,
          col = "greenyellow", border = NA)
}))
abline(h = tail(velasgentime, n=1), col="red", lwd=2)
abline(h = veladgentime[,length(v.par)]["good"], col= "deepskyblue", lwd=2)
abline(h = veladgentime[,length(v.par)]["bad"], col= "gold2", lwd=2)
abline(h = veladgentime[,length(v.par)]["ave"], col="pink", lwd=2)
lines(0:(time.window-1), stochgentime[1:time.window], lwd=2)
```