# THE EVOLUTIONARY CONSEQUENCES OF BEHAVIORAL PLASTICITY
#
# Evolutionary simulation model of the effect of behavioural 
# buffering on rates of evolution. R script re-runs all 
# simulations and recreates figures included in the paper
#
# Carlos A. Botero
# University of Texas at Austin
# Supported by an award from NSF ORC # 2413199
#
# 4 Feb 2026
##########################################################################
library(pbmcapply)
library(minpack.lm)
library(dplyr)

# clear memory
rm(list=ls())

# Simulation parameter values
Burnin.Gen <- 500 # # of generations for burn-in period
Baseline.Gen <- 1000 # number of generations for baseline simulation
Diasp.Gen <- 500 # Number of generations the diaspora is allowed to evolve

E <- 0.5 # environmental target for burn-in period
E2 <- 0.3 # environmental target for second batch of Baseline.Gen generations

# Beverton-Holt parameters
N <- 5000 # target carrying capacity of the simulated environment
R0 <- 2 # assumed baseline reproduction

# Tune Beverton-Holt constant to regulate populations at the desired carrying capacity
alpha <- (R0 - 1) / N 

# number of replicate simulations
numReps <- 100 

# This parameter captures the expected spread of variation introduced by mutation 
# across the entire set of loci coding for the quantitative genetic trait
mu.s <- 0.05 

# decay parameter for exponential decay fitness functions (use negative values)
d <- -1

# new target of selection for a diaspora taken from the population at the end of 
# the initial evolution
deltaE <- -1

# Number of individuals taken from last generation to start new population
diasp.size <- 100

################################################################################
# Evolutionary simulation algorithm

# Implement sexual reproduction assuming that we are dealing with a quantitative
# genetic trait controlled by many loci. As in many classic studies, I use here
# the mid-parent approach to approximate the offspring genotype as the mean of
# the parents plus a small random deviation to account for mutation. I assume as
# well that individuals with higher fitness are more likely to be chosen as
# mates
sexual.reproduction <- function(pop, mu.s) {
  num_off <- pop[,'NumOff']
  total_offspring <- sum(num_off)
  
  if(total_offspring == 0) return(NULL)
  
  n <- nrow(pop)
  
  # Create vector of parent1 indices according to NumOff
  parent1 <- rep(1:n, times = num_off)
  
  # Sample parent2 randomly, weighted by fitness
  prob <- pop[,'W'] / sum(pop[,'W'])
  parent2 <- sample(1:n, total_offspring, replace = TRUE, prob = prob)
  
  # Approximate offspring genotype as mid-parent value plus small offset from
  # aggregate mutations
  offspring_genotype <- (pop[parent1,'Genotype'] + pop[parent2,'Genotype']) / 2
  
  mutations <- rnorm(length(offspring_genotype), 0, mu.s)
  offspring_genotype <- offspring_genotype + mutations
  
  # Create offspring matrix
  offspring <- matrix(NA, nrow=length(offspring_genotype), ncol=4)
  colnames(offspring) <- c('Genotype','Mismatch','W','NumOff')
  offspring[,'Genotype'] <- offspring_genotype
  
  return(offspring)
}


replicate.run <- function(this.rep, deltaE, mu.s, d, diasp.size, b,
                          N, Baseline.Gen, Diasp.Gen, E, E2, R0, numReps) { 
  extinct <- FALSE
  
  # output list
  these.results <- list(
    "All.reps" = NA,
    "MeanGenotypeLastGenBurnin" = NA,
    "GenotypicVarianceBurnin" = NA,
    "ExtinctBurnin" = NA,
    "MeanGenotypeLastGenBaseline" = NA,
    "GenotypicVarianceBaseline" = NA,
    "ExtinctBaseline" = NA,
    "All.reps.diasp" = NA,
    "MeanGenotypeLastGenDiasp" = NA,
    "GenotypicVarianceDiasp" = NA,
    "ExtinctDiaspora" = NA,
    "PopSizes" = NULL
  )
  
  # Initialize population
  my.pop <- matrix(NA, N, 4)
  colnames(my.pop) <- c('Genotype', 'Mismatch', 'W', 'NumOff')
  my.pop[,'Genotype'] <- runif(N)
  
  # Storage for mean genotypes
  meanGenotypeVector <- numeric(Burnin.Gen + Baseline.Gen + Diasp.Gen)
  
  # burn-in period
  for (gen in 1:Burnin.Gen) {
    n <- nrow(my.pop)
    
    # compute mismatches
    my.pop[,'Mismatch'] <- abs((E - my.pop[,'Genotype'])*(1-b))
    
    # compute absolute fitness
    my.pop[,'W'] <- R0*exp(d * my.pop[,'Mismatch'])
    
    # Implement Beverton-Holt density dependence
    my.pop[,'W'] <- my.pop[,'W'] / (1 + alpha * n)
    
    # Translate fitness into expected number of offspring
    my.pop[,'NumOff'] <- rpois(n, my.pop[,'W'])
    
    # Create offspring based on sexual reproduction
    offspring.pop <- sexual.reproduction(my.pop, mu.s)
    
    if(is.null(offspring.pop)) {
      extinct <- TRUE
      break
    }
    
    my.pop <- offspring.pop
    meanGenotypeVector[gen] <- mean(my.pop[,'Genotype'])
    these.results$PopSizes <- c(these.results$PopSizes, nrow(my.pop))
  }
  
  # Now that populations have attained a natural level of standing genetic
  # variation, we begin the simulation of an adaptive process
  if(!extinct) {
    # store burn-in results for further analysis
    these.results$All.reps <- meanGenotypeVector[1:Burnin.Gen]
    these.results$MeanGenotypeLastGenBurnin <- mean(my.pop[,'Genotype'])
    these.results$GenotypicVarianceBurnin <- var(my.pop[,'Genotype'])
    these.results$ExtinctBurnin <- FALSE
    
    for (gen in (Burnin.Gen+1):(Burnin.Gen+Baseline.Gen)) {
      n <- nrow(my.pop)
  
      my.pop[,'Mismatch'] <- abs((E2 - my.pop[,'Genotype'])*(1-b))
      my.pop[,'W'] <- R0*exp(d * my.pop[,'Mismatch'])
      my.pop[,'W'] <- my.pop[,'W'] / (1 + alpha * n)
      my.pop[,'NumOff'] <- rpois(n, my.pop[,'W'])
      
      offspring.pop <- sexual.reproduction(my.pop, mu.s)
      
      if(is.null(offspring.pop)) {
        extinct <- TRUE
        these.results$ExtinctBaseline <- TRUE
        break
      }
      
      my.pop <- offspring.pop
      meanGenotypeVector[gen] <- mean(my.pop[,'Genotype'])
      these.results$PopSizes <- c(these.results$PopSizes, nrow(my.pop))
    }
    
    if(!extinct) {
      these.results$All.reps <- meanGenotypeVector[1:(Burnin.Gen+Baseline.Gen)]
      these.results$MeanGenotypeLastGenBaseline <- mean(my.pop[,'Genotype'])
      these.results$GenotypicVarianceBaseline <- var(my.pop[,'Genotype'])
      these.results$ExtinctBaseline <- FALSE
    }
    
    # Third stage evolution (take a small diaspora and expose it to a dramatic
    # change in its environment)
    newE <- E2 + deltaE
    
    if(!extinct) {
      # extract diaspora from last simulated generation
      my.pop <- my.pop[sample(1:nrow(my.pop), diasp.size),]
      #initialize
      meanGenotypeVectorDiaspora <- numeric(Diasp.Gen+1)
      # store starting mean population value
      meanGenotypeVectorDiaspora[1] <- mean(my.pop[,'Genotype'])
      
      for(gen in 1:Diasp.Gen) {
        my.pop[,'Mismatch'] <- abs((newE - my.pop[,'Genotype'])*(1-b))
        my.pop[,'W'] <- R0*exp(d * my.pop[,'Mismatch'])
        my.pop[,'W'] <- my.pop[,'W'] / (1 + alpha * nrow(my.pop))
        my.pop[,'NumOff'] <- rpois(nrow(my.pop), my.pop[,'W'])
        
        offspring.pop <- sexual.reproduction(my.pop, mu.s)
        
        if(is.null(offspring.pop)) {
          extinct <- TRUE
          break
        }
        
        my.pop <- offspring.pop
        meanGenotypeVectorDiaspora[gen+1] <- mean(my.pop[,'Genotype'])
      }
      
      these.results$All.reps.diasp <- meanGenotypeVectorDiaspora
      if(!extinct){
        these.results$All.reps <- meanGenotypeVector
        these.results$MeanGenotypeLastGenDiasp <- mean(my.pop[,'Genotype'])
        these.results$GenotypicVarianceDiasp <- var(my.pop[,'Genotype'])
        these.results$ExtinctDiaspora <- FALSE
      } else {
        these.results$ExtinctDiaspora <- TRUE
      }
    }
  } else {
    these.results$ExtinctBurnin <- TRUE
  }
  
  return(these.results)
}

simulation.set <- function(deltaE, mu.s, d, diasp.size) {
  for (b in seq(0, 1, by = 0.05)) {
    # run replicate simulations in parallel
    myres <- pbmclapply(1:numReps, function(thisrep) 
      replicate.run(thisrep, deltaE, mu.s, d, diasp.size, b,
                    N, Baseline.Gen, Diasp.Gen, E, E2, R0, numReps),
      mc.cores = detectCores()-4,
      mc.preschedule = FALSE)
    
    # extract results
    All.reps <- list()
    All.reps.diasp <- list()
    
    my.results <- data.frame('b' = b, 'Replicate' = 1:length(myres),
                             
                             'MeanGenotypeLastGenBurnin' = rep(NA,length(myres)),
                             'GenotypicVarianceBurnin' = rep(NA,length(myres)),
                             'ExtinctBurnin' = rep(NA,length(myres)),
                             
                             'MeanGenotypeLastGenBaseline' = rep(NA,length(myres)),
                             'GenotypicVarianceBaseline' = rep(NA,length(myres)),
                             'ExtinctBaseline' = rep(NA,length(myres)),
                             
                             'MeanGenotypeLastGenDiasp' = rep(NA,length(myres)),
                             'GenotypicVarianceDiasp' = rep(NA,length(myres)),
                             'ExtinctDiaspora' = rep(NA,length(myres)))
    
    for (i in 1:length(myres)) {
      All.reps[[i]] <- myres[[i]]$All.reps
      All.reps.diasp[[i]] <- myres[[i]]$All.reps.diasp
      
      my.results$MeanGenotypeLastGenBurnin[i] <- myres[[i]]$MeanGenotypeLastGenBurnin
      my.results$GenotypicVarianceBurnin[i] <- myres[[i]]$GenotypicVarianceBurnin
      my.results$ExtinctBurnin[i] <- myres[[i]]$ExtinctBurnin
      
      my.results$MeanGenotypeLastGenBaseline[i] <- myres[[i]]$MeanGenotypeLastGenBaseline
      my.results$GenotypicVarianceBaseline[i] <- myres[[i]]$GenotypicVarianceBaseline
      my.results$ExtinctBaseline[i] <- myres[[i]]$ExtinctBaseline
      
      my.results$MeanGenotypeLastGenDiasp[i] <- myres[[i]]$MeanGenotypeLastGenDiasp
      my.results$GenotypicVarianceDiasp[i] <- myres[[i]]$GenotypicVarianceDiasp
      my.results$ExtinctDiaspora[i] <- myres[[i]]$ExtinctDiaspora
      
    }
    
    # save results
    my.fname <- paste0("SimResults_mu.s.", mu.s,
                       "_diasp.size.", diasp.size,
                       "_d.", d, "_deltaE.", deltaE,
                       "_b.", b,".Rdata")
    save(All.reps, All.reps.diasp, my.results, file = my.fname)
  }
}
# ################################################################################
# # Run SIMULATIONS
mydir <- '/results'
setwd(mydir)

# run baseline simulations
simulation.set(deltaE = deltaE, mu.s = mu.s, d = d, diasp.size = diasp.size)

########
# sensitivity analysis
for (my.d in c(-0.5, -0.64, -0.77, -0.9, -1.18, -1.32, -1.45, -1.59, -1.73, -1.86, -2) ) {
  simulation.set(deltaE = deltaE, mu.s = mu.s, d = my.d, diasp.size = diasp.size)
}

for (my.mu.s in c(0.025, 0.0325, 0.04, 0.0475, 0.0625, 0.07, 0.0775, 0.085, 0.0925, 0.1) ) {
  simulation.set(deltaE = deltaE, mu.s = my.mu.s, d = d, diasp.size = diasp.size)
}

for (my.deltaE in c(-1.3, -1.6) ) {
  simulation.set(deltaE = my.deltaE, mu.s = mu.s, d = d, diasp.size = diasp.size)
}

for (my.diasp.size in c(50, 200) ) {
  simulation.set(deltaE = deltaE, mu.s = mu.s, d = d, diasp.size = my.diasp.size)
}

##########################################################################
# FIGURES
##########################################################################
# clear memory
rm(list=ls())

mydir <- '/results'
setwd(mydir)

# Simulation parameter values
Burnin.Gen <- 500 # # of generations for burn-in period
Baseline.Gen <- 1000 # number of generations for baseline simulation
Diasp.Gen <- 500 # Number of generations the diaspora is allowed to evolve

N <- 5000 # population size
MaxGen <- 1000 # number of generations per replicate

diasp.gens <- 500 # Number of generations the diaspora is allowed to evolve

E <- 0.3 # environmental target for initial MaxGen generations
E2 <- 0.1 # environmental target for second batch of MaxGen generations

numReps <- 100 # number of replicate simulations

mu.s <- 0.05 # baseline mutation rate

# Let's begin by exploring the general evolutionary trajectories and how they
# change with b
my.bs <- seq(0, 1, by = 0.05)
my.ds <- c(-0.5, -1.0, -2.0)

all.results <- NULL
all.trajectories <- NULL

for (d in my.ds) {
  for (i in 1:length(my.bs)) {
    my.fname <- paste0("SimResults_mu.s.", mu.s,"_diasp.size.100", 
                       "_d.", d,"_deltaE.-1_b.", my.bs[i], ".Rdata")
    load(file = my.fname)
    all.results <- rbind(all.results, cbind(data.frame('d' = d, 'b' = my.bs[i], 
                                                       my.results)))
    
    # extract information on evolutionary trajectories
    trajectories <- do.call(rbind, All.reps)
    means <- colMeans(trajectories)
    sds <- apply(trajectories, 2, sd)
    all.trajectories <- rbind(all.trajectories, 
                              data.frame("d" = d,
                                         "b" = as.character(my.bs[i]),
                                         "LowerBound" = means - 1.96*sds,
                                         "Mean" = means,
                                         "UpperBound" = means + 1.96*sds))
  }
}

baseline.trajectories <- all.trajectories

################################################################################
# FIG. 1. Fitness curves and evolutionary trajectories
pdf(file = "Fig.1.pdf", width = 6, height = 6, pointsize = 12)

layout(matrix(c(2:7,1,1), 4, 2, byrow = T), heights = c(1, 1, 1, 0.2),
       widths = c(1, 2.5))

bs.to.plot <- c(0, 0.3, 0.6, 0.9, 1)
my.colors <- c('#66c2a5', '#fc8d62', '#8da0cb', '#e78ac3', '#a6d854')

# Color legend 
par(mar = c(0, 0, 0, 0)) # Set margins
plot(1, type = "n", xlab = "", ylab = "", 
     xlim = c(0, length(bs.to.plot) + 2), ylim = c(0, 1),
     axes = FALSE)

# Add labels and colored squares horizontally
for (i in seq_along(bs.to.plot)) {
  text((i+0.75) - 0.2, 0.5, labels = bquote(b == .(bs.to.plot[i])), adj = 1) # Text to the left of the square
  rect((i+0.75) - 0.1, 0.3, i + 0.8, 0.7, col = my.colors[i], border = "black")    # Colored square
}

# now data
par(mar = c(5.1, 4.5, 2.1, 1.1))
delta <- seq (0,3, by = 0.01)
my.labels <- rbind(c('A', 'B'), c('C', 'D'), c('E', 'F'))

for (my.index in 1:length(my.ds)) {
  d <- my.ds[my.index]
  
  # first plot selection curves
  for (i in 1:length(bs.to.plot)) {
    W <- exp(d*delta*(1-bs.to.plot[i]))
    
    if (i == 1) {
      plot(delta, W, type = 'l', 
           col = my.colors[i], lty = 1,
           ylim = c(0,1), lwd = 2,
           xlab = expression(paste('Mismatch (|',italic('E-I'),'|)')), 
           ylab = expression(paste('Fitness (',italic('W'),')')))
    } else {
      lines(delta, W, col = my.colors[i], lwd = 2)
    }
  }
  
  mtext(my.labels[my.index,1], adj = -0.45, line = 0.5, cex = 1)
  
  # Now plot observed trajectories for baseline simulation (first MaxGen generations)
  this.trajectory <- all.trajectories[all.trajectories$d == d & 
                                        all.trajectories$b == bs.to.plot[1],]
  plot(this.trajectory$Mean[1:(Burnin.Gen+Baseline.Gen)], 
       xlab = 'Generation', ylab = expression('Mean genotype, '*bar(G)), 
       xlim = c(0,Burnin.Gen+Baseline.Gen), 
       ylim = c(floor(min(all.trajectories$UpperBound[1:(Burnin.Gen+Baseline.Gen)])*10)/10,
                ceiling(max(all.trajectories$UpperBound[1:(Burnin.Gen+Baseline.Gen)])*10)/10),
       type = 'l', col = my.colors[1])
  polygon(c(1:(Burnin.Gen+Baseline.Gen), (Burnin.Gen+Baseline.Gen):1), 
          c(this.trajectory$LowerBound[1:(Burnin.Gen+Baseline.Gen)], 
            rev(this.trajectory$UpperBound[1:(Burnin.Gen+Baseline.Gen)])),
          col = adjustcolor(my.colors[1], alpha.f = 0.3),
          border = FALSE)
  
  for (i in 2:length(bs.to.plot)) {
    this.trajectory <- all.trajectories[all.trajectories$d == d & 
                                          all.trajectories$b == bs.to.plot[i],]
    
    lines(this.trajectory$Mean[1:(Burnin.Gen+Baseline.Gen)], col = my.colors[i])
    polygon(c(1:(Burnin.Gen+Baseline.Gen), (Burnin.Gen+Baseline.Gen):1), 
            c(this.trajectory$LowerBound[1:(Burnin.Gen+Baseline.Gen)], 
              rev(this.trajectory$UpperBound[1:(Burnin.Gen+Baseline.Gen)])),
            col = adjustcolor(my.colors[i], alpha.f = 0.3),
            border = FALSE)
  }
  
  lines(c(500,500), c(-100,100), lty = 2)
  
  mtext(my.labels[my.index,2], adj = -0.12, line = 0.5, cex = 1)
  
}

dev.off()


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

# FIG. 2. Standing genetic variation as a function of cognitive buffering.

# retrieve baseline example
these.results <- NULL

for (b in my.bs) {
  load(file = paste0("SimResults_mu.s.", mu.s,"_diasp.size.100", 
                     "_d.-1_deltaE.-1_b.", b, ".Rdata"))
  
  these.results <- rbind(these.results, cbind(data.frame('mu.s' = 0.05, 'b' = b, 
                                                         my.results)))
}

pdf(file = "Fig.2.pdf", width = 3, height = 3, pointsize = 12)

layout(1)
par(mar = c(5.1, 5.1, 2.1, 2.1))
plot(these.results$b, these.results$GenotypicVarianceBaseline, 
     xlab = "Cognitive buffer, b", ylab = expression(""*s[italic(G1500)]^2*""),
     pch = 16, col = adjustcolor("#66c2a5", alpha.f = 0.3))

mymod <- lm(GenotypicVarianceBaseline ~ b, data = these.results)

x <- seq(0,1,by = 0.001)
y <- mymod$coefficients['(Intercept)'] + 
  mymod$coefficients['b']*x 
summary(mymod) 

lines(x,y,col = '#e78ac3', lw = 2)

dev.off()

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

# FIG. 3. Sensitivity analysis
pdf(file = "Fig.3.pdf", width = 5.25, height = 4.2, pointsize = 12)

layout(mat = matrix(c(2,3,4, 1,1,1, 6,7,8, 5,5,5), 4, 3, byrow = T), 
       heights = c(1, 0.2, 1, 0.2))

####
# first row: changes in strength of selection
my.d.colors <- c('#66c2a5', '#e78ac3', '#8da0cb')

# Color legend for d values
par(mar = c(0, 2, 0, 0)) # Set margins
plot(1, type = "n", xlab = "", ylab = "", 
     xlim = c(0, length(my.ds) + 2), ylim = c(0, 1),
     axes = FALSE)

# Add labels and colored squares horizontally
for (i in seq_along(my.ds)) {
  text((i+0.75) - 0.2, 0.35, labels = bquote(d == .(my.ds[i])), adj = 1) # Text to the left of the square
  rect((i+0.75) - 0.1, 0.1, i + 0.75, 0.5, col = my.d.colors[i], border = "black")    # Colored square
}

# now data
par(mar = c(4.1, 4.5, 2.1, 1.1))

# NB: The trajectory for b = 1 cannot be approximated by an exponential curve so 
# we remove it from this analysis

rates.of.evolution <- data.frame('d' = rep(my.ds, each = length(my.bs[my.bs<1])),
                                 'b'= my.bs[my.bs<1], 'DecayParameter' = NA)

for (d in my.ds) {
  these.trajectories <- all.trajectories[all.trajectories$d == d,]
  
  for (b in my.bs[my.bs<1]) {
    y <- these.trajectories$Mean[which(these.trajectories$b == b)]
    y <- y[(Burnin.Gen + 1):(Burnin.Gen + Baseline.Gen)]
    x <- 1:length(y)
    
    # Define initial guesses
    c_start <- 0.2  # Approximate asymptote
    a_start <- max(y) - c_start  # Initial value minus asymptote
    b_start <- 0.1  # Rough guess for decay rate
    
    # Fit the model
    fit <- nlsLM(y ~ c + a * exp(-d * x),
                 start = list(c = c_start, a = a_start, d = b_start), 
                 control = list('maxiter' = 100))
    rates.of.evolution$DecayParameter[which(rates.of.evolution$d== d &
                                              rates.of.evolution$b == b)] <- summary(fit)$coefficients['d', 'Estimate']
  }
}

rates.of.evolution <- na.omit(rates.of.evolution)

plot(rates.of.evolution$b[which(rates.of.evolution$d == my.ds[1])], 
     rates.of.evolution$DecayParameter[which(rates.of.evolution$d == my.ds[1])],
     pch = 16, ylim = range(rates.of.evolution$DecayParameter),
     col = my.d.colors[1], xlab = "Cognitive buffer, b", 
     ylab = 'Rate of evolution')

for (i in 2:length(my.ds)) {
  points(rates.of.evolution$b[which(rates.of.evolution$d == my.ds[i])], 
         rates.of.evolution$DecayParameter[which(rates.of.evolution$d == my.ds[i])],
         pch = 16, col = my.d.colors[i])
}

mtext('A', adj = -0.45, line = 0.5, cex = 1)

mymod <- lm(DecayParameter ~ b * d, data = rates.of.evolution)
summary(mymod)

#####
# 2) What about effects on standing genetic variation?
minval <- min(all.results$GenotypicVarianceBaseline)
maxval <- max(all.results$GenotypicVarianceBaseline)

# let's plot a few representative examples first...
for (i in 1:length(my.ds)) { 
  d <- my.ds[i]
  this.data <- all.results[all.results$d== d,]
  
  for ( j in 1:length(my.bs[my.bs<1])) {
    if (j == 1 & i ==1) {
      plot(rep(my.bs[j], length(all.results$GenotypicVarianceBaseline[all.results$d == d & 
                                                                        all.results$b ==my.bs[j]])), 
           all.results$GenotypicVarianceBaseline[all.results$d == d & 
                                                   all.results$b ==my.bs[j]], 
           pch = 16, xlab = "Cognitive buffer, b", 
           ylab = expression(""*s[italic(G1500)]^2*""), xlim = c(0,1),
           ylim = c(minval, maxval), col = adjustcolor(my.d.colors[i], alpha.f = 0.3))
    } else {
      points(rep(my.bs[j], length(all.results$GenotypicVarianceBaseline[all.results$d == d & 
                                                                          all.results$b == my.bs[j]])), 
             all.results$GenotypicVarianceBaseline[all.results$d == d & 
                                                     all.results$b == my.bs[j]], pch = 16,
             col = adjustcolor(my.d.colors[i], alpha.f = 0.3))
    }
  }
}

mtext('B', adj = -0.45, line = 0.5, cex = 1)

#####
# 3) How does the relationship between behavioral plasticity and standing
# genetic variation change with the strength of selection? 

# compute slope of this relationship for every level of d
my.ds <- c(-0.5, -0.64, -0.77, -0.9, -1, -1.18, -1.32, -1.45, -1.59, -1.73, -1.86, -2)
my.slopes <- numeric(length(my.ds))

for (i in 1:length(my.ds)) { 
  d <- my.ds[i]
  this.data <- all.results[all.results$d== d,]
  these.results <- NULL
  
  # load data for all b levels under that value of d
  for (b in my.bs) {
    load(file = paste0("SimResults_mu.s.", mu.s,"_diasp.size.100", 
                       "_d.", d,"_deltaE.-1_b.", b, ".Rdata"))
    
    these.results <- rbind(these.results, cbind(data.frame('mu.s' = 0.05, 'b' = b, 
                                                           my.results)))
  }
  
  # compute slope for change in genotypic variance as a function of plasticity
  mymod <- lm(GenotypicVarianceBaseline ~ b, data = these.results)
  my.slopes[i] <- summary(mymod)$coefficients['b','Estimate']
}

plot(my.ds, my.slopes, pch = 16, xlab = expression("Fitness decay, "*italic(d)), 
     ylab = expression("Slope for SGV | "*italic(b)*""), 
     xlim = range(my.ds))

mymod <- lm(my.slopes ~ my.ds)
summary(mymod)

abline(mymod)

mtext('C', adj = -0.45, line = 0.5, cex = 1)

####
# second row: changes in mutation rates...
my.mus <- c(0.5*mu.s, mu.s, 2*mu.s)
my.mu.colors <- c('#fbb4ae', '#b3cde3', '#ccebc5')

# Color legend for d values
par(mar = c(0, 2, 0, 0)) # Set margins
plot(1, type = "n", xlab = "", ylab = "", 
     xlim = c(0, length(my.mus) + 2), ylim = c(0, 1),
     axes = FALSE)

# Add labels and colored squares horizontally
for (i in seq_along(my.mus)) {
  text((i+0.75) - 0.2, 0.35, labels = bquote(mu == .(my.mus[i])), adj = 1) # Text to the left of the square
  rect((i+0.75) - 0.1, 0.1, i + 0.75, 0.5, col = my.mu.colors[i], border = "black")    # Colored square
}

# now data
par(mar = c(4.1, 4.5, 2.1, 1.1))
rates.of.evolution <- data.frame('mu.s' = rep(my.mus, each = numReps*length(my.bs)),
                                 'b' = rep(my.bs, numReps*length(my.mus)),
                                 'Replicate'= rep(1:numReps, each = length(my.bs)), 
                                 'DecayParameter' = NA)

all.results.mu <- NULL

for (this.mu.s in my.mus) {
  for (b in my.bs[my.bs<1]) {
    load(file = paste0("SimResults_mu.s.", this.mu.s, "_diasp.size.100", 
                       "_d.-1_deltaE.-1_b.", b, ".Rdata"))
    
    all.results.mu <- rbind(all.results.mu, cbind(data.frame('mu.s' = this.mu.s, 'b' = b, 
                                                             my.results)))
    
    trajectories <- do.call(rbind, lapply(All.reps, `[`, (Burnin.Gen + 1):(Burnin.Gen + Baseline.Gen)))
    
    for (i in 1:dim(trajectories)[1] ) {
      y <- trajectories[i,]
      x <- 1:length(y)
      
      # Define initial guesses
      c_start <- 0.2  # Approximate asymptote
      a_start <- max(y) - c_start  # Initial value minus asymptote
      b_start <- 0.1  # Rough guess for decay rate
      
      # Fit the model
      fit <- nlsLM(y ~ c + a * exp(-d * x),
                   start = list(c = c_start, a = a_start, d = b_start),
                   control = list(maxiter = 150))
      
      rates.of.evolution$DecayParameter[which(rates.of.evolution$mu.s == this.mu.s &
                                                rates.of.evolution$b == b &
                                                rates.of.evolution$Replicate == i)] <- summary(fit)$coefficients['d', 'Estimate']
    }
  }
}

# Let's compare first how rates change with increases in b
mean_decay_df <- rates.of.evolution %>%
  group_by(mu.s, b) %>%
  summarize(mean_DecayParameter = mean(DecayParameter, na.rm = TRUE), .groups = "drop")

# eliminate values for b == 1
mean_decay_df <- mean_decay_df[mean_decay_df$b != 1,]

for (i in 1:length(my.mus)) {
  this.data <- mean_decay_df[which(mean_decay_df$mu.s == my.mus[i]),]
  
  if (i == 1) {
    plot(jitter(this.data$b, factor = 0.2), 
         jitter(this.data$mean_DecayParameter, factor = 0.2), 
         xlim = range(mean_decay_df$b), 
         ylim = range(mean_decay_df$mean_DecayParameter),
         pch = 16, col = my.mu.colors[i], 
         xlab = "Cognitive buffer, b", 
         ylab = 'Rate of evolution')
  } else {
    points(jitter(this.data$b, factor = 0.3), 
           jitter(this.data$mean_DecayParameter, factor = 0.3), 
           pch = 16, col = my.mu.colors[i])
  }
}

mtext('D', adj = -0.45, line = 0.5, cex = 1)

#### 2) What about standing genetic variation?
minval <- min(all.results.mu$GenotypicVarianceBaseline)
maxval <- max(all.results.mu$GenotypicVarianceBaseline)

for (i in 1:length(my.mus)) { 
  this.mu <- my.mus[i]
  this.data <- all.results.mu[all.results.mu$mu.s == this.mu,]
  
  for ( j in 1:length(my.bs)) {
    if (j == 1 & i == 1) {
      plot(rep(my.bs[j], length(all.results.mu$GenotypicVarianceBaseline[all.results.mu$mu.s == this.mu & 
                                                                           all.results.mu$b ==my.bs[j]])), 
           all.results.mu$GenotypicVarianceBaseline[all.results.mu$mu.s == this.mu & 
                                                      all.results.mu$b ==my.bs[j]], 
           pch = 16, xlab = "Cognitive buffer, b", 
           ylab = expression(""*s[italic(G1500)]^2*""), xlim = c(0,1),
           ylim = c(minval, maxval), col = my.mu.colors[i])
    } else {
      points(rep(my.bs[j], length(all.results.mu$GenotypicVarianceBaseline[all.results.mu$mu.s == this.mu & 
                                                                             all.results.mu$b == my.bs[j]])), 
             all.results.mu$GenotypicVarianceBaseline[all.results.mu$mu.s == this.mu & 
                                                        all.results.mu$b == my.bs[j]], pch = 16,
             col = my.mu.colors[i])
    }
  }
}

mtext('E', adj = -0.45, line = 0.5, cex = 1)

#### 3) How does the relationship between behavioral plasticity and standing
# genetic variation change with mutation rate? 

# compute slope of this relationship for every level of mu.s
my.mus <- c(0.025, 0.0325, 0.04, 0.0475, 0.05, 0.0625, 0.07, 0.0775, 0.085, 0.0925, 0.1)
my.slopes <- numeric(length(my.mus))

for (i in 1:length(my.mus)) { 
  this.mu <- my.mus[i]
  this.data <- all.results[all.results$mu.s== this.mu,]
  these.results <- NULL
  
  # load data for all b levels under that value of d
  for (b in my.bs) {
    load(file = paste0("SimResults_mu.s.", this.mu,"_diasp.size.100", 
                       "_d.-1_deltaE.-1_b.", b, ".Rdata"))
    
    these.results <- rbind(these.results, cbind(data.frame('mu.s' = this.mu, 'b' = b, 
                                                           my.results)))
  }
  
  # compute slope for change in genotypic variance as a function of plasticity
  mymod <- lm(GenotypicVarianceBaseline ~ b, data = these.results)
  my.slopes[i] <- summary(mymod)$coefficients['b','Estimate']
}

plot(my.mus, my.slopes, pch = 16, xlab = expression("Mutation rate, "*italic(mu)*""), 
     ylab = expression("Slope for SGV | "*italic(b)*""), 
     xlim = range(my.mus))

mymod <- lm(my.slopes ~ my.mus + I(my.mus^2))
summary(mymod)

x <- seq(min(my.mus), max(my.mus), by = 0.005)
y <- summary(mymod)$coefficients['(Intercept)', 'Estimate'] +
  summary(mymod)$coefficients['my.mus', 'Estimate']*x +
  summary(mymod)$coefficients['I(my.mus^2)', 'Estimate']*(x^2)
lines(x,y) 

mtext('F', adj = -0.45, line = 0.5, cex = 1)


dev.off()

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

# FIG. 4. Cognitive buffering increases population survival when faced with
# environmental change.
pdf(file = "Fig.4.pdf", width = 6, height = 2.5, pointsize = 12)

layout(mat = matrix(c(1,3,5,2,4,6), 2, 3, 
                    byrow = T), heights = c(1, 0.15))

####

# Ediasp
my.deltaE <- c(-1, -1.3, -1.6)
my.deltaE.colors <- c("#a6cee3", "#fec44f", "#b2df8a")

# plot data
par(mar = c(4.1, 4.5, 2.1, 1.1))
proportion.extinct <- data.frame('b' = my.bs,
                                 'Baseline' = rep(NA, length(my.bs)),
                                 'Diaspora' = rep(NA, length(my.bs)))

for (i in 1:length(my.deltaE)) {
  Ediasp <- my.deltaE[i]
  for (b in my.bs) {
    load(file = paste0("SimResults_mu.s.", mu.s,"_diasp.size.100", 
                       "_d.-1_deltaE.", Ediasp, "_b.", b, ".Rdata"))
    
    # compute proportion of replicate simulations that went extinct
    proportion.extinct$Baseline[proportion.extinct$b == b] <- length(which(my.results$ExtinctBaseline == TRUE)) / dim(my.results)[1]
    proportion.extinct$II.Stage[proportion.extinct$b == b] <- length(which(my.results$Extinct2ndStage == TRUE)) / dim(my.results)[1]
    proportion.extinct$Diaspora[proportion.extinct$b == b] <- length(which(my.results$ExtinctDiaspora == TRUE)) / dim(my.results)[1]
  }
  
  # plot
  if (Ediasp == my.deltaE[1]) {
    plot(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
         xlab = "Cognitive buffer, b", 
         ylab = "P of succesful colonization",
         ylim = c(0,1), pch = 16, 
         col = adjustcolor(my.deltaE.colors[i], alpha.f = 0.45))
    lines(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
          col = adjustcolor(my.deltaE.colors[i], alpha.f = 0.45))
  } else {
    points(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
           pch = 16, col = adjustcolor(my.deltaE.colors[i], 
                                       alpha.f = 0.45))
    lines(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
          col = adjustcolor(my.deltaE.colors[i], alpha.f = 0.45))
  }
}

mtext('A', adj = -0.45, line = 0.5, cex = 1)

# now color legend
par(mar = c(0, 0, 0, 0)) # Set margins
plot(1, type = "n", xlab = "", ylab = "", 
     xlim = c(0, length(my.deltaE)), ylim = c(0, 1),
     axes = FALSE)

# Add labels and colored squares horizontally
for (i in seq_along(my.deltaE)) {
  text(i - 0.2, 0.35, labels = bquote(Delta*"E ="*.(my.deltaE[i])), adj = 1) # Text to the left of the square
  rect(i - 0.1, 0.3, i, 0.5, col = my.deltaE.colors[i], border = "black")    # Colored square
}

####
# Strength of selection
my.ds <- c(-0.5, -1, -2)

# first data
par(mar = c(4.1, 4.5, 2.1, 1.1))
proportion.extinct <- data.frame('b' = my.bs,
                                 'Baseline' = rep(NA, length(my.bs)),
                                 'Diaspora' = rep(NA, length(my.bs)))

for (i in 1:length(my.ds)) {
  d <- my.ds[i]
  for (b in my.bs) {
    load(file = paste0("SimResults_mu.s.", mu.s,"_diasp.size.100", 
                       "_d.", d, "_deltaE.-1_b.", b, ".Rdata"))
    
    # compute proportion of replicate simulations that went extinct
    proportion.extinct$Baseline[proportion.extinct$b == b] <- length(which(my.results$ExtinctBaseline == TRUE)) / dim(my.results)[1]
    proportion.extinct$Diaspora[proportion.extinct$b == b] <- length(which(my.results$ExtinctDiaspora == TRUE)) / dim(my.results)[1]
  }
  
  # plot
  if (d == my.ds[1]) {
    plot(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
         xlab = "Cognitive buffer, b", 
         ylab = "P of succesful colonization",
         ylim = c(0,1), pch = 16, 
         col = adjustcolor(my.d.colors[i], alpha.f = 0.45))
    lines(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
          col = adjustcolor(my.d.colors[i], alpha.f = 0.45))
  } else {
    points(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
           pch = 16, col = adjustcolor(my.d.colors[i], 
                                       alpha.f = 0.45))
    lines(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
          col = adjustcolor(my.d.colors[i], alpha.f = 0.45))
  }
}

mtext('B', adj = -0.45, line = 0.5, cex = 1)

# Color legend
par(mar = c(0, 0, 0, 0)) # Set margins
plot(1, type = "n", xlab = "", ylab = "", 
     xlim = c(0, length(my.ds)), ylim = c(0, 1),
     axes = FALSE)

# Add labels and colored squares horizontally
for (i in seq_along(my.ds)) {
  text(i - 0.2, 0.35, labels = bquote(d*" = "*.(my.ds[i])), adj = 1) # Text to the left of the square
  rect(i - 0.1, 0.3, i, 0.5, col = my.d.colors[i], border = "black")    # Colored square
}

#####
# mutation rate
my.mus <- c(0.025, 0.05, 0.1)

# first data
par(mar = c(4.1, 4.5, 2.1, 1.1))
proportion.extinct <- data.frame('b' = my.bs,
                                 'Baseline' = rep(NA, length(my.bs)),
                                 'Diaspora' = rep(NA, length(my.bs)))

for (i in 1:length(my.mus)) {
  this.mu.s <- my.mus[i]
  for (b in my.bs) {
    load(file = paste0("SimResults_mu.s.", this.mu.s, "_diasp.size.100", 
                       "_d.-1_deltaE.-1_b.", b, ".Rdata"))
    
    # compute proportion of replicate simulations that went extinct
    proportion.extinct$Baseline[proportion.extinct$b == b] <- length(which(my.results$ExtinctBaseline == TRUE)) / dim(my.results)[1]
    proportion.extinct$Diaspora[proportion.extinct$b == b] <- length(which(my.results$ExtinctDiaspora == TRUE)) / dim(my.results)[1]
  }
  
  # plot
  if (this.mu.s == my.mus[1]) {
    plot(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
         xlab = "Cognitive buffer, b", 
         ylab = "P of succesful colonization",
         ylim = c(0,1), pch = 16, 
         col = adjustcolor(my.mu.colors[i], alpha.f = 0.45))
    lines(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
          col = adjustcolor(my.mu.colors[i], alpha.f = 0.45))
  } else {
    points(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
           pch = 16, col = adjustcolor(my.mu.colors[i], 
                                       alpha.f = 0.45))
    lines(proportion.extinct$b, (1-proportion.extinct$Diaspora), 
          col = adjustcolor(my.mu.colors[i], alpha.f = 0.45))
  } 
}

mtext('C', adj = -0.45, line = 0.5, cex = 1)

# Color legend
par(mar = c(0, 0, 0, 0)) # Set margins
plot(1, type = "n", xlab = "", ylab = "", 
     xlim = c(0, length(my.mus)), ylim = c(0, 1),
     axes = FALSE)

# Add labels and colored squares horizontally
for (i in seq_along(my.mus)) {
  text(i - 0.2, 0.35, labels = bquote(mu*" = "*.(my.mus[i])), adj = 1) # Text to the left of the square
  rect(i - 0.1, 0.3, i, 0.5, col = my.mu.colors[i], border = "black")    # Colored square
}

dev.off()