####################################################################################################################################################
# R code to simulate the dynamic model communities used in the study
# "Temporal population variability in local forest communities has mixed effects on tree species richness across a latitudinal gradient"
# by Fung, T., Chisholm, R.A., ..., Zimmerman, J.
#####################################################################################################################################################

####################################################################################################################################
# The code below simulates the model commmunity conceptualized in Fig. S1 in the Supporting Information.
# This model was used to assess bias in the two metrics of temporal population variability considered (see Appendix S2).
####################################################################################################################################

# The function "model_sim_1" simulates species dynamics for model communities with no immigration and different values of the per-capita 
# speciation probability, as specified by the vector "nus". Each community has "J" individuals, and dynamics are simulated starting from 
# 1 species, for "t_max" time-steps. The first species has a fitness value that is drawn randomly from a lognormal distribution with mean 1 
# and variance "A". In each time-step, with probability 1/"tau", the fitness of every species in the model is redrawn randomly from the 
# lognormal distribution. The simulations are initiated with the random seed number "seed". 
#
# The output is a list with three elements:
# (i) The first element "mean_S_Vec" is a vector of the mean number of species over the last "t_thres" time-steps, for each community. 
# (ii) The second element "mean_deltaN_Vec" is a vector of the first metric of temporal population variability (referred to as "Metric 1" in Appendix S2), 
#      for a subsample of "J_subsample" individuals in each community. This is calculated using species dynamics in the last "t_thres" time-steps, with the 
#     number of time-steps in a census interval defined by "t_calc_deltaN".
# (iii) The third element "mean_deltaN2_Vec" is a vector of the second metric of temporal population variability (referred to as "Metric 2" in Appendix S2), 
#       for a subsample of "J_subsample" individuals in each community. This is calculated using species dynamics in the last "t_thres" time-steps.
#
# This function is used to generate results to draw Figs. S2-S5 in Appendix S2.

model_sim_1 <- function(seed,J,nus,A,tau,t_max,t_thres,t_calc_deltaN,J_subsample){
  
  set.seed(seed)
  # "mu_lognormal" and "sigma_lognormal" are the mean and s.d. of the lognormal distribution of fitnesses.
  # "mu" and "sigma" are the mean and s.d. of the corresponding normal distribution.
  mu_lognormal = 1
  sigma_lognormal = sqrt(A)
  mu = log((mu_lognormal^2)/sqrt((mu_lognormal^2)+(sigma_lognormal^2)))
  sigma = sqrt(log(1+((sigma_lognormal^2)/(mu_lognormal^2))))
  n_nus = length(nus)
  # Create empty vectors for holding output values from the model simulations.
  # "mean_S_Vec" is vector with ith entry corresponding to the mean number of species over the last "t_thres" time-steps, for the ith community.
  # "mean_deltaN_Vec" is vector with ith entry corresponding to the first metric of temporal population variability, calculated using species dynamics over the last "t_thres" time-steps, for a subsample of "J_subsample" individuals in the ith community.
  # "mean_deltaN_per_N1_class_count_Mat" is a matrix with the ith row and jth column showing the number of census intervals with at least one species with initial abundance (abundance at the start of a census interval) j, 
  #    considering the last "t_thres" time-steps for a subsample of "J_subsample" individuals in the ith community. These are used to calculate entries in "mean_deltaN_per_N1_class_Mat".
  # "mean_deltaN_per_N1_class_Mat" is a matrix with the ith row and jth column showing the mean temporal population variability (calculated using eq. 1 in the main text)
  #    for species with initial abundance (abundance at the start of a census interval) j, for a subsample of "J_subsample" individuals in the ith community. These are used to calculate entries in "mean_deltaN2_Vec".
  # "mean_deltaN2_Vec" is vector with ith entry corresponding to the second metric of temporal population variability, calculated using species dynamics over the last "t_thres" time-steps, for a subsample of "J_subsample" individuals in the ith community.
  mean_S_Vec = numeric(n_nus)
  mean_deltaN_Vec = numeric(n_nus)
  mean_deltaN_per_N1_class_count_Mat = matrix(0,n_nus,J)
  mean_deltaN_per_N1_class_Mat = matrix(0,n_nus,J)
  mean_deltaN2_Vec = numeric(n_nus)
  
  # Loop over the values of per-speciation probability in vector "n_nus".
  for(i in 1:n_nus){
    # "counter1" keeps track of the number of census intervals during the last "t_thres" time-steps.
    counter1 = 0
    nu = nus[i]
    cat('nu = ',nu,'\n',sep='')
    
    # "N_Vec" is the vector of species abundances (in a model community).
    # Start with 1 species with "J" individuals.
    N_Vec = rep(J,1)
    # NVec_calc_deltaN_start is the vector of species abundances at the start of a census interval. It is used in calculating entries in "mean_deltaN_Vec" and "mean_deltaN2_Vec".
    NVec_calc_deltaN_start = N_Vec
    # "f_Vec" is the vector of fitnesses for each species.
    f_Vec = rlnorm(1,meanlog=mu,sdlog=sigma)
    # "S" is the total number of species corresponding to "N_Vec".
    S = length(N_Vec)
    # Loop over "t_max" time-steps for ith community.
    for(t in 1:t_max){
      if(t%%1e5==0){ 
        cat('t = ',t,'\n',sep='')
      }
      # Redraw species fitnesses with probability 1/"tau".
      ran1 = runif(1)
      if(ran1<(1/tau)){
        f_Vec = rlnorm(length(N_Vec),meanlog=mu,sdlog=sigma)
      }
      if(t==t_thres-t_calc_deltaN){
        NVec_calc_deltaN_start = N_Vec
      }
      # If "t" is equal or greater to "t_thres", then update entries in "mean_S_Vec", "mean_deltaN_Vec", "mean_deltaN_per_N1_class_count_Mat" and "mean_deltaN_per_N1_class_Mat" accordingly.
      # After updating entries, entries in "N_Vec" and "f_Vec" corresponding to extinct species are removed.
      if(t>=t_thres){
        mean_S_Vec[i] = mean_S_Vec[i]+S
        if(t%%t_calc_deltaN==0){
          counter1 = counter1+1
          if(J==J_subsample){
            NVec_minus_NVec_calc_deltaN_start = abs(N_Vec-NVec_calc_deltaN_start)
            zero_ind = which(NVec_calc_deltaN_start==0)
            if(length(zero_ind)>0){
              NVec_calc_deltaN_start = NVec_calc_deltaN_start[-zero_ind]
              NVec_minus_NVec_calc_deltaN_start = NVec_minus_NVec_calc_deltaN_start[-zero_ind]
            }
            mean_deltaN_Vec[i] = mean_deltaN_Vec[i]+mean(NVec_minus_NVec_calc_deltaN_start)
            dat1 = data.frame(NVec_calc_deltaN_start,NVec_minus_NVec_calc_deltaN_start)
            mean_deltaN_per_N1_class = tapply(dat1$NVec_minus_NVec_calc_deltaN_start,dat1$NVec_calc_deltaN_start,mean)
            names_mean_deltaN_per_N1_class = as.numeric(names(mean_deltaN_per_N1_class))
            vals_mean_deltaN_per_N1_class = as.numeric(mean_deltaN_per_N1_class)
            mean_deltaN_per_N1_class_Mat[i,names_mean_deltaN_per_N1_class] = mean_deltaN_per_N1_class_Mat[i,names_mean_deltaN_per_N1_class]+vals_mean_deltaN_per_N1_class
            table_NVec_calc_deltaN_start = table(NVec_calc_deltaN_start)
            names_table_NVec_calc_deltaN_start = as.numeric(names(table_NVec_calc_deltaN_start))
            mean_deltaN_per_N1_class_count_Mat[i,names_table_NVec_calc_deltaN_start] = mean_deltaN_per_N1_class_count_Mat[i,names_table_NVec_calc_deltaN_start]+1
          }else{
            NVec_calc_deltaN_start_sppID_Vec = rep(1:length(NVec_calc_deltaN_start),NVec_calc_deltaN_start)
            table_NVec_calc_deltaN_start_sub = table(sample(NVec_calc_deltaN_start_sppID_Vec,J_subsample,replace=FALSE))
            names_NVec_calc_deltaN_start_sub = as.numeric(names(table_NVec_calc_deltaN_start_sub))
            vals_NVec_calc_deltaN_start_sub = as.numeric(table_NVec_calc_deltaN_start_sub)
            N_Vec_sppID_Vec = rep(1:length(N_Vec),N_Vec)
            table_NVec_sub = table(sample(N_Vec_sppID_Vec,J_subsample,replace=FALSE))
            names_NVec_sub = as.numeric(names(table_NVec_sub))
            vals_NVec_sub = as.numeric(table_NVec_sub)
            S_sub = length(vals_NVec_calc_deltaN_start_sub)
            mean_S_Vec[i] = mean_S_Vec[i]+S_sub
            match_ind = match(names_NVec_calc_deltaN_start_sub,names_NVec_sub)
            match_ind_NA_ind = which(is.na(match_ind))
            match_ind_noNA_ind = which(!is.na(match_ind))
            NVec_minus_NVec_calc_deltaN_start_sub = numeric(length(vals_NVec_calc_deltaN_start_sub))
            NVec_minus_NVec_calc_deltaN_start_sub[match_ind_noNA_ind] = abs(vals_NVec_calc_deltaN_start_sub[match_ind_noNA_ind]-vals_NVec_sub[match_ind[match_ind_noNA_ind]])
            mean_deltaN_Vec[i] = mean_deltaN_Vec[i]+mean(abs(vals_NVec_calc_deltaN_start_sub[match_ind_noNA_ind]-vals_NVec_sub[match_ind[match_ind_noNA_ind]]))
            if(length(match_ind_NA_ind)>0){
              NVec_minus_NVec_calc_deltaN_start_sub[match_ind_NA_ind] = abs(vals_NVec_calc_deltaN_start_sub[match_ind_NA_ind])
            }
            dat1 = data.frame(vals_NVec_calc_deltaN_start_sub,NVec_minus_NVec_calc_deltaN_start_sub)
            mean_deltaN_per_N1_class = tapply(dat1$NVec_minus_NVec_calc_deltaN_start_sub,dat1$vals_NVec_calc_deltaN_start_sub,mean)
            names_mean_deltaN_per_N1_class = as.numeric(names(mean_deltaN_per_N1_class))
            vals_mean_deltaN_per_N1_class = as.numeric(mean_deltaN_per_N1_class)
            mean_deltaN_per_N1_class_Mat[i,names_mean_deltaN_per_N1_class] = mean_deltaN_per_N1_class_Mat[i,names_mean_deltaN_per_N1_class]+vals_mean_deltaN_per_N1_class
            mean_deltaN_per_N1_class_count_Mat[i,vals_NVec_calc_deltaN_start_sub] = mean_deltaN_per_N1_class_count_Mat[i,vals_NVec_calc_deltaN_start_sub]+1
          }
          # Remove entries in "N_Vec" and "f_Vec" corresponding to extinct species.
          extinct_ind = which(N_Vec==0)
          if(length(extinct_ind)>0){
            N_Vec = N_Vec[-extinct_ind]
            f_Vec = f_Vec[-extinct_ind]
          }
          NVec_calc_deltaN_start = N_Vec
        }
      }
      # Randomly choose an individual to die.
      death_ind = sample(1:length(N_Vec),1,prob=N_Vec)
      N_Vec[death_ind] = N_Vec[death_ind]-1
      # If a species goes extinct, update "S", "N_Vec" and "f_Vec" accordingly.
      if(N_Vec[death_ind]==0){
        S = S-1
        if(t<t_thres-t_calc_deltaN){
          N_Vec = N_Vec[-death_ind]
          f_Vec = f_Vec[-death_ind]
        }
      }
      # Randomly choose a replacement individual, via speciation or a local birth.
      # If a new species appears via speciation, update "S", "N_Vec", "f_Vec" and "NVec_calc_deltaN_start" accordingly.
      # If a local birth occurs, update "N_Vec" accordingly.
      ran1 = runif(1)
      if(ran1<nu){
        S = S+1
        N_Vec = c(N_Vec,1)
        f_Vec = c(f_Vec,rlnorm(1,meanlog=mu,sdlog=sigma))
        NVec_calc_deltaN_start = c(NVec_calc_deltaN_start,0)
      }else{
        prob_birth = (N_Vec*f_Vec)/sum(N_Vec*f_Vec)
        birth_ind = sample(1:length(N_Vec),1,prob=prob_birth)
        N_Vec[birth_ind] = N_Vec[birth_ind]+1
      }  
    }
    # Update entries in "mean_S_Vec", "mean_deltaN_Vec" and "mean_deltaN_per_N1_class_Mat".
    if(J==J_subsample){
      mean_S_Vec[i] = mean_S_Vec[i]/(t_max-t_thres+1)
    }else{
      mean_S_Vec[i] = mean_S_Vec[i]/counter1
    }
    mean_deltaN_Vec[i] = mean_deltaN_Vec[i]/counter1
    mean_deltaN_per_N1_class_Mat[i,] = mean_deltaN_per_N1_class_Mat[i,]/mean_deltaN_per_N1_class_count_Mat[i,]
  }
  # Update entries in "mean_deltaN2_Vec".
  # "init_N_common" is the set of initial abundances with at least one recorded species in each model community, in the last "t_thres" time-steps.
  init_N1_common = which(!is.na(colMeans(mean_deltaN_per_N1_class_Mat)))
  mean_deltaN2_Vec = rowMeans(mean_deltaN_per_N1_class_Mat[,init_N1_common])
  
  return(list(mean_S_Vec,mean_deltaN_Vec,mean_deltaN2_Vec))
  
}


####################################################################################################################################
# The code below simulates the model commmunity conceptualized in Fig. 3 in the Main Text.
# This model was used to determine the regimes of temporal environmental variance corresponding to the observed tree communities
# (see Appendix S5).
####################################################################################################################################

# The function "model_sim_2" simulates species dynamics for a model community without speciation or immigration.
# The simulation starts with an initial vector of species abundances "N_Vec_1stcensus", which corresponds to the first census of a 
# rarefied dataset for an observed tree community. Species dynamics are then simulated for a length of time corresponding to all censuses 
# of the observed tree community, i.e. the same number of deaths corresponding to these censuses. During this time, the number of births 
# is also the same as in the observed tree community. The vectors "B_Vec_censuses" and D_Vec_censuses" show the number of births and deaths 
# in between each pair of consecutive censuses. After the census period, species dynamics are simulated for a further "t_max_a" time-steps, 
# during which the total number of individuals is fixed. The initial set of species have fitnesses that are drawn randomly from a lognormal 
# distribution with mean 1 and variance "A". In each time-step, with probability 1/"tau", the fitness of every species in the model is redrawn 
# randomly from the lognormal distribution. The simulations are initiated with the random seed number "seed".

# The output is a list with four elements:
# (i) The first element "mean_DeltaN" is the first metric of temporal population variability (referred to as "Metric 1" in Appendix S2), 
#     for the model community. This is calculated using species abundances over the census period, and the vector of the time between 
#     censuses, given by "DeltaT_Vec_censuses".
# (ii) The second element "mean_cdf" is the cumulative distribution function (cdf) of pairwise temporal correlations between species abundances 
#      for the model community, averaged over all pairs of consecutive censuses over the census period. The cdf is given as values of the 
#      cumulative probability over 1,000 values of the pairwise temporal correlation equally spaced between -1 and 1.
# (iii) The third element "S_init" is the number of species at the start of the simulation.
# (iv) The third element "S_fin" is the the number of species at the end of the simulation.
#
# This function is used to generate results to draw Fig. 5 in the Main Text and Figs. 13-18 in Appendix S2.

model_sim_2 <- function(seed,A,tau,t_max_a,N_Vec_1stcensus,B_Vec_censuses,D_Vec_censuses,DeltaT_Vec_censuses){
  
  set.seed(seed)
  # "n_censuses" is the number of censuses of the corresponding observed tree community.
  n_censuses = length(B_Vec_censuses)
  lambda = 1/tau
  # "mu_lognormal" and "sigma_lognormal" are the mean and s.d. of the lognormal distribution of fitnesses.
  # "mu" and "sigma" are the mean and s.d. of the corresponding normal distribution.
  mu_lognormal = 1
  sigma_lognormal = sqrt(A)
  mu = log((mu_lognormal^2)/sqrt((mu_lognormal^2)+(sigma_lognormal^2)))
  sigma = sqrt(log(1+((sigma_lognormal^2)/(mu_lognormal^2))))
  # "t_max" is maximum number of time-steps to run in the simulation.
  t_max = sum(D_Vec_censuses)+t_max_a
  # "census_sim_ind" refers to current census in simulation.
  census_sim_ind = 1
  # "Bs_over_Ds" is ratio of births to deaths in between each pair of censuses.
  # "B_census_sim_ind" is observed B (number of births) for the current census.
  # "D_census_sim_ind" is observed D (number of deaths) for the current census.
  # "B_over_D_census_sim_ind" is observed B/D for the current census.
  # "B_sim" keeps track of B since beginning of a census (in the simulation).
  # "D_sim" keeps track of D since beginning of a census (in the simulation).
  Bs_over_Ds = B_Vec_censuses/D_Vec_censuses
  B_census_sim_ind = B_Vec_censuses[census_sim_ind]
  D_census_sim_ind = D_Vec_censuses[census_sim_ind]
  B_over_D_census_sim_ind = Bs_over_Ds[census_sim_ind]
  B_sim = 0
  D_sim = 0
  # "N_Vec" is the vector of species abundances (in a model community).
  # "N_Vec_init_census" is "N_Vec" at the beginning of a census.
  # "N_Vec_all_censuses" is a matrix with the ith column being "N_Vec" at the ith census.
  N_Vec = N_Vec_1stcensus
  N_Vec_init_census = N_Vec
  N_Vec_all_censuses = matrix(0,length(N_Vec),n_censuses)
  N_Vec_all_censuses[,1] = N_Vec
  # "S" is the total number of species.
  S = length(N_Vec)
  S_init = S
  # "f_Vec" is the vector of fitnesses for each species.
  f_vec = rlnorm(length(N_Vec),meanlog=mu,sdlog=sigma)
  mean_DeltaN = 0
  for(t in 1:t_max){
    if(t%%1e5==0){ 
      cat('t = ',t,'\n',sep='')
    }
    # Redraw species fitnesses with probability 1/"tau".
    ran1 = runif(1)
    if(ran1<lambda){
      f_vec = rlnorm(length(N_Vec),meanlog=mu,sdlog=sigma)
    }
    # Update "mean_DeltaN" if still in census period and at the time of a census
    if(census_sim_ind<=(n_censuses-1)){
      if(D_sim==D_census_sim_ind){
        N_Vec_init_census_pos_ind = which(N_Vec_init_census>0)
        N_Vec_init_census_pos = N_Vec_init_census[N_Vec_init_census_pos_ind]
        N_Vec_init_pos = N_Vec[N_Vec_init_census_pos_ind]
        NVecinitcensuspos_NVecinitpos_divide_DeltaT = abs(N_Vec_init_census_pos-N_Vec_init_pos)/DeltaT_Vec_censuses[census_sim_ind]
        mean_DeltaN = mean_DeltaN+mean(NVecinitcensuspos_NVecinitpos_divide_DeltaT)
        N_Vec_all_censuses[,census_sim_ind+1] = N_Vec
        census_sim_ind = census_sim_ind+1
        B_census_sim_ind = Bs_tot_rarefied[census_sim_ind]
        D_census_sim_ind = Ds_tot_rarefied[census_sim_ind]
        B_over_D_census_sim_ind = Bs_over_Ds[census_sim_ind]
        D_sim = 0
        B_sim = 0
        N_Vec_init_census = N_Vec
      }
    }
    if(census_sim_ind<=(n_censuses-1)){
      # Randomly choose an individual to die.
      # Update "N_Vec" accordingly.
      death_ind = sample(1:length(N_Vec),1,prob=N_Vec)
      N_Vec[death_ind] = N_Vec[death_ind]-1
      D_sim = D_sim+1
      # If a species goes extinct, update "S".
      if(N_Vec[death_ind]==0){
        S = S-1
      }
      # Number of births is determined by "B_over_D_census_sim_ind".
      # Randomly choose identity of newborn individuals. 
      # Update "N_Vec" accordingly. 
      if(D_sim==D_census_sim_ind){
        n_births = B_census_sim_ind-B_sim
      }else{
        if((B_sim/D_sim)<B_over_D_census_sim_ind){
          n_births = floor(B_over_D_census_sim_ind*D_sim-B_sim)
        }else{
          n_births = 0
        }
      }
      if(n_births>0){
        prob_birth = (N_Vec*f_vec)/sum(N_Vec*f_vec)
        for(o in 1:n_births){
          birth_ind = sample(1:length(N_Vec),1,prob=prob_birth)
          N_Vec[birth_ind] = N_Vec[birth_ind]+1
        }
      }
      B_sim = B_sim+n_births
    }else{
      # After the census period, total number of individuals remains fixed at whatever value it had at the end of the census period.
      # Randomly choose an individual to die.
      # Update "N_Vec" accordingly.
      death_ind = sample(1:length(N_Vec),1,prob=N_Vec)
      N_Vec[death_ind] = N_Vec[death_ind]-1
      # If a species goes extinct, update "S".
      if(N_Vec[death_ind]==0){
        S = S-1
      }
      # Randomly choose a replacement individual, via a local birth.
      # Update "N_Vec" accordingly.
      prob_birth = (N_Vec*f_vec)/sum(N_Vec*f_vec)
      birth_ind = sample(1:length(N_Vec),1,prob=prob_birth)
      N_Vec[birth_ind] = N_Vec[birth_ind]+1
    }
  }
  mean_DeltaN = mean_DeltaN/(n_censuses-1)
  S_fin = S
  # Using "N_Vec_all_censuses", "mean_cdf" is computed.
  S1 = nrow(N_Vec_all_censuses)
  temp_cor_mat = matrix(0,S1,S1)
  mean_N_spp = rowMeans(N_Vec_all_censuses)
  sd_N_spp = sqrt(rowMeans(N_Vec_all_censuses^2)-(mean_N_spp^2))
  for(i in 1:S1){
    for(j in 1:S1){
      Ns_sp_i = as.numeric(N_Vec_all_censuses[i,])
      Ns_sp_j = as.numeric(N_Vec_all_censuses[j,])
      mean_N_sp_i = mean_N_spp[i]
      mean_N_sp_j = mean_N_spp[j]
      sd_N_sp_i = sd_N_spp[i]
      sd_N_sp_j = sd_N_spp[j]
      if(max(sd_N_sp_i,sd_N_sp_j)>0){
        temp_cor_ij = (1/n_censuses)*(1/(sd_N_sp_i*sd_N_sp_j))*sum((Ns_sp_i-mean_N_sp_i)*(Ns_sp_j-mean_N_sp_j))
      }else{
        temp_cor_ij = NA
      }
      temp_cor_mat[i,j] = temp_cor_ij
    }
  }
  temp_cor_mat_numeric = as.numeric(temp_cor_mat)
  temp_cor_mat_numeric_noNA = temp_cor_mat_numeric[which(!is.na(temp_cor_mat_numeric))]
  temp_cor_cdf = ecdf(temp_cor_mat_numeric_noNA)
  xs = seq(-1,1,length.out=1e3)
  mean_cdf = temp_cor_cdf(xs)
  
  return(list(mean_DeltaN,mean_cdf,S_init,S_fin))
  
}