#          ##### Inclusive fitness forces of selection in an age-structured population #####                    
# Supplementary R Code for the numerical solutions in Figures 2 & 3 of the main manuscript #
#                           Submitted to Nature Communications Biology                     #
#                Author - Mark Roper. Contact email - markroper67@gmail.com                #
############################################################################################

#Start by clearing the work environment and downloading the required packages
rm(list=ls())
#install.packages(c("ggplot2", "popbio", "gridExtra"))
library(ggplot2)
library(popbio)
library(gridExtra)

#ggplot custom theme to reproduce figures

theme_m <-  theme(panel.background=element_blank(), 
                  strip.background=element_blank(), 
                  axis.line=element_line("black"), 
                  axis.ticks=element_line("black"), 
                  axis.text=element_text(colour="black", size=18), 
                  axis.title=element_text(size=20),
                  panel.grid=element_blank(),
                  strip.text=element_text(size=10),
                  legend.key=element_rect(fill="white"), 
                  axis.line.x=element_line("black"),
                  axis.line.y=element_line("black"), 
                  plot.title = element_text(color="#666666", face="bold", size=14))


##### Workflow #####

## The following script provides code for the two numerical solutions found in Figures 2 and 3 of the main manuscript.
# Each example goes through the same 5 stages:
# 1) The number of age classes and background demography are defined.
# 2) Social transfers are modelled (see Appendix C in the Supplementary Information for details).
# 3) Relatedness is quantified numerically (see Appendix A in the Supplementary Information for details).
# 4) An Inclusive Fitness matrix is computed.
# 5) Forces of selection are derived using Equations 8 and 10 in the main text, and Hamilton's indicators.

## The first example considers the case when post-reproductive individuals provide help to younger age classes (see Figure 2 in the main manuscript).
## The second example considers the case when pre-reproductive individuals provide help to older reproductive age classes (see Figure 3 in the main manuscript).


## I choose to clean the global environment between examples for personal preference, but variables are named such that the user does not have to do so.
# Variables in example 1 have the prefix EX1_ (example 1), whereas in example 2 have the prefix EX2_ (example 2).

#################
##  EXAMPLE 1  ##
#################

### STAGE 1 - Age classes and background demography

## 1.1: Define the number of age classes and store age classes in the vector EX1_age which runs from 1 to the last age class (omega) and which all individuals die.
EX1_omega <- 50 #last age class
EX1_age <- c(1:EX1_omega) 

## 1.2: Define demographic age rates of survival and reproduction (p(x) and b(x)).

## 1.2.1: Mortality risk at age x is parameterised with the Siler function (see [12] of the main manuscript).
siler <- function (alpha1, alpha2, beta1, beta2) {
  
  mort <- alpha1*exp((-beta1)*x) + alpha2*exp(beta2*x)
  return(mort)
}

## 1.2.2: The probability of survival for an individual aged x is equal to exp(-mu(x)). 
#This is computed in the loop below using the chosen parameters and stored in the vector 'EX1_age_specific_survival'.


EX1_age_specific_survival <- rep(0,EX1_omega)
for (i in 1:length(EX1_age_specific_survival)){
  
  x = EX1_age[i]
  EX1_age_specific_survival[i] <- exp(-siler(0.4,0.1,0.6,0)) 
  
}
#Set the survival rate of the last age class equal to 0
EX1_age_specific_survival[EX1_omega] <- 0
## 1.2.3: Rate of reproduction is modelled according to [14] of the main manuscript.
repro <- function(x, epsilon, phi, kappa){ #epsilon = age at maturity, kappa = age at which reproduction ceases
  
  if(x < epsilon | x > kappa){
    return(0)
  } else{
    return((x-epsilon)*exp(-phi*(x-epsilon)))
  }
}
# 1.2.4
#This is computed in the loop below using the chosen parameters and stored in the vector 'EX1_age_specific_rateofreproduction'.
EX1_age_specific_rateofreproduction <- rep(0, EX1_omega)
for (i in 1:length(EX1_age_specific_rateofreproduction)){
  
  x <- EX1_age[i]
  EX1_age_specific_rateofreproduction[i] <- repro(x, 15, 0.125, 40)
  
}


## 1.3: Using the background demography, we can quanitfy the asymptotic frequency of each age class under the assumpiton of a stationary population.

# 1.3.1: We first need to calculate lx, the probability of survival to age x. l(1) is assumed to be 1.
EX1_age_lx <- rep(0,EX1_omega)
EX1_age_lx[1] <- 1 
for (i in 2:length(EX1_age_lx)){
  x <- EX1_age[i]
  surv <- EX1_age_specific_survival[1:(x-1)]
  lx <- 1
    for (j in 1:length(surv)){
      lx <- lx*surv[j]
    }
  EX1_age_lx[i] <- lx
}
# 1.3.2: We can then calculate the asymptotic frequency for each age class (see [13] in the main manuscript).
#This is computed in the loop below and stored in the vector 'EX1_asymptotic_frequency'.
EX1_asymptotic_frequency <- rep(0,EX1_omega)
for (i in 1:length(EX1_asymptotic_frequency)){
  lx <- EX1_age_lx[i]
  EX1_asymptotic_frequency[i] <- lx/sum(EX1_age_lx)
}


### STAGE 2 - Sociality

# 2.1: We first define the propensity matrices (see Appendix C) which defines for an age class their relative (to other age classes) propensity to contribute to the survival and reproduction to every other age class
# 2.1.2: We do this first for survival and store the matrix as 'EX1_surv_social_propensity'
# The function EX1_help_surv populates the matrix. 
# In this example all post-reproductive age classes (ages 41-50) contribute to the survival of juveniles with the same relative propensity, which we set to 1. 
EX1_surv_social_propensity <- matrix(0, EX1_omega, EX1_omega)

EX1_help_surv <- function(surv_social_propensity, ages, epsilon, kappa){ #epsilon = age at maturity, kappa = age at which reproduction ceases
  for (i in 1:length(ages)){
    
    if(ages[i] < epsilon){
        surv_social_propensity[,ages[i]] <- 0} 
   
     if(ages[i] > epsilon & ages[i] <= kappa){
        surv_social_propensity[,ages[i]] <- 0}
    
    if(ages[i] > kappa){
        surv_social_propensity[(1:epsilon),ages[i]] <- 1} 
  }
  
return(surv_social_propensity)
  }

EX1_surv_social_propensity <- EX1_help_surv(EX1_surv_social_propensity, EX1_age, epsilon = 15, kappa = 40)


# 2.1.3 We then do the same for reproduction and store the matrix as 'EX1_repro_social_propensity'
# In this example we assume there are no social effects on reproduction
EX1_repro_social_propensity <- matrix(0, EX1_omega, EX1_omega)
EX1_help_repro <- function(repro_social_propensity, ages, epsilon, kappa){ 
  for (i in 1:length(ages)){
    
    if(ages[i] < epsilon){
      repro_social_propensity[,ages[i]] <- 0} 
    
    if(ages[i] > epsilon & ages[i] <= kappa){
      repro_social_propensity[,ages[i]] <- 0} 
    
    if(ages[i] > kappa){
      repro_social_propensity[,ages[i]] <- 0} 
  }
  
  return(repro_social_propensity)
}

EX1_repro_social_propensity <- EX1_help_repro(EX1_repro_social_propensity, EX1_age, epsilon = 15, kappa = 40)

# 2.2: The next step is to model the stripping away from the demographic rates (EX1_age_specific_survival and EX1_age_specific_rateofreproduction) the fractions that are due to the social environment and distribute them across age classes according to inclusive fitness methodology.
# 2.2.1: We first do this for survival using the stripping function and store in the vector 'EX1_DIRECT_age_specific_survival.
# In this example, we assume that 30% of each juvenile age classes survival rate is due to the social environment 

EX1_DIRECT_age_specific_survival <- rep(0, EX1_omega)
EX1_surv_strip <- function (stripped_surv, demo_surv, ages, epsilon, kappa) {
   for (i in 1:length(ages)){
    if(ages[i] <= epsilon){
      stripped_surv[ages[i]] <- demo_surv[ages[i]]*0.7    
    }
    if(ages[i] > epsilon & ages[i] <= kappa){
      stripped_surv[ages[i]] <- demo_surv[ages[i]]
    }
    if(ages[i] > kappa){
      stripped_surv[ages[i]] <- demo_surv[ages[i]]
    }
  }
  return(stripped_surv)
}
EX1_DIRECT_age_specific_survival <- EX1_surv_strip(EX1_DIRECT_age_specific_survival, EX1_age_specific_survival, EX1_age, epsilon = 15, kappa = 40)

#2.2.2: We then create and populate a vector that contains the number of genetic offspring equivalents that need to be distributed to other age classes.
# This vector is populated by subtracting the stripped survival rates from the demographic survival rates (see [2] in the main manuscript) and is stored in the vector 'EX1_survTransfersToDistribute'.
EX1_survTransfersToDistribute <- rep(0, EX1_omega)
EX1_survTransfersToDistribute <- EX1_age_specific_survival - EX1_DIRECT_age_specific_survival

#2.2.3: These offspring that exist due to the social environment are then distributed according to [C3] in Suppplementary Information Appendix C and using EX1_surv_social_propensity.
# The indirect genetic offspring contributions to survival then populate the matrix EX1_indirectSurvContributions
EX1_indirectSurvContributions <- matrix(0, EX1_omega, EX1_omega)

for (i in 1:length(EX1_age)){
  for (j in 1:length(EX1_age)){
    
    Stripped <- EX1_survTransfersToDistribute[j]
    fx <- EX1_asymptotic_frequency[i]
    Ax <- EX1_surv_social_propensity[j,i]
    
    AX <- EX1_surv_social_propensity[j,]
    
    SUM <- EX1_asymptotic_frequency*AX
    TOTAL <- sum(EX1_asymptotic_frequency*AX)
    
    EX1_indirectSurvContributions[j,i] <- Stripped*((fx*Ax)/TOTAL)
    if(is.nan(EX1_indirectSurvContributions[j,i])){EX1_indirectSurvContributions[j,i]<-0}
    
  }
}

# 2.2.4: We check all offspring have been distributed correctly and that there is no double accounting.
EX1_checkS <- matrix(0, EX1_omega, 3)
EX1_checkS[,1] <- EX1_age
EX1_checkS[,2] <- EX1_survTransfersToDistribute
for (i in 1:length(EX1_age)){
  EX1_checkS[i,3] <- sum(EX1_indirectSurvContributions[i,])
}
EX1_checkS <- as.data.frame(EX1_checkS) ## all offspring accounted

# 2.2.5: We then repeat stages 2.2.1-2.2.4 for social reproduction.
# In this example we assume that 0% of each reproductive age classes rate of reproduction is due to the social environment

EX1_DIRECT_age_specific_reproduction <- rep(0, EX1_omega)
EX1_repro_strip <- function (stripped_repro, demo_repro, ages, epsilon, kappa) {
  
  for (i in 1:length(ages)){
    
    if(ages[i] > epsilon & ages[i] <= kappa){
      stripped_repro[ages[i]] <- demo_repro[ages[i]]  
    }
    
  }
  return(stripped_repro)
}

EX1_DIRECT_age_specific_reproduction <- EX1_repro_strip(EX1_DIRECT_age_specific_reproduction, EX1_age_specific_rateofreproduction, EX1_age, epsilon = 15, kappa = 40)

EX1_reproTransfersToDistribute <- rep(0, EX1_omega)
EX1_reproTransfersToDistribute <- EX1_age_specific_rateofreproduction - EX1_DIRECT_age_specific_reproduction


EX1_indirectReproContributions <- matrix(0, EX1_omega, EX1_omega)

for (i in 1:length(EX1_age)){
  for (j in 1:length(EX1_age)){
    
    Stripped <- EX1_reproTransfersToDistribute[j]
    fx <- EX1_asymptotic_frequency[i]
    Ax <- EX1_repro_social_propensity[j,i]
    
    AX <- EX1_repro_social_propensity[j,]
    
    SUM <- EX1_asymptotic_frequency*AX
    TOTAL <- sum(EX1_asymptotic_frequency*AX)
    
    EX1_indirectReproContributions[j,i] <- Stripped*((fx*Ax)/TOTAL)
    if(is.nan(EX1_indirectReproContributions[j,i])){EX1_indirectReproContributions[j,i]<-0}
    
  }
}

EX1_checkR <- matrix(0, EX1_omega, 3)
EX1_checkR[,1] <- EX1_age
EX1_checkR[,2] <- EX1_reproTransfersToDistribute
for (i in 1:length(EX1_age)){
  EX1_checkR[i,3] <- sum(EX1_indirectReproContributions[i,])
}
EX1_checkR <- as.data.frame(EX1_checkR) ## In this example, we don't model any social effects on reproduction, so there are no offpspring to distribute. 



### STAGE 3 - Relatedness

# The quantification of relatedness follows steps that are outlined in detail in Appendix A of the Supplementary Information.

#3.1: First, we define the mean number of offspring produced by an individual across age classes.
# This is calculated by weighting the background demographic rate of reproduction for each age class by the asymptotic frequency of the age class.
# We store this value in the variable EX1_b_bar
EX1_b_bar <- sum(EX1_asymptotic_frequency*EX1_age_specific_rateofreproduction)

#3.2: Using this, we can then quantify, for each age class, the expected proportion of offspring after dispersal at the local patch that are 1) offspring (demographic) of an individual aged x, 2) offspring of other individuals on the patch, or 3) offspring that have dispersed from elsewhere
# We store these proportions in a dataframe using the function EX1_offspring_proportions.
# The function takes serveral parameters: here, we decide on the parameter values for N (number of individuals on the patch), d (dispersal rate), and c (cost of dispersal).
# The dataframe is then stored as EX1_offspringPostDispersal and is used for further calculations. 

EX1_offspring_proportions <- function(N, d, c, reproduction){
  
  offspring <- matrix(0, nrow = EX1_omega, ncol = 3)
  offspring <- as.data.frame(offspring)
  colnames(offspring) <- c("hx", "kx", "dispersed") 
  
  hx <- rep(0, EX1_omega)
  kx <- rep(0, EX1_omega)
  
  for (i in 1:length(hx)){
    
    hx_term1 <- reproduction[i]*(1-d)
    hx_term2 <- reproduction[i]*(1-d) + (N-1)*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c)
    hx[i] <- hx_term1/hx_term2
    
    kx_term1 <- (N-1)*EX1_b_bar*(1-d) 
    kx_term2 <- reproduction[i]*(1-d) + (N-1)*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c)
    kx[i] <- kx_term1/kx_term2
    
  }
  
  offspring$hx <- hx
  offspring$kx <- kx
  offspring$dispersed <- 1 - hx - kx
  
  return(offspring)
}

#Set parameter values for N, d and c
N <- 4
d <- 0.5 
c <- 0
EX1_offspringPostDispersal <- EX1_offspring_proportions(N, d, c, EX1_age_specific_rateofreproduction)

# 3.3: We quantify the mean proportions of offspring after dispersal that are the demographic offspring of a focal individual (h_bar) or the offspring of other individuals on the patch (k_bar)

hx <- EX1_offspringPostDispersal$hx
h_bar <- sum(EX1_asymptotic_frequency*hx)
kx <- EX1_offspringPostDispersal$kx
k_bar <- sum(EX1_asymptotic_frequency*kx)

# 3.4: We can then calculate rhat(1) 
# With this, we can then calculate numerically, for each age x, the relatedness of an individual aged x to another random breeder on their patch 
# We store these relatedness values in the vector EX1_rhatx

EX1_rhatx <- rep(0, EX1_omega)

EX1_rhatx1_term1 <- EX1_asymptotic_frequency[1]*(1-d)*(1-d)*(h_bar^2)
EX1_rhatx1_term2 <- sum(EX1_asymptotic_frequency[2:EX1_omega]*hx[1:(EX1_omega-1)])
EX1_rhatx1_term3 <- EX1_asymptotic_frequency[1]*(1-d)*(1-d)*(1-(h_bar^2))
EX1_rhatx1_term4 <- sum(EX1_asymptotic_frequency[2:EX1_omega]*kx[1:(EX1_omega-1)])

EX1_rhatx[1] <- ((EX1_rhatx1_term1 + EX1_rhatx1_term2)/(1-(EX1_rhatx1_term3+EX1_rhatx1_term4)))

for (i in 1:length(EX1_rhatx)){
  x <- EX1_age[i]
  
  younger <- EX1_age[1:(x-1)]
  sameage <- x
  older <- EX1_age[(x+1):EX1_omega]
  
  if (x != 1){
    r_younger <- rep(0, length(younger))
    for (ageClass in 1:length(r_younger)){
      y <- younger[ageClass]
      r_younger[ageClass] <- EX1_asymptotic_frequency[y]*((1-d)*(hx[x-y] + kx[x-y]*EX1_rhatx[x-y]))
    }
  }else{ r_younger <- 0}
 
  r_sameageT1 <- EX1_asymptotic_frequency[x]*(1-d)*(1-d)
  r_sameageT2 <- (h_bar^2) + EX1_rhatx[1]*(1-(h_bar^2))
  r_sameage <- r_sameageT1*r_sameageT2
  
  if(x!=EX1_omega){
  r_older <- rep(0, length(older))
  for (ageClass in 1:length(r_older)){
    y <- older[ageClass]
    r_older[ageClass] <- EX1_asymptotic_frequency[y]*((1-d)*(hx[y-x] + kx[y-x]*EX1_rhatx[1]))
  }
  }else{r_older <- 0}
  
  if (x!=1){
  EX1_rhatx[i] <- sum(r_younger) + r_sameage + sum(r_older)}
}

# 3.5: We can then also quantify the relatedness of a focal individual aged x to another individual on the patch including itself
EX1_rx <- rep(0,EX1_omega)
for (i in 1:length(EX1_rx)){
  EX1_rx[i] <- (1/N) + ((N-1)/N)*EX1_rhatx[i]
}

# 3.6: We can then plot the age-specific relatedness of a individual to other individuals on the patch using a relatedness dataframe.

EX1_relatedness <- data.frame("x" = EX1_age, "rhatx" = EX1_rhatx, "r" = EX1_rx)
EX1_r_plot <- ggplot(EX1_relatedness, aes(x, rhatx))
EX1_r_plot <- EX1_r_plot + geom_line(col = "black", size = 2) + labs(x = "Age", y = "Relatedness to patch members") + theme_m
EX1_r_plot 
ggsave(EX1_r_plot, file = "Fig2D.pdf",height = 5, width = 8)

### STAGE 4 - Computation of the fitness matrix (W)
# Here, the overall goal is to sum together the different components of inclusive fitness using sections 1-3 to create a fitness matrix of which the elements represent age-specific inclusive fitness contributions - [] in the main manuscript.

# 4.1: Create an empty matrix to populate
EX1_W <- matrix(0, EX1_omega, EX1_omega)

# 4.2: The top row of the fitness matrix represents inclusive fitness contributions of each age x (column) to the offspring class.
# This is a sum of the direct (stripped demographic rate of reproduction) and expected indirect genetic contributions for each age x.

#4.2.1: First, we can extract the indirect genetic offspring contributions of age class x, where each contribution is weighted by the expected relatedness between an individual aged x and an offspring born to a random breeder on the patch.
# We store indirect contributions in the vector EX1_age_specific_indirect_mx

EX1_age_specific_indirect_mx <- rep(0,EX1_omega)
for (i in 1:length(EX1_age_specific_indirect_mx)){
  
  x <- EX1_age[i]
  offspring <- EX1_indirectReproContributions[,i]
  relatedness <- EX1_rhatx[i]
  
  EX1_age_specific_indirect_mx[i] <- sum(offspring*relatedness)
  
}


#4.2.2: We can then quantify effective inclusive fitness (expected genetic offspring equivalents that establish onto a patch).
# For this we need the average survival rate of a random breeder on the patch, which we store in the variable EX1_p_bar
EX1_p_bar <- sum(EX1_asymptotic_frequency*EX1_age_specific_survival)
# We can then quanitfy and store effective inclusive fitness in the vector EX1_age_specific_effectivemx
EX1_age_specific_effectivemx <- rep(0, EX1_omega)

for (i in 1:length(EX1_age_specific_effectivemx)){
  
  x <- EX1_age[i]
  
  Fx <- EX1_DIRECT_age_specific_reproduction[x] + EX1_age_specific_indirect_mx[x]
  
  gx_term1 <- (1 - EX1_age_specific_survival[x]) + (N-1)*(1 - EX1_p_bar)
  gx_term2 <- EX1_age_specific_rateofreproduction[x]*(1-d) + (N-1)*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c)
  gx <- gx_term1/gx_term2
  
  g_bar <- N*(1 - EX1_p_bar)/(N*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c))
  
  effective_mx <- Fx*(1-d)*gx + Fx*(1-c)*d*g_bar
  
  EX1_age_specific_effectivemx[i] <- effective_mx
  
}  
#We can then populate the top row of the fitness matrix (W) with expected inclusive fitness contributions to the offspring class for each age x.
EX1_W[1,] <- EX1_age_specific_effectivemx

#4.3: The sub-diagonal of the fitness matrix W represents inclusive fitness contributions of individuals aged x to their own survival (stripped survival). 
EX1_survW <- matrix(0, EX1_omega, EX1_omega)
for (i in 1:EX1_omega){
  for (j in 1:EX1_omega){
    
    if(i == j+1){
      
      age1 <- EX1_age[j]
      survp <- EX1_DIRECT_age_specific_survival[age1]
      EX1_survW[i,j] <- survp
    }
  }
}
#4.3.1: We can then populate the sub-diagonal of W using the vector EX1_survW
EX1_W <- EX1_W + EX1_survW

#4.4: Finally, we can populate all other elements of W that represent age class x's inclusive fitness contributions to the survival of other age classes.
EX1_survTransferW <- matrix(0, EX1_omega, EX1_omega)
for (i in 1:EX1_omega){
  for (j in 1:EX1_omega){
    
    if(j == 1 | j == i+1){
      EX1_survTransferW[j,i] <- 0
    }
    
    else{
      x <- EX1_age[i]
      y <- EX1_age[j]
      
      survTransfer <- EX1_indirectSurvContributions[j,i]
      relatedness <- EX1_rhatx[i]
      
      EX1_survTransferW[j,i] <- survTransfer*relatedness
      
    }
  }
}
##4.4.1: We can then populate all remaining elements of W with the vector EX1_survTransferW
EX1_W <- EX1_W + EX1_survTransferW

### STAGE 5 - Forces of selection.
# We now have everything we need to compute forces of selection according to [8] and [10] in the main manuscript.
# Below, we compute forces of selection acting on a mutant allele that 1) alters survival rate at age x ([8]), and 2) alters rate of reproduction at age x ([10]).
# After deriving, we store the forces of selection in the vectors EX1_FOS_px and EX1_FOS_bx


# 5.1: First
EX1_asymptotic_frequency_W <- eigen.analysis(EX1_W)$stable.stage
EX1_inclusive_repro_value <- eigen.analysis(EX1_W)$repro.value 

EX1_mx <- rep(0, EX1_omega)
for (i in 1:length(EX1_mx)){
  
  EX1_mx[i] = EX1_age_specific_rateofreproduction[i]*((1-EX1_p_bar)/EX1_b_bar) 
}


# We can then calculate Hamilton's scalar, and the scalar vector of the modified asymptotic frequency and inclusive reproductive value
Ham_scalar_T <- sum(EX1_age_lx*EX1_mx*EX1_age)
scalar <- sum(EX1_asymptotic_frequency_W*EX1_inclusive_repro_value)


# 5.2: Calculating the force of selection on survival rate at age x according to [8] in the main manuscript

# 5.2.1 We first need to quantify hdotx and kdotx, that represent the proportions of offspring at the local patch after dispersal that are the expected direct and indirect genetic offspring of an individual aged x
EX1_hdotx <- rep(0, EX1_omega)
for (i in 1:length(EX1_hdotx)){
  
  T1 <- EX1_W[1,i]*(1-d)
  T2 <- EX1_age_specific_rateofreproduction[i]*(1-d) + (N-1)*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c)
  EX1_hdotx[i] <- T1/T2
}

EX1_kdotx <- rep(0, EX1_omega)
for (i in 1:length(EX1_kdotx)){
  Fbar <- sum(EX1_asymptotic_frequency*EX1_W[1,])
  
  T1 <- (N-1)*Fbar*(1-d)
  T2 <- EX1_age_specific_rateofreproduction[i]*(1-d) + (N-1)*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c)
  EX1_kdotx[i] <- T1/T2
  
}


# 5.2.2 We then calculate the age-specific force of selection on survival at age x and store in the vector EX1_FOS_px
EX1_FOS_px <- rep(0,EX1_omega)
EX1_FOS_px_Ham <- rep(0, EX1_omega)
for (i in 1:length(EX1_FOS_px)){
  
  A <- EX1_hdotx[i]
  r <- EX1_rhatx[i]
  B <- EX1_kdotx[i]
  
  term1 <- (EX1_asymptotic_frequency_W[i]*(EX1_inclusive_repro_value[1]*(A + r*B)))
  irvnext <- EX1_inclusive_repro_value[i+1] 
  if (is.na(irvnext)){irvnext <- 0}
  
  EX1_FOS_px[i] <- ((EX1_asymptotic_frequency_W[i]*irvnext) - term1)/scalar
  
  if (i == EX1_omega){next}
  future_lx <- EX1_age_lx[(i+1):EX1_omega]
  future_mx <- EX1_mx[(i+1):EX1_omega]
  
  EX1_FOS_px_Ham[i] <- (sum(future_lx*future_mx))/Ham_scalar_T
  
}


# 5.2.3: We then calculate the age-specific force of selection on reproduction at age x and store in the vector EX1_FOS_bx

EX1_FOS_bx <- rep(0,EX1_omega)
EX1_FOS_bx_Ham <- rep(0, EX1_omega)
for (i in 1:length(EX1_FOS_bx)){
  
  term1 <- (EX1_asymptotic_frequency_W[i]*EX1_inclusive_repro_value[1])/scalar
  
  gx_term1 <- (1 - EX1_age_specific_survival[i]) + (N-1)*(1 - EX1_p_bar)
  gx_term2 <- EX1_age_specific_rateofreproduction[i]*(1-d) + (N-1)*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c)
  gx <- gx_term1/gx_term2
  gbar <- N*(1 - EX1_p_bar)/(N*EX1_b_bar*(1-d) + N*EX1_b_bar*d*(1-c))
  hx <- EX1_offspringPostDispersal$hx[i]
  Ix <- (EX1_age_specific_indirect_mx[i]*(1-d))/gx_term2 
  r <- EX1_rhatx[i]
  kdotx <- EX1_kdotx[i]
  
  term2 <- ((1-d)*gx)*((1-hx)-Ix-(r*kdotx))
  
  term3 <- (1-c)*d*gbar
  
  EX1_FOS_bx[i] <-term1*(term2 + term3)
  
  EX1_FOS_bx_Ham[i] <- EX1_age_lx[i]/Ham_scalar_T
  
}



EX1_forceofselection <- data.frame("x" = EX1_age, "px" = EX1_FOS_px, "bx" = EX1_FOS_bx, "pxHam" = EX1_FOS_px_Ham, "bxHam" = EX1_FOS_bx_Ham)
Fig_FOS_px <- ggplot(EX1_forceofselection, aes(x, px))
Fig_FOS_px <- Fig_FOS_px + geom_line(col="darkblue", size = 1.5, linetype = "longdash") + labs(x="Age", y ="Force of selection on survival") + theme_m
Fig_FOS_px <- Fig_FOS_px + geom_hline(yintercept=0, linetype="dashed", color = "black") + geom_vline(xintercept = 40, linetype="dashed", color="black")
Fig_FOS_px <- Fig_FOS_px + geom_line(aes(x, pxHam), col = "darkblue", size = 1.5)
Fig_FOS_px

Fig_FOS_bx <- ggplot(EX1_forceofselection, aes(x, bx))
Fig_FOS_bx <- Fig_FOS_bx + geom_line(col="red", size = 1.5, linetype = "longdash") + labs(x="Age", y ="Force of selection on reproduction") + theme_m
Fig_FOS_bx <- Fig_FOS_bx + geom_line(aes(x, bxHam), col = "red", size = 1.5)
Fig_FOS_bx <- Fig_FOS_bx + geom_hline(yintercept=0, linetype="dashed", color = "black")
Fig_FOS_bx

FIG2C <- grid.arrange(Fig_FOS_px, Fig_FOS_bx, ncol = 2)
ggsave(FIG2C, file = "Fig2C.pdf", width = 8, height = 5 )


### DEMOGRAHPY PLOTS
surv_transfers = rep(0, EX1_omega)
for (i in 1:EX1_omega){
  
  surv_transfers[i] <- sum(EX1_survTransferW[,i])
  
}

repro_transfers = rep(0, EX1_omega)
for (i in 1:EX1_omega){
  
  repro_transfers[i] <- sum(EX1_indirectReproContributions[,i])
  
}


EX1_age_specific_survival[50] <- EX1_age_specific_survival[49] #so the plot doesn't drastically go to zero / doesn't impact calculations
EX1_DIRECT_age_specific_survival[50] <- EX1_DIRECT_age_specific_survival[49]

EX1_background_demography <- data.frame("x"= EX1_age, "px" = EX1_age_specific_survival, "pdotx" = EX1_DIRECT_age_specific_survival, "bdotx" = EX1_DIRECT_age_specific_reproduction, "mx" = EX1_mx, "eff_mx" = EX1_age_specific_effectivemx, "surv_transfers" = surv_transfers, "repro_transfers" = repro_transfers)


EX1_surv_px <- ggplot(EX1_background_demography, aes(x, px))
EX1_surv_px <- EX1_surv_px + geom_line(col = "darkblue", size = 2) + labs(x = "Age", y = "Survival") + theme_m
EX1_surv_px <- EX1_surv_px + geom_line(aes(x, pdotx), size = 2, linetype = "dotted", col = "darkblue")

EX1_surv_transfers <- ggplot(EX1_background_demography, aes(x, surv_transfers))
EX1_surv_transfers <- EX1_surv_transfers + geom_line(col = "darkblue", size = 2) + labs(x = "Age", y = "Survival transfers") + theme_m
EX1_surv_transfers

EX1_repro <- ggplot(EX1_background_demography, aes(x, mx))
EX1_repro <- EX1_repro  + geom_line(col = "red", size = 2) + labs(x = "Age", y = "Reproduction") + theme_m
EX1_repro <- EX1_repro + geom_line(aes(x, eff_mx), size = 2, linetype ="dotted", col ="red")


EX1_repro_transfers <- ggplot(EX1_background_demography, aes(x, repro_transfers))
EX1_repro_transfers <- EX1_repro_transfers+ geom_line(col = "red", size = 2) + labs(x = "Age", y = "Reproduction transfers") + theme_m
EX1_repro_transfers 


Fig2B <- grid.arrange(EX1_surv_px, EX1_surv_transfers, EX1_repro, EX1_repro_transfers, nrow = 2, ncol = 2)
ggsave(Fig2B, file = "Fig2B.pdf", height = 5, width = 8)

########################################################################################################################

#################
##  EXAMPLE 2  ##
#################
rm(list=ls())
#ggplot theme to reproduce figures

theme_m <-  theme(panel.background=element_blank(), 
                  strip.background=element_blank(), 
                  axis.line=element_line("black"), 
                  axis.ticks=element_line("black"), 
                  axis.text=element_text(colour="black", size=18), 
                  axis.title=element_text(size=20),
                  panel.grid=element_blank(),
                  strip.text=element_text(size=10),
                  legend.key=element_rect(fill="white"), 
                  axis.line.x=element_line("black"),
                  axis.line.y=element_line("black"), 
                  plot.title = element_text(color="#666666", face="bold", size=14))

### STAGE 1 - Age classes and background demography

## 1.1: Define the number of age classes and store age classes in the vector EX2_age which runs from 1 to the last age class (omega) and which all individuals die.
EX2_omega <- 20 #last age class
EX2_age <- c(1:EX2_omega) 

## 1.2: Define demographic age rates of survival and reproduction (p(x) and b(x)).

## 1.2.1: Mortality risk at age x is parameterised with the Siler function (see [11] of the main manuscript.
siler <- function (alpha1, alpha2, beta1, beta2) {
  
  mort <- alpha1*exp((-beta1)*x) + alpha2*exp(beta2*x)
  return(mort)
}

## 1.2.2: The probability of survival for an individual aged x is equal to exp(-mu(x)). 
#This is computed in the loop below using the chosen parameters and stored in the vector 'EX2_age_specific_survival'.

EX2_age_specific_survival <- rep(0,EX2_omega)
for (i in 1:length(EX2_age_specific_survival)){
  
  x = EX2_age[i]
  EX2_age_specific_survival[i] <- exp(-siler(0.4,0.1,0.6,0)) 
  
}
#Set the survival rate of the last age class equal to 0
EX2_age_specific_survival[EX2_omega] <- 0
## 1.2.3: Rate of reproduction is modelled according to [13] of the main manuscript.
repro <- function(x, epsilon, phi, kappa){ #epsilon = age at maturity, kappa = age at which reproduction ceases
  
  if(x < epsilon | x > kappa){
    return(0)
  } else{
    return((x-epsilon)*exp(-phi*(x-epsilon)))
  }
}
# 1.2.4
#This is computed in the loop below using the chosen parameters and stored in the vector 'EX2_age_specific_rateofreproduction'.
EX2_age_specific_rateofreproduction <- rep(0, EX2_omega)
for (i in 1:length(EX2_age_specific_rateofreproduction)){
  
  x <- EX2_age[i]
  EX2_age_specific_rateofreproduction[i] <- repro(x, 5, 0.2, 21)
  
}
plot(EX2_age, EX2_age_specific_rateofreproduction)

## 1.3: Using the background demography, we can quanitfy the asymptotic frequency of each age class under the assumpiton of a stationary population.

# 1.3.1: We first need to calculate lx, the probability of survival to age x. l(1) is assumed to be 1.
EX2_age_lx <- rep(0,EX2_omega)
EX2_age_lx[1] <- 1 
for (i in 2:length(EX2_age_lx)){
  x <- EX2_age[i]
  surv <- EX2_age_specific_survival[1:(x-1)]
  lx <- 1
  for (j in 1:length(surv)){
    lx <- lx*surv[j]
  }
  EX2_age_lx[i] <- lx
}
# 1.3.2: We can then calculate the asymptotic frequency for each age class (see [12] in the main manuscript).
#This is computed in the loop below and stored in the vector 'EX2_asymptotic_frequency'.
EX2_asymptotic_frequency <- rep(0,EX2_omega)
for (i in 1:length(EX2_asymptotic_frequency)){
  lx <- EX2_age_lx[i]
  EX2_asymptotic_frequency[i] <- lx/sum(EX2_age_lx)
}


### STAGE 2 - Sociality

# 2.1: We first define the propensity matrices (see Appendix C) which defines for an age class their relative (to other age classes) propensity to contribute to the survival and reproduction to every other age class
# 2.1.2: We do this first for survival and store the matrix as 'EX2_surv_social_propensity'
# The function EX2_help_surv populates the matrix. 
# In this example no age classes contribute to the survival of other age classes.
EX2_surv_social_propensity <- matrix(0, EX2_omega, EX2_omega)

EX2_help_surv <- function(surv_social_propensity, ages, epsilon, kappa){ #epsilon = age at maturity, kappa = age at which reproduction ceases
  for (i in 1:length(ages)){
    
    if(ages[i] < epsilon){
      surv_social_propensity[,ages[i]] <- 0} 
    
    if(ages[i] > epsilon & ages[i] <= kappa){
      surv_social_propensity[,ages[i]] <- 0}
    
    if(ages[i] > kappa){
      surv_social_propensity[ages[i]] <- 0} 
  }
  
  return(surv_social_propensity)
}

EX2_surv_social_propensity <- EX2_help_surv(EX2_surv_social_propensity, EX2_age, epsilon = 5, kappa = 21)


# 2.1.3 We then do the same for reproduction and store the matrix as 'EX2_repro_social_propensity'
# In this example we assume that each pre-reproductive age class has the same relative propensity (1) to contribute to the reproduction of reproductive age classes
EX2_repro_social_propensity <- matrix(0, EX2_omega, EX2_omega)
EX2_help_repro <- function(repro_social_propensity, ages, epsilon, kappa){ 
  for (i in 1:length(ages)){
    
    if(ages[i] <= epsilon){
      repro_social_propensity[((epsilon+1):(kappa-1)),ages[i]] <- 1} 
    
    if(ages[i] > epsilon & ages[i] <= kappa){
      repro_social_propensity[,ages[i]] <- 0} 
    
    if(ages[i] > kappa){
      repro_social_propensity[,ages[i]] <- 0} 
  }
  
  return(repro_social_propensity)
}

EX2_repro_social_propensity <- EX2_help_repro(EX2_repro_social_propensity, EX2_age, epsilon = 5, kappa = 21)

# 2.2: The next step is to model the stripping away from the demographic rates (EX2_age_specific_survival and EX2_age_specific_rateofreproduction) the fractions that are due to the social environment and distribute them across age classes according to inclusive fitness methodology.
# 2.2.1: We first do this for survival and store in the vector 'EX2_DIRECT_age_specific_survival.
# In this example no age classes contribute to the survival of other age classes.

EX2_DIRECT_age_specific_survival <- rep(0, EX2_omega)
EX2_surv_strip <- function (stripped_surv, demo_surv, ages, epsilon, kappa) {
  for (i in 1:length(ages)){
    if(ages[i] <= epsilon){
      stripped_surv[ages[i]] <- demo_surv[ages[i]]   
    }
    if(ages[i] > epsilon & ages[i] <= kappa){
      stripped_surv[ages[i]] <- demo_surv[ages[i]]
    }
    if(ages[i] > kappa){
      stripped_surv[ages[i]] <- demo_surv[ages[i]]
    }
  }
  return(stripped_surv)
}
EX2_DIRECT_age_specific_survival <- EX2_surv_strip(EX2_DIRECT_age_specific_survival, EX2_age_specific_survival, EX2_age, epsilon = 5, kappa = 21)

#2.2.2: We then create and populate a vector that contains the number of genetic offspring equivalents that need to be distributed to other age classes.
# This vector is populated by subtracting the stripped survival rates from the demographic survival rates (see [2] in the main manuscript) and is stored in the vector 'EX1_survTransfersToDistribute'.
EX2_survTransfersToDistribute <- rep(0, EX2_omega)
EX2_survTransfersToDistribute <- EX2_age_specific_survival - EX2_DIRECT_age_specific_survival

#2.2.3: These offspring that exist due to the social environment are then distributed according to [C3] in Suppplementary Information Appendix C and using EX1_surv_social_propensity.
# The indirect genetic offspring contributions to survival then populate the matrix EX1_indirectSurvContributions
EX2_indirectSurvContributions <- matrix(0, EX2_omega, EX2_omega)

for (i in 1:length(EX2_age)){
  for (j in 1:length(EX2_age)){
    
    Stripped <- EX2_survTransfersToDistribute[j]
    fx <- EX2_asymptotic_frequency[i]
    Ax <- EX2_surv_social_propensity[j,i]
    
    AX <- EX2_surv_social_propensity[j,]
    
    SUM <- EX2_asymptotic_frequency*AX
    TOTAL <- sum(EX2_asymptotic_frequency*AX)
    
    EX2_indirectSurvContributions[j,i] <- Stripped*((fx*Ax)/TOTAL)
    if(is.nan(EX2_indirectSurvContributions[j,i])){EX2_indirectSurvContributions[j,i]<-0}
    
  }
}

# 2.2.4: We check all offspring have been distributed correctly and that there is no double accounting.
EX2_checkS <- matrix(0, EX2_omega, 3)
EX2_checkS[,1] <- EX2_age
EX2_checkS[,2] <- EX2_survTransfersToDistribute
for (i in 1:length(EX2_age)){
  EX2_checkS[i,3] <- sum(EX2_indirectSurvContributions[i,])
}
EX2_checkS <- as.data.frame(EX2_checkS) ## In this example, we don't model any social effects on survival, so there are no offpspring to distribute.

# 2.2.5: We then repeat stages 2.2.1-2.2.4 for social reproduction.
# In this example we assume that 30% of each reproductive age classes rate of reproduction is due to the social environment

EX2_DIRECT_age_specific_reproduction <- rep(0, EX2_omega)
EX2_repro_strip <- function (stripped_repro, demo_repro, ages, epsilon, kappa) {
  
  for (i in 1:length(ages)){
    
    if(ages[i] > epsilon & ages[i] <= kappa){
      stripped_repro[ages[i]] <- demo_repro[ages[i]]*0.7  
    }
    
  }
  return(stripped_repro)
}

EX2_DIRECT_age_specific_reproduction <- EX2_repro_strip(EX2_DIRECT_age_specific_reproduction, EX2_age_specific_rateofreproduction, EX2_age, epsilon = 5, kappa = 21)

EX2_reproTransfersToDistribute <- rep(0, EX2_omega)
EX2_reproTransfersToDistribute <- EX2_age_specific_rateofreproduction - EX2_DIRECT_age_specific_reproduction


EX2_indirectReproContributions <- matrix(0, EX2_omega, EX2_omega)

for (i in 1:length(EX2_age)){
  for (j in 1:length(EX2_age)){
    
    Stripped <- EX2_reproTransfersToDistribute[j]
    fx <- EX2_asymptotic_frequency[i]
    Ax <- EX2_repro_social_propensity[j,i]
    
    AX <- EX2_repro_social_propensity[j,]
    
    SUM <- EX2_asymptotic_frequency*AX
    TOTAL <- sum(EX2_asymptotic_frequency*AX)
    
    EX2_indirectReproContributions[j,i] <- Stripped*((fx*Ax)/TOTAL)
    if(is.nan(EX2_indirectReproContributions[j,i])){EX2_indirectReproContributions[j,i]<-0}
    
  }
}

EX2_checkR <- matrix(0, EX2_omega, 3)
EX2_checkR[,1] <- EX2_age
EX2_checkR[,2] <- EX2_reproTransfersToDistribute
for (i in 1:length(EX2_age)){
  EX2_checkR[i,3] <- sum(EX2_indirectReproContributions[i,])
}
EX2_checkR <- as.data.frame(EX2_checkR) ## All offspring accounted.

### STAGE 3 - Relatedness

# The quantification of relatedness follows steps that are outlined in detail in Appendix A of the Supplementary Information.

#3.1: First, we define the mean number of offspring produced by an individual across age classes.
# This is calculated by weighting the background demographic rate of reproduction for each age class by the asymptotic frequency of the age class.
# We store this value in the variable EX2_b_bar
EX2_b_bar <- sum(EX2_asymptotic_frequency*EX2_age_specific_rateofreproduction)

#3.2: Using this, we can then quantify, for each age class, the expected proportion of offspring after dispersal at the local patch that are 1) offspring (demographic) of an individual aged x, 2) offspring of other individuals on the patch, or 3) offspring that have dispersed from elsewhere
# We store these proportions in a dataframe using the function EX2_offspring_proportions.
# The function takes serveral parameters: here, we decide on the parameter values for N (number of individuals on the patch), d (dispersal rate), and c (cost of dispersal).
# The dataframe is then stored as EX2_offspringPostDispersal and is used for further calculations. 

EX2_offspring_proportions <- function(N, d, c, reproduction){
  
  offspring <- matrix(0, nrow = EX2_omega, ncol = 3)
  offspring <- as.data.frame(offspring)
  colnames(offspring) <- c("hx", "kx", "dispersed") 
  
  hx <- rep(0, EX2_omega)
  kx <- rep(0, EX2_omega)
  
  for (i in 1:length(hx)){
    
    hx_term1 <- reproduction[i]*(1-d)
    hx_term2 <- reproduction[i]*(1-d) + (N-1)*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c)
    hx[i] <- hx_term1/hx_term2
    
    kx_term1 <- (N-1)*EX2_b_bar*(1-d) 
    kx_term2 <- reproduction[i]*(1-d) + (N-1)*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c)
    kx[i] <- kx_term1/kx_term2
    
  }
  
  offspring$hx <- hx
  offspring$kx <- kx
  offspring$dispersed <- 1 - hx - kx
  
  return(offspring)
}

N <- 4
d <- 0.5 
c <- 0
EX2_offspringPostDispersal <- EX2_offspring_proportions(N, d, c, EX2_age_specific_rateofreproduction)

# 3.3: We quantify the mean proportions of offspring after dispersal that are the demographic offspring of a focal individual (h_bar) or the offspring of other individuals on the patch (k_bar)

EX2_hx <- EX2_offspringPostDispersal$hx
EX2_h_bar <- sum(EX2_asymptotic_frequency*EX2_hx)
EX2_kx <- EX2_offspringPostDispersal$kx
EX2_k_bar <- sum(EX2_asymptotic_frequency*EX2_kx)

# 3.4: We can then calculate rhat(1) according Appendix A
# With this, we can then calculate numerically, for each age x, the relatedness of an individual aged x to another random breeder on their patch

EX2_rhatx <- rep(0, EX2_omega)

EX2_rhatx1_term1 <- EX2_asymptotic_frequency[1]*(1-d)*(1-d)*(EX2_h_bar^2)
EX2_rhatx1_term2 <- sum(EX2_asymptotic_frequency[2:EX2_omega]*EX2_hx[1:(EX2_omega-1)])
EX2_rhatx1_term3 <- EX2_asymptotic_frequency[1]*(1-d)*(1-d)*(1-(EX2_h_bar^2))
EX2_rhatx1_term4 <- sum(EX2_asymptotic_frequency[2:EX2_omega]*EX2_kx[1:(EX2_omega-1)])

EX2_rhatx[1] <- ((EX2_rhatx1_term1 + EX2_rhatx1_term2)/(1-(EX2_rhatx1_term3+EX2_rhatx1_term4)))

for (i in 1:length(EX2_rhatx)){
  x <- EX2_age[i]
  
  younger <- EX2_age[1:(x-1)]
  sameage <- x
  older <- EX2_age[(x+1):EX2_omega]
  
  if (x != 1){
    r_younger <- rep(0, length(younger))
    for (ageClass in 1:length(r_younger)){
      y <- younger[ageClass]
      r_younger[ageClass] <- EX2_asymptotic_frequency[y]*((1-d)*(EX2_hx[x-y] + EX2_kx[x-y]*EX2_rhatx[x-y]))
    }
  }else{ r_younger <- 0}
  
  r_sameageT1 <- EX2_asymptotic_frequency[x]*(1-d)*(1-d)
  r_sameageT2 <- (EX2_h_bar^2) + EX2_rhatx[1]*(1-(EX2_h_bar^2))
  r_sameage <- r_sameageT1*r_sameageT2
  
  if(x!=EX2_omega){
    r_older <- rep(0, length(older))
    for (ageClass in 1:length(r_older)){
      y <- older[ageClass]
      r_older[ageClass] <- EX2_asymptotic_frequency[y]*((1-d)*(EX2_hx[y-x] + EX2_kx[y-x]*EX2_rhatx[1]))
    }
  }else{r_older <- 0}
  
  if (x!=1){
    EX2_rhatx[i] <- sum(r_younger) + r_sameage + sum(r_older)}
}

# 3.5: We can then also quantify the relatedness of a focal individual aged x to another individual on the patch including itself
EX2_rx <- rep(0,EX2_omega)
for (i in 1:length(EX2_rx)){
  EX2_rx[i] <- (1/N) + ((N-1)/N)*EX2_rhatx[i]
}

# 3.6: We can then plot the age-specific relatedness of a individual to other individuals on the patch using a relatedness dataframe.

EX2_relatedness <- data.frame("x" = EX2_age, "rhatx" = EX2_rhatx, "r" = EX2_rx)
EX2_r_plot <- ggplot(EX2_relatedness, aes(x, rhatx))
EX2_r_plot <- EX2_r_plot + geom_line(col = "black", size = 2) + labs(x = "Age", y = "Relatedness") + theme_m
EX2_r_plot 
ggsave(EX2_r_plot, file = "Fig3d.pdf",height = 5, width = 8)

### STAGE 4 - Computation of the fitness matrix (W)
# Here, the overall goal is to sum together the different components of inclusive fitness using sections 1-3 to create a fitness matrix of which the elements represent age-specific inclusive fitness contributions

# 4.1: Create an empty matrix to populate
EX2_W <- matrix(0, EX2_omega, EX2_omega)

# 4.2: The top row of the fitness matrix represents inclusive fitness contributions of each age x (column) to the offspring class.
# This is a sum of the direct (stripped demographic rate of reproduction) and expected indirect genetic contributions for each age x.

#4.2.1: First, we can extract the indirect genetic offspring contributions of age class x, where each contribution is weighted by the expected relatedness between an individual aged x and an offspring born to a random breeder on the patch.
# We store indirect contributions in the vector EX2_age_specific_indirect_mx

EX2_age_specific_indirect_mx <- rep(0,EX2_omega)
for (i in 1:length(EX2_age_specific_indirect_mx)){
  
  x <- EX2_age[i]
  offspring <- EX2_indirectReproContributions[,i]
  relatedness <- EX2_rhatx[i]
  
  EX2_age_specific_indirect_mx[i] <- sum(offspring*relatedness)
  
}


#4.2.2: We can then quantify effective inclusive fitness (expected genetic offspring equivalents that establish onto a patch) using [5] in the main manuscript.
# For this we need the average survival rate of a random breeder on the patch, which we store in the variable EX2_p_bar
EX2_p_bar <- sum(EX2_asymptotic_frequency*EX2_age_specific_survival)
# We can then quanitfy and store effective inclusive fitness in the vector EX1_age_specific_effectivemx
EX2_age_specific_effectivemx <- rep(0, EX2_omega)

for (i in 1:length(EX2_age_specific_effectivemx)){
  
  x <- EX2_age[i]
  
  Fx <- EX2_DIRECT_age_specific_reproduction[x] + EX2_age_specific_indirect_mx[x]
  
  gx_term1 <- (1 - EX2_age_specific_survival[x]) + (N-1)*(1 - EX2_p_bar)
  gx_term2 <- EX2_age_specific_rateofreproduction[x]*(1-d) + (N-1)*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c)
  gx <- gx_term1/gx_term2
  
  g_bar <- N*(1 - EX2_p_bar)/(N*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c))
  
  effective_mx <- Fx*(1-d)*gx + Fx*(1-c)*d*g_bar
  
  EX2_age_specific_effectivemx[i] <- effective_mx
  
}  
#We can then populate the top row of the fitness matrix (W) with expected inclusive fitness contributions to the offspring class for each age x.
EX2_W[1,] <- EX2_age_specific_effectivemx

#4.3: The sub-diagonal of the fitness matrix W represents inclusive fitness contributions of individuals aged x to their own survival (stripped survival). 
EX2_survW <- matrix(0, EX2_omega, EX2_omega)
for (i in 1:EX2_omega){
  for (j in 1:EX2_omega){
    
    if(i == j+1){
      
      age1 <- EX2_age[j]
      survp <- EX2_DIRECT_age_specific_survival[age1]
      EX2_survW[i,j] <- survp
    }
  }
}
#4.3.1: We can then populate the sub-diagonal of W using the vector EX1_survW
EX2_W <- EX2_W + EX2_survW

#4.4: Finally, we can populate all other elements of W that represent age class x's inclusive fitness contributions to the survival of other age classes.
EX2_survTransferW <- matrix(0, EX2_omega, EX2_omega)
for (i in 1:EX2_omega){
  for (j in 1:EX2_omega){
    
    if(j == 1 | j == i+1){
      EX2_survTransferW[j,i] <- 0
    }
    
    else{
      x <- EX2_age[i]
      y <- EX2_age[j]
      
      survTransfer <- EX2_indirectSurvContributions[j,i]
      relatedness <- EX2_rhatx[i]
      
      EX2_survTransferW[j,i] <- survTransfer*relatedness
      
    }
  }
}
##4.4.1: We can then populate all remaining elements of W with the vector EX2_survTransferW (in this case a matrix of zeros)
EX2_W <- EX2_W + EX2_survTransferW



### STAGE 5 - Forces of selection.
# We now have everything we need to compute forces of selection according to [8] and [10] in the main manuscript.
# Below, we compute forces of selection acting on a mutant allele that 1) alters survival rate at age x ([8]), and 2) alters rate of reproduction at age x ([10]).
# After deriving, we store the forces of selection in the vectors EX2_FOS_px and EX2_FOS_bx


# 5.1.1: First, we need to calculate the age-specific vector of inclusive reproductive value as the dominant left eigenvector of the fitness matrix W
EX2_asymptotic_frequency_W <- eigen.analysis(EX2_W)$stable.stage
EX2_inclusive_repro_value <- eigen.analysis(EX2_W)$repro.value 

EX2_mx <- rep(0, EX2_omega)
for (i in 1:length(EX2_mx)){
  
  EX2_mx[i] = EX2_age_specific_rateofreproduction[i]*((1-EX2_p_bar)/EX2_b_bar)  
}


Ham_scalar_T <- sum(EX2_age_lx*EX2_mx*EX2_age)


# 5.1.2: We can then calcualte the scalar vector of the asymptotic frequency and inclusive reproductive value
scalar <- sum(EX2_asymptotic_frequency_W*EX2_inclusive_repro_value)

# 5.2: Calculating the force of selection on survival rate at age x 

# 5.2.1 We first need to quantify hdotx and kdotx, that represent the proportions of offspring at the local patch after dispersal that are the expected direct and indirect genetic offspring of an individual aged x
EX2_hdotx <- rep(0, EX2_omega)
for (i in 1:length(EX2_hdotx)){
  
  T1 <- EX2_W[1,i]*(1-d)
  T2 <- EX2_age_specific_rateofreproduction[i]*(1-d) + (N-1)*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c)
  EX2_hdotx[i] <- T1/T2
}

EX2_kdotx <- rep(0, EX2_omega)
for (i in 1:length(EX2_kdotx)){
  Fbar <- sum(EX2_asymptotic_frequency*EX2_W[1,])
  
  T1 <- (N-1)*Fbar*(1-d)
  T2 <- EX2_age_specific_rateofreproduction[i]*(1-d) + (N-1)*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c)
  EX2_kdotx[i] <- T1/T2
  
}


# 5.2.2 We then calculate the age-specific force of selection on survival at age x and store in the vector EX1_FOS_px
EX2_FOS_px <- rep(0,EX2_omega)
EX2_FOS_px_Ham <- rep(0,EX2_omega)
for (i in 1:length(EX2_FOS_px)){
  
  A <- EX2_hdotx[i]
  r <- EX2_rhatx[i]
  B <- EX2_kdotx[i]
  
  term1 <- (EX2_asymptotic_frequency_W[i]*(EX2_inclusive_repro_value[1]*(A + r*B)))
  irvnext <- EX2_inclusive_repro_value[i+1] 
  if (is.na(irvnext)){irvnext <- 0}
  
  EX2_FOS_px[i] <- ((EX2_asymptotic_frequency_W[i]*irvnext) - term1)/scalar
  
  if (i == 20){next}
  future_lx <- EX2_age_lx[(i+1):EX2_omega]
  future_mx <- EX2_mx[(i+1):EX2_omega]
  
  EX2_FOS_px_Ham[i] <- (sum(future_lx*future_mx))/Ham_scalar_T
  
}


# 5.2.3: We then calculate the age-specific force of selection on reproduction at age x and store in the vector EX1_FOS_bx

EX2_FOS_bx <- rep(0,EX2_omega)
EX2_FOS_bx_Ham <- rep(0,EX2_omega)
for (i in 1:length(EX2_FOS_bx)){
  
  term1 <- (EX2_asymptotic_frequency_W[i]*EX2_inclusive_repro_value[1])/scalar
  
  gx_term1 <- (1 - EX2_age_specific_survival[i]) + (N-1)*(1 - EX2_p_bar)
  gx_term2 <- EX2_age_specific_rateofreproduction[i]*(1-d) + (N-1)*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c)
  gx <- gx_term1/gx_term2
  gbar <- N*(1 - EX2_p_bar)/(N*EX2_b_bar*(1-d) + N*EX2_b_bar*d*(1-c))
  hx <- EX2_offspringPostDispersal$hx[i]
  Ix <- (EX2_age_specific_indirect_mx[i]*(1-d))/gx_term2   ####
  r <- EX2_rhatx[i]
  kdotx <- EX2_kdotx[i]
  
  term2 <- ((1-d)*gx)*((1-hx)-Ix-(r*kdotx))
  
  term3 <- (1-c)*d*gbar
  
  EX2_FOS_bx[i] <-term1*(term2 + term3)
  
  EX2_FOS_bx_Ham[i] <- EX2_age_lx[i]/Ham_scalar_T
  
}


# Forces of selection figure

EX2_forceofselection <- data.frame("x" = EX2_age, "px" = EX2_FOS_px, "bx" = EX2_FOS_bx, "pxHam" = EX2_FOS_px_Ham, "bxHam" = EX2_FOS_bx_Ham)
EX2_Fig_FOS_px <- ggplot(EX2_forceofselection, aes(x, px))
EX2_Fig_FOS_px <- EX2_Fig_FOS_px + geom_line(col="darkblue", size = 1.5, linetype = "longdash") + labs(x="Age", y ="Force of selection on survival") + theme_m
EX2_Fig_FOS_px <- EX2_Fig_FOS_px + geom_hline(yintercept=0, linetype="dashed", color = "black")
EX2_Fig_FOS_px <- EX2_Fig_FOS_px + geom_line(aes(x, pxHam), col = "darkblue", size = 1.5)
EX2_Fig_FOS_px
EX2_Fig_FOS_bx <- ggplot(EX2_forceofselection, aes(x, bx))
EX2_Fig_FOS_bx <- EX2_Fig_FOS_bx + geom_line(col="red", size = 1.5, linetype = "longdash") + labs(x="Age", y ="Force of selection on reproduction") + theme_m
EX2_Fig_FOS_bx <- EX2_Fig_FOS_bx + geom_hline(yintercept=0, linetype="dashed", color = "black")
EX2_Fig_FOS_bx <- EX2_Fig_FOS_bx + geom_line(aes(x, bxHam), col = "red", size = 1.5)
EX2_Fig_FOS_bx
FIG3C <- grid.arrange(EX2_Fig_FOS_px, EX2_Fig_FOS_bx, ncol = 2)
ggsave(FIG3C, file = "Fig3c.pdf", width = 8, height = 5 )


### DEMOGRAHPY PLOTS
surv_transfers = rep(0, EX2_omega)
for (i in 1:EX2_omega){
  
  surv_transfers[i] <- sum(EX2_survTransferW[,i])
  
}

repro_transfers = rep(0, EX2_omega)
for (i in 1:EX2_omega){
  
  if (i < 6){
    repro_transfers[i] <- EX2_age_specific_effectivemx[i]
  }else{
    repro_transfers[i] <- 0
  }
    
  
}


EX2_age_specific_survival[20] <- EX2_age_specific_survival[19]
EX2_DIRECT_age_specific_survival[20] <- EX2_DIRECT_age_specific_survival[19]

EX2_background_demography <- data.frame("x"= EX2_age, "px" = EX2_age_specific_survival, "pdotx" = EX2_DIRECT_age_specific_survival, "bdotx" = EX2_DIRECT_age_specific_reproduction, "mx" = mx, "eff_mx" = EX2_age_specific_effectivemx, "surv_transfers" = surv_transfers, "repro_transfers" = repro_transfers)


EX2_surv_px <- ggplot(EX2_background_demography, aes(x, px))
EX2_surv_px <- EX2_surv_px + geom_line(col = "darkblue", size = 2) + labs(x = "Age", y = "Survival") + theme_m
EX2_surv_px <- EX2_surv_px + geom_line(aes(x, pdotx), size = 2, linetype = "dotted", col = "darkblue")

EX2_surv_transfers <- ggplot(EX2_background_demography, aes(x, surv_transfers))
EX2_surv_transfers <- EX2_surv_transfers + geom_line(col = "darkblue", size = 2) + labs(x = "Age", y = "Survival transfers") + theme_m
EX2_surv_transfers

EX2_repro <- ggplot(EX2_background_demography, aes(x, mx))
EX2_repro <- EX2_repro + geom_line(col = "red", size = 2) + labs(x = "Age", y = "Reproduction") + theme_m
EX2_repro <- EX2_repro + geom_line(aes(x, eff_mx), size = 2, linetype ="dotted", col ="red")


EX2_repro_transfers <- ggplot(EX2_background_demography, aes(x, repro_transfers))
EX2_repro_transfers <- EX2_repro_transfers+ geom_line(col = "red", size = 2) + labs(x = "Age", y = "Reproduction transfers") + theme_m
EX2_repro_transfers 


Fig3B <- grid.arrange(EX2_surv_px, EX2_surv_transfers, EX2_repro, EX2_repro_transfers, nrow = 2, ncol = 2)
ggsave(Fig3B, file = "Fig3B.pdf", height = 5, width = 8)