###
### R Script for diversification model selection in 
### Baird et al 2021: "50 Million Years of Beetle Evolution Along the Antarctic Polar Front"
### using RPANDA (Morlon et al; https://cran.r-project.org/web/packages/RPANDA/index.html)
###
### Note: only exponentially-variable (not linear) models were tested, on recommendation
### by RPANDA developers and due to issues with negative speciation rates when using
### linear models
################################################################################


library(RPANDA)
Loading required package: picante
Loading required package: ape
Loading required package: vegan
Loading required package: permute
Loading required package: lattice
This is vegan 2.5-7
Loading required package: nlme


tree<-read.tree(file="BEAST_dated_final_tree.newick")
tot_time<-max(node.age(tree)$ages)



############### Constant speciation, no extinction #################

f.lamb <-function(t,y){y[1]}
f.mu<-function(t,y){0}
lamb_par<-c(0.1)
mu_par<-c()
result_spcst_noext <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,fix.mu=TRUE,dt=1e-3)
result_spcst_noext



############ Constant speciation, no extinction, different initial parameters ##############

lamb_par<-c(0.01)
result_spcst_noext2 <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,fix.mu=TRUE,dt=1e-3)
result_spcst_noext2



############## Constant speciation, constant extinction ################

f.lamb <-function(t,y){y[1]}
f.mu<-function(t,y){y[1]}
lamb_par<-c(0.14)
mu_par<-c(0.001)
result_spcst_excst <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,cst.mu=TRUE,dt=1e-3)
result_spcst_excst



############## Constant speciation, constant extinction, different initial parameters ################

lamb_par<-c(0.5)
mu_par<-c(0.01)
result_spcst_excst2 <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,cst.mu=TRUE,dt=1e-3)
result_spcst_excst2



###########################################################
################  TIME-DEPENDENT MODELS  ####################
###########################################################


############# Speciation variable (expo), no extinction ###############

f.lamb <-function(t,y){y[1] * exp(y[2] * t)}
mu_par<-c()
f.mu<-function(t,y){0}
lamb_par<-c(0.05, 0.01)
result_spexp_noext <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,expo.lamb=TRUE,fix.mu=TRUE,dt=1e-3)
result_spexp_noext



############# Speciation variable (expo), no extinction, different initial parameters ############

lamb_par<-c(0.1, 0.001)
result_spexp_noext2 <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,expo.lamb=TRUE,fix.mu=TRUE,dt=1e-3)
result_spexp_noext2 



############## Speciation variable (expo), constant extinction ################

f.lamb<-function(t,y){y[1]*exp(y[2]*t)}
f.mu<-function(t,y){y[1]}
lamb_par<-c(0.3,0.001)
mu_par<-c(0.001)
result_spexp_extcst <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,expo.lamb=TRUE,cst.mu=TRUE,dt=1e-3)
result_spexp_extcst 



############# Speciation variable (expo), constant extinction, different initial parameters ############

lamb_par<-c(0.1,0.01)
mu_par<-c(0.005)
result_spexp_extcst2 <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,expo.lamb=TRUE,cst.mu=TRUE,dt=1e-3)
result_spexp_extcst2



############## Speciation variable (expo), extinction variable (expo) ###############

f.lamb<-function(t,y){y[1] * exp(y[2] * t)}
f.mu<-function(t,y){y[1] * exp(y[2] * t)}
lamb_par<-c(0.3, 0.001)
mu_par<-c(0.01, 0.001)
result_spexp_extexp <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,expo.lamb=TRUE,expo.mu=TRUE,dt=1e-3)
result_spexp_extexp 



############# Speciation variable (expo), extinction variable (expo), different initial parameters #########

lamb_par<-c(0.1, 0.01)
 mu_par<-c(0.01, 0.01)
result_spexp_extexp2 <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,expo.lamb=TRUE,expo.mu=TRUE,dt=1e-3)
result_spexp_extexp2



################ Constant speciation, extinction variable (expo) ###################

f.lamb <-function(t,y){y[1]}
f.mu<-function(t,y){y[1] * exp(y[2] * t)}
lamb_par<-c(0.1)
mu_par<-c(0.001, 0.01)
result_spcst_extexp <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,dt=1e-3)
result_spcst_extexp



################ Constant speciation, extinction variable (expo), different initial parameters ##############

lamb_par<-c(0.3)
mu_par<-c(0.1, 0.001)
result_spcst_extexp2 <- fit_bd(tree,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,dt=1e-3)
result_spcst_extexp2 



##########################################################
################  TEMP-DEPENDENT MODELS  ##################
##########################################################


############### Speciation variable (expo), no extinction ###############

env_data<-read.csv("env_data.csv")
dof<-smooth.spline(env_data[,1], env_data[,2])$df

f.lamb<-function(t,x,y){y[1]*exp(y[2]*x)}
f.mu<-function(t,x,y){0}
lamb_par<-c(0.05,0.01)
mu_par<- c()
result_spexp_noext_temp <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,fix.mu=TRUE,df=dof,dt=1e-3)
result_spexp_noext_temp 



############## Speciation variable (expo), no extinction, different initial parameters ###################

lamb_par<-c(0.3,0.001)
result_spexp_noext_temp2 <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,fix.mu=TRUE,df=dof,dt=1e-3)
result_spexp_noext_temp2



################# Speciation variable (expo), constant extinction ######################

f.lamb<-function(t,x,y){y[1]*exp(y[2]*x)}
f.mu<-function(t,x,y){y[1]}
lamb_par<-c(0.3,0.001)
mu_par<-c(0.01)
result_spexp_extcst_temp <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.mu=TRUE,df=dof,dt=1e-3)
result_spexp_extcst_temp 



################ Speciation variable (expo), constant extinction, different initial parameters ############

lamb_par<-c(0.1,0.01)
mu_par<-c(0.001)
result_spexp_extcst_temp2 <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.mu=TRUE,df=dof,dt=1e-3)
result_spexp_extcst_temp2 



################# Speciation variable (expo), extinction variable (expo) #########################

f.lamb<-function(t,x,y){y[1]*exp(y[2]*x)}
f.mu<-function(t,x,y){y[1]*exp(y[2]*x)}
lamb_par<-c(0.1,0.01)
mu_par<-c(0.01,0.01)
result_spexp_extexp_temp <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,df=dof,dt=1e-3)
result_spexp_extexp_temp 



############### Speciation variable (expo), extinction variable (expo), different initial parameters ##################

lamb_par<-c(0.3,0.001)
mu_par<-c(0.01,0.001)
result_spexp_extexp_temp2 <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,df=dof,dt=1e-3)
result_spexp_extexp_temp2 




#################### Constant speciation, extinction variable (expo) #########################

f.lamb <-function(t,x,y){y[1]}
f.mu <-function(t,x,y){y[1] * exp(y[2] * x)}
lamb_par<-c(0.3)
mu_par<-c(0.01,0.01)
result_spcst_extexp_temp <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,df=dof,dt=1e-3)
result_spcst_extexp_temp



################# Constant speciation, extinction variable (expo), different initial parameters ##################

mu_par<-c(0.1,0.001)
lamb_par<-c(0.5)
result_spcst_extexp_temp2 <- fit_env(tree,env_data,tot_time,f.lamb,f.mu,lamb_par,mu_par,f=24/36,cst.lamb=TRUE,df=dof,dt=1e-3)
result_spcst_extexp_temp2 



################# Plotting results for top 3 models ########################



plot_fit_bd(result_spexp_noext,tot_time)
plot_fit_bd(result_spcstnoext,tot_time)
plot_fit_env(result_spexp_noext_temp,env_data,tot_time)