Fitting hidden rate models_corHMM.Rmd

---
title: "hidden rates models_using_corHMM"
author: "Sridhar Halali"
date: "2023-12-13"
output: html_document
editor_options: 
  chunk_output_type: console
---

# Loading required packages
```{r setup, message=FALSE}
library(phytools)
library(tidyverse)
library(ape)
library(openxlsx)
library(plotrix)
library(corHMM)
library(conflicted)
library(geiger)

```

# Setting the working directory
```{r}
# Mac directory
setwd("")

```

# Importing processed diapause data and the pruned phylogeny (see 'ver2_fitMK & ASR_phytools' script for preprocessing steps)
```{r}
# Importing the pruned diapause classification file
dat_diapause <- read.xlsx("diapause_classification_data_PROCESSED.xlsx")

# Convert all character columns to factor
dat_diapause <- dat_diapause %>% mutate_if(is.character, as.factor)

dat_diapause <- data.frame(dat_diapause, row.names = 1)

# Importing pruned phylogeny
ButterflyPhylo <- read.nexus("butterfly_phylogeny_pruned.trees")
ButterflyPhylo

# Chekcking if tips match
name.check(ButterflyPhylo, dat_diapause)

# Sorting rownames in the order of tip labels
dat_diapause <- dat_diapause[ButterflyPhylo$tip.label,]

# Extracting required variables and assigning column values to rownames
DiapauseFive <- setNames(dat_diapause$diapause_five_state, rownames(dat_diapause)) 
DiapauseThree <- setNames(dat_diapause$diapause_three_state, rownames(dat_diapause)) 
DiapauseBinary <- setNames(dat_diapause$diapause_binary, rownames(dat_diapause)) 

```


# Using corHMM() from corHMM package for fitting hidden rate models
```{r}
# Preparing the data set. Converting factors to numeric as required by the package

# Five state classification 
dat_FiveState <- data.frame(Genus.species= names(DiapauseFive), DiapauseFiveState = DiapauseFive)
dat_FiveState$DiapauseFiveState <- as.numeric(dat_FiveState$DiapauseFiveState)

# Three state classification
dat_ThreeState <- data.frame(Genus.species= names(DiapauseThree), DiapauseThreeState = DiapauseThree)
dat_ThreeState$DiapauseThreeState <- as.numeric(dat_ThreeState$DiapauseThreeState)

# Binary state classification
dat_BinaryState <- data.frame(Genus.species= names(DiapauseBinary),DiapauseBinaryState = DiapauseBinary)
dat_BinaryState$DiapauseBinaryState <- as.numeric(dat_BinaryState$DiapauseBinaryState)

```


# Fitting standard Mk (just compare with phytools) hidden rates models only the all rates different models since this was the best fitting model 
Note that 'root.p' argument is for choosing the root prior. Similar to the analyses carried out in phytools, we are fitting both 'flat' (which is NULL in corHMM) and 'fitzjohn' (which is maddfitz in corHMM)


##* Five state diapause classification 
```{r}
# All rates different model (flat root prior)
mod_5State_ARD_flat <- corHMM(ButterflyPhylo, data= dat_FiveState, rate.cat = 1, rate.mat=NULL, 
                       model = "ARD", node.states = "marginal", root.p=  "NULL", nstarts = 5, n.cores = 4)

# All rates different model (fitzjohn root prior)                         
mod_5State_ARD_fitz <- corHMM(phy= ButterflyPhylo, data= dat_FiveState, rate.cat = 1,  rate.mat=NULL ,
                           model = "ARD", node.states = "marginal", root.p="maddfitz", nstarts = 5, n.cores = 4)
```

##* Three state diapause classification 
```{r}
# flat root prior
mod_3State_ARD_flat <- corHMM(ButterflyPhylo, data= dat_ThreeState, rate.cat = 1, rate.mat=NULL, 
                       model = "ARD", node.states = "marginal", root.p=  "NULL", nstarts = 5, , n.cores = 4)

# fitzjohn root prior
mod_3State_ARD_fitz <- corHMM(ButterflyPhylo, data= dat_ThreeState, rate.cat = 1, rate.mat=NULL, 
                       model = "ARD", node.states = "marginal", root.p=  "maddfitz", nstarts = 5, , n.cores = 4)
```


##* Binary state diapause classification 
```{r}
# flat root prior
mod_2State_ARD_flat <- corHMM(ButterflyPhylo, data= dat_BinaryState, rate.cat = 1, rate.mat=NULL, 
                       model = "ARD", node.states = "marginal", root.p=  "NULL", nstarts = 5, , n.cores = 4)

# fitzjohn root prior
mod_2State_ARD_fitz <- corHMM(ButterflyPhylo, data= dat_BinaryState, rate.cat = 1, rate.mat=NULL, 
                       model = "ARD", node.states = "marginal", root.p=  "maddfitz", nstarts = 5, , n.cores = 4)

```

# Saving fitted objects as .RDS files
```{r}
saveRDS(mod_5State_ARD_flat, file="corHMM_5State_stdMK_ARD_flat.RDS")
saveRDS(mod_5State_ARD_fitz, file="corHMM_5State_stdMK_ARD_fitz.RDS")

saveRDS(mod_3State_ARD_flat, file="corHMM_3State_stdMK_ARD_flat.RDS")
saveRDS(mod_3State_ARD_fitz, file="corHMM_3State_stdMK_ARD_fitz.RDS")

saveRDS(mod_2State_ARD_flat, file="corHMM_2State_stdMK_ARD_flat.RDS")
saveRDS(mod_2State_ARD_fitz, file="corHMM_2State_stdMK_ARD_fitz.RDS")

```


# Fitting full hidden rate model
Hidden rate model with two rate classes with ARD model (ARD/ARD). Transitions between rate classes (R1 & R2) occur at different rates. This model is same as when using fitHRM function in phytools with corHMM_model=T. root.p=NULL is as the flat prior in phytools and root.p=maddfitz is same as the fitzjohn prior in phytools**


##* Five state diapause classification
```{r}
# This is only for checking the model. corHMM's by default full model doesn't require this to be done but is required for custom models
RateCat1_five <- getStateMat4Dat(dat_FiveState)$rate.mat # matrix for rate categories
RateCat1_five

RateCat2_five <- getStateMat4Dat(dat_FiveState)$rate.mat
RateCat2_five

RateClassMat_five <- getRateCatMat(2) # matrix for rate classes
RateClassMat_five


# flat root prior
mod_5state_ARD_full_hrm_flat <- corHMM(phy= ButterflyPhylo, data= dat_FiveState, rate.cat = 2,  
                          rate.mat=NULL, model = "ARD", node.states = "marginal",  
                          root.p="NULL", get.tip.states = T, nstarts = 5, n.cores = 6)

# fitzjohn root prior
mod_5state_ARD_full_hrm_fitz <- corHMM(phy= ButterflyPhylo, data= dat_FiveState, rate.cat = 2,  
                          rate.mat=NULL, model = "ARD", node.states = "marginal",  
                          root.p="maddfitz", get.tip.states = T, nstarts = 5, n.cores = 6)

# Saving the object
saveRDS(mod_5state_ARD_full_hrm_flat, file = "corhmm_5State_full_hrm_flat.RDS")
saveRDS(mod_5state_ARD_full_hrm_fitz, file = "corhmm_5State_full_hrm_fitz.RDS")

saveRDS(mod_5state_ARD_full_hrm_flat, file = "run2_corhmm_5State_full_hrm_flat.RDS")
saveRDS(mod_5state_ARD_full_hrm_fitz, file = "run2_corhmm_5State_full_hrm_fitz.RDS")

```

##* Three state diapause classification
```{r}
# flat root prior
mod_3state_ARD_hrm_full_flat <- corHMM(phy= ButterflyPhylo, data= dat_ThreeState, rate.cat = 2,  
                          rate.mat=NULL, model = "ARD", node.states = "marginal",  
                          root.p="NULL", get.tip.states = T, nstarts = 5, n.cores = 6)

# fitzjohn root prior
mod_3state_ARD_hrm_full_fitz <- corHMM(phy= ButterflyPhylo, data= dat_ThreeState, rate.cat = 2,  
                          rate.mat=NULL, model = "ARD", node.states = "marginal",  
                          root.p="maddfitz", get.tip.states = T, nstarts = 5, n.cores = 6)


saveRDS(mod_3state_ARD_hrm_full_flat, file = "corhmm_3State_full_hrm_flat.RDS")
saveRDS(mod_3state_ARD_hrm_full_fitz, file = "corhmm_3State_full_hrm_fitz.RDS")

saveRDS(mod_3state_ARD_hrm_full_flat, file = "run2_corhmm_3State_full_hrm_flat.RDS")
saveRDS(mod_3state_ARD_hrm_full_fitz, file = "run2_corhmm_3State_full_hrm_fitz.RDS")
```


##* Binary state diapause classification
```{r}
# flat root prior
mod_2state_ARD_hrm_full_flat <- corHMM(phy= ButterflyPhylo, data= dat_BinaryState, rate.cat = 2,  
                          rate.mat=NULL, model = "ARD", node.states = "marginal",  
                          root.p="NULL", get.tip.states = T, nstarts = 5, n.cores = 6)

# fitzjohn root prior
mod_2state_ARD_hrm_full_fitz <- corHMM(phy= ButterflyPhylo, data= dat_BinaryState, rate.cat = 2,  
                          rate.mat=NULL, model = "ARD", node.states = "marginal",  
                          root.p="maddfitz", get.tip.states = T, nstarts = 5, n.cores = 6)

saveRDS(mod_2state_ARD_hrm_full_flat, file = "corhmm_2State_full_hrm_flat.RDS")
saveRDS(mod_2state_ARD_hrm_full_fitz, file = "corhmm_2State_full_hrm_fitz.RDS")

saveRDS(mod_2state_ARD_hrm_full_flat, file = "run2_corhmm_2State_full_hrm_flat.RDS")
saveRDS(mod_2state_ARD_hrm_full_fitz, file = "run2_corhmm_2State_full_hrm_fitz.RDS")
```


# Fitting the 'umbral' hidden rates model
See Methods section for the details on this model

##* Five state diapause classification
```{r}
# Using index matrix of the five state model for creating the index matrix for the umbral model
umbralMat_five <- mod_5state_ARD_full_hrm_flat$index.mat
umbralMat_five

# Setting all values to 0
umbralMat_five[] <- 0
umbralMat_five

# Now modifying the model matrix- no transitions are allowed between the rate classes (R1 & R2)
umbralMatFive_final <- matrix(c(
  0,	1,	0,	0,	0,	2,	0,	0,	0,	0,
  3,	0,	4,	0,	0,	0,	5,	0,	0,	0,
  0,	6,	0,	7,	0,	0,	0,	8,	0,	0,
  0,	0,	9,	0,	10,	0,	0,	0,	11,	0,
  0,	0,	0,	12,	0,	0,	0,	0,	0,	13,
  14, 0, 0, 0 , 0, 0, 0, 0, 0, 0,
  0, 15, 0, 0 , 0, 0, 0, 0, 0, 0,
  0, 0, 16, 0 , 0, 0, 0, 0, 0, 0,
  0, 0, 0, 17, 0, 0, 0, 0, 0, 0,
  0, 0, 0, 0, 18, 0, 0, 0, 0, 0), 
  10, 10, 
  byrow = T,
  dimnames = list(colnames(umbralMat_five),
                  rownames(umbralMat_five)))

umbralMatFive_final

plotMKmodel(umbralMatFive_final, rate.cat = 2)


# Now fitting the model using the flat and fitzjohn root prior
mod_5state_ARD_Umbral_hrm_flat <- corHMM(phy= ButterflyPhylo, data= dat_FiveState, rate.cat = 2,  
                       rate.mat=umbralMatFive_final, model = "ARD", node.states = "marginal",  
                       root.p="NULL", get.tip.states = T, nstarts = 5, n.cores = 6)

mod_5state_ARD_Umbral_hrm_fitz <- corHMM(phy= ButterflyPhylo, data= dat_FiveState, rate.cat = 2,  
                       rate.mat=umbralMatFive_final, model = "ARD", node.states = "marginal",  
                       root.p="maddfitz", get.tip.states = T, nstarts = 5, n.cores = 6)


saveRDS(mod_5state_ARD_Umbral_hrm_flat, file="corhmm_5State_Umbral_hrm_flat.RDS")
saveRDS(mod_5state_ARD_Umbral_hrm_fitz, file="corhmm_5State_Umbral_hrm_fitz.RDS")

saveRDS(mod_5state_ARD_Umbral_hrm_flat, file="run2_corhmm_5State_Umbral_hrm_flat.RDS")
saveRDS(mod_5state_ARD_Umbral_hrm_fitz, file="run2_corhmm_5State_Umbral_hrm_fitz.RDS")
```


##* Three state diapause classification
Using steps as above
```{r}
# Using structure of the three state model matrix for the purposes
umbralMat_three <- mod_3state_ARD_hrm_full_flat$index.mat
umbralMat_three

# Setting all values to 0
umbralMat_three[] <- 0
umbralMat_three

# Now modifying the model matrix- no transitions are allowed between the rate classes (R1 & R2)
umbralMatThree_final <- matrix(c(
  0, 1, 0, 2, 0, 0, 
  3, 0, 4, 0, 5, 0, 
  0, 6, 0, 0, 0, 7, 
  8, 0, 0, 0, 0, 0, 
  0, 9, 0, 0, 0, 0, 
  0, 0, 10, 0, 0, 0),
  byrow = T,
  6, 6,
  dimnames = list(colnames(umbralMat_three),
                  rownames(umbralMat_three)))

umbralMatThree_final

plotMKmodel(umbralMatThree_final, rate.cat = 2)


# Now fitting the model using the flat and fitzjohn root prior
mod_3state_ARD_Umbral_hrm_flat <- corHMM(phy= ButterflyPhylo, data= dat_ThreeState, rate.cat = 2,  
                       rate.mat=umbralMatThree_final, model = "ARD", node.states = "marginal",  
                       root.p="NULL", get.tip.states = T, nstarts = 5, n.cores = 6)

mod_3state_ARD_Umbral_hrm_fitz <- corHMM(phy= ButterflyPhylo, data= dat_ThreeState, rate.cat = 2,  
                       rate.mat=umbralMatThree_final, model = "ARD", node.states = "marginal",  
                       root.p="maddfitz", get.tip.states = T, nstarts = 5, n.cores = 6)

saveRDS(mod_3state_ARD_Umbral_hrm_flat, file="corhmm_3State_Umbral_hrm_flat.RDS")
saveRDS(mod_3state_ARD_Umbral_hrm_fitz, file="corhmm_3State_Umbral_hrm_fitz.RDS")

saveRDS(mod_3state_ARD_Umbral_hrm_flat, file="run2_corhmm_3State_Umbral_hrm_flat.RDS")
saveRDS(mod_3state_ARD_Umbral_hrm_fitz, file="run2_corhmm_3State_Umbral_hrm_fitz.RDS")
```


##* Binary state diapause classification
Using steps as above
```{r}
# Using structure of the five state model matrix for the purposes
umbralMat_binary <- mod_2state_ARD_hrm_full_flat$index.mat
umbralMat_binary

# Setting all values to 0
umbralMat_binary[] <- 0
umbralMat_binary

# now modifying the model matrix- no transitions are allowed between the rate classes (R1 & R2)
umbralMatBinary_final <- matrix(c(
  0, 1, 2, 0, 
  3, 0, 0, 4, 
  5, 0, 0, 0, 
  0, 6, 0, 0),
  byrow = T,
  4, 4,
  dimnames = list(colnames(umbralMat_binary),
                  rownames(umbralMat_binary)
))

umbralMatBinary_final


# Now fitting the model using the flat and fitzjohn root prior

# Binary State
mod_2state_ARD_Umbral_hrm_flat <- corHMM(phy= ButterflyPhylo, data= dat_BinaryState, rate.cat = 2,  
                       rate.mat=umbralMatBinary_final, model = "ARD", node.states = "marginal",  
                       root.p="NULL", get.tip.states = T, nstarts = 5, n.cores = 6)

mod_2state_ARD_Umbral_hrm_fitz <- corHMM(phy= ButterflyPhylo, data= dat_BinaryState, rate.cat = 2,  
                       rate.mat=umbralMatBinary_final, model = "ARD", node.states = "marginal",  
                       root.p="maddfitz", get.tip.states = T, nstarts = 5, n.cores = 6)

saveRDS(mod_2state_ARD_Umbral_hrm_flat, file= "corhmm_2State_Umbral_hrm_flat.RDS")
saveRDS(mod_2state_ARD_Umbral_hrm_fitz, file= "corhmm_2State_Umbral_hrm_fitz.RDS")

saveRDS(mod_2state_ARD_Umbral_hrm_flat, file= "run2_corhmm_2State_Umbral_hrm_flat.RDS")
saveRDS(mod_2state_ARD_Umbral_hrm_fitz, file= "run2_corhmm_2State_Umbral_hrm_fitz.RDS")

```


# Loading all fitted models saved as .RDS files 
```{r}
# Five states
mod_5State_ARD_StdMK_flat <- readRDS("corHMM_5State_stdMK_ARD_flat.RDS")
mod_5state_ARD_full_hrm_flat <- readRDS("corhmm_5State_full_hrm_flat.RDS")
mod_5state_ARD_Umbral_hrm_flat <- readRDS("corhmm_5State_Umbral_hrm_flat.RDS")

mod_5State_ARD_StdMK_fitz <- readRDS("corHMM_5State_stdMK_ARD_fitz.RDS")
mod_5state_ARD_full_hrm_fitz <- readRDS("corhmm_5State_full_hrm_fitz.RDS")
mod_5state_ARD_Umbral_hrm_fitz <- readRDS("corhmm_5State_Umbral_hrm_fitz.RDS")

# Three states
mod_3State_ARD_StdMK_flat <- readRDS("corHMM_3State_stdMK_ARD_flat.RDS")
mod_3state_ARD_full_hrm_flat <- readRDS("corhmm_3State_full_hrm_flat.RDS")
mod_3state_ARD_Umbral_hrm_flat <- readRDS("corhmm_3State_Umbral_hrm_flat.RDS")

mod_3State_ARD_StdMk_fitz <- readRDS("corHMM_3State_stdMK_ARD_fitz.RDS")
mod_3state_ARD_full_hrm_fitz <- readRDS("corhmm_3State_full_hrm_fitz.RDS")
mod_3state_ARD_Umbral_hrm_fitz <- readRDS("corhmm_3State_Umbral_hrm_fitz.RDS")


# Binary states
mod_2State_ARD_StdMK_flat <- readRDS("corHMM_2State_stdMK_ARD_flat.RDS")
mod_2state_ARD_full_hrm_flat <- readRDS("corhmm_2State_full_hrm_flat.RDS")
mod_2state_ARD_Umbral_hrm_flat <- readRDS("corhmm_2State_Umbral_hrm_flat.RDS")

mod_2State_ARD_StdMK_fitz <- readRDS("corHMM_2State_stdMK_ARD_fitz.RDS")
mod_2state_ARD_full_hrm_fitz <- readRDS("corhmm_2State_full_hrm_fitz.RDS")
mod_2state_ARD_Umbral_hrm_fitz <- readRDS("corhmm_2State_Umbral_hrm_fitz.RDS")

```


# Table S5:  Assessing model fits- extracting log likelihood and AIC values
```{r}
ModelFitHRM <- data.frame(
  
  state = rep(c("five", "three", "binary"), times=c(6, 6, 6)),
  
  model = rep(c("Std_Mk", "full_HRM", "Umbral_model"), 2),

  root_prior= rep(c("flat", "maddfitz"), times=c(3, 3)),

  logLik = c(
    # five states
    mod_5State_ARD_StdMK_flat$loglik, 
    mod_5state_ARD_full_hrm_flat$loglik, 
    mod_5state_ARD_Umbral_hrm_flat$loglik, 
    
    mod_5State_ARD_StdMK_fitz$loglik, 
    mod_5state_ARD_full_hrm_fitz$loglik, 
    mod_5state_ARD_Umbral_hrm_fitz$loglik,
    
    # Three states
    mod_3State_ARD_StdMK_flat$loglik, 
    mod_3state_ARD_full_hrm_flat$loglik, 
    mod_3state_ARD_Umbral_hrm_flat$loglik, 
    
    mod_3State_ARD_StdMk_fitz$loglik, 
    mod_3state_ARD_full_hrm_fitz$loglik, 
    mod_3state_ARD_Umbral_hrm_fitz$loglik,
    
    # binary state
    mod_2State_ARD_StdMK_flat$loglik, 
    mod_2state_ARD_full_hrm_flat$loglik, 
    mod_2state_ARD_Umbral_hrm_flat$loglik, 
    
    mod_2State_ARD_StdMK_fitz$loglik, 
    mod_2state_ARD_full_hrm_fitz$loglik, 
    mod_2state_ARD_Umbral_hrm_fitz$loglik),
  
  AIC = c(
    # five states
    mod_5State_ARD_StdMK_flat$AIC, 
    mod_5state_ARD_full_hrm_flat$AIC, 
    mod_5state_ARD_Umbral_hrm_flat$AIC, 
    
    mod_5State_ARD_StdMK_fitz$AIC, 
    mod_5state_ARD_full_hrm_fitz$AIC, 
    mod_5state_ARD_Umbral_hrm_fitz$AIC,
    
    
    # Three states
    mod_3State_ARD_StdMK_flat$AIC, 
    mod_3state_ARD_full_hrm_flat$AIC, 
    mod_3state_ARD_Umbral_hrm_flat$AIC, 
    
    mod_3State_ARD_StdMk_fitz$AIC, 
    mod_3state_ARD_full_hrm_fitz$AIC, 
    mod_3state_ARD_Umbral_hrm_fitz$AIC,
    
    # binary state
    mod_2State_ARD_StdMK_flat$AIC, 
    mod_2state_ARD_full_hrm_flat$AIC, 
    mod_2state_ARD_Umbral_hrm_flat$AIC, 
    
    mod_2State_ARD_StdMK_fitz$AIC, 
    mod_2state_ARD_full_hrm_fitz$AIC, 
    mod_2state_ARD_Umbral_hrm_fitz$AIC))
   
ModelFitHRM

#write.xlsx(ModelFitHRM, file="corhmm_full_hrm_fits_2May24.xlsx")

```


# Plotting transition rates for best fitting hidden rates model (which is 'umbral' hidden rates model)
In principle, for the five state classification, AIC values between the standard MK and umbral hidden rate model are very close. However, we are still choosing the best fitting hidden rates model to represent transition rates and ancestral states using the hidden rates model

```{r}
par(mfrow=c(1,1))

#============= Modifying headers of Q matrix of Five state Umbral model =============#
# Flat root prior
FiveState_Umbral_flat_qMat <- as.Qmatrix(mod_5state_ARD_Umbral_hrm_flat)

colnames(FiveState_Umbral_flat_qMat) <- c("Adult", "Egg", "Larva", "no_diapause", "Pupa", 
                                           "Adult*", "Egg*", "Larva*", "no_diapause*", "Pupa*")

rownames(FiveState_Umbral_flat_qMat) <- c("Adult", "Egg", "Larva", "no_diapause", "Pupa", 
                                           "Adult*", "Egg*", "Larva*", "no_diapause*", "Pupa*")

# Fitzjohn root prior
FiveState_Umbral_fitz_qMat <- as.Qmatrix(mod_5state_ARD_Umbral_hrm_fitz)

colnames(FiveState_Umbral_fitz_qMat) <- c("Adult", "Egg", "Larva", "no_diapause", "Pupa", 
                                           "Adult*", "Egg*", "Larva*", "no_diapause*", "Pupa*")

rownames(FiveState_Umbral_fitz_qMat) <- c("Adult", "Egg", "Larva", "no_diapause", "Pupa", 
                                           "Adult*", "Egg*", "Larva*", "no_diapause*", "Pupa*")

# Figure S5A: Plotting rates
pdf("plotMK_Five_Umbral_flat.pdf", height=7, width=7)
plot(FiveState_Umbral_flat_qMat, show.zeros=F, color=T, spacer=0.2, cex.rates=0.8, 
     lwd=3, main= "hrm_FiveState_Umbral_flat_prior")
dev.off()


# Figure S5D: Plotting rates
pdf("plotMK_Five_Umbral_fitz.pdf", height=7, width=7)
plot(FiveState_Umbral_fitz_qMat, show.zeros=F, color=T, spacer=0.2, cex.rates=0.8, 
     lwd=3, main= "hrm_TFiveState_Umbral_fitz_prior")
dev.off()


#============= Modifying headers of Q matrix of Five state Umbral model =============#

# Flat root prior
ThreeState_Umbral_flat_qMat <- as.Qmatrix(mod_3state_ARD_Umbral_hrm_flat)

colnames(ThreeState_Umbral_flat_qMat) <- c("Adult", "Juvenile", "no_diapause",
                                           "Adult*", "Juvenile*", "no_diapause*")

rownames(ThreeState_Umbral_flat_qMat) <- c("Adult", "Juvenile", "no_diapause",
                                           "Adult*", "Juvenile*", "no_diapause*")


# fitzjohn root prior
ThreeState_Umbral_fitz_qMat <- as.Qmatrix(mod_3state_ARD_Umbral_hrm_fitz)

colnames(ThreeState_Umbral_fitz_qMat) <- c("Adult", "Juvenile", "no_diapause",
                                           "Adult*", "Juvenile*", "no_diapause*")

rownames(ThreeState_Umbral_fitz_qMat) <- c("Adult", "Juvenile", "no_diapause",
                                           "Adult*", "Juvenile*", "no_diapause*")

# Figure S5B: Plotting rates
pdf("plotMK_Three_Umbral_flat.pdf", height=7, width=7)

plot(ThreeState_Umbral_flat_qMat, show.zeros=T, color=T, spacer=0.2, cex.rates=0.9, 
     lwd=3, main= "hrm_ThreeState_Umbral_flat_prior")

dev.off()


# Figure S5E: Plotting rates
pdf("plotMK_Three_Umbral_fitz.pdf", height=7, width=7)

plot(ThreeState_Umbral_fitz_qMat, show.zeros=T, color=T, spacer=0.2, cex.rates=0.9, 
     lwd=3, main= "hrm_ThreeState_Umbral_fitz_prior")

dev.off()


#============= Modifying headers of Q matrix of binary state Umbral model =============#

# Flat root prior
BinaryState_Umbral_flat_qMat <- as.Qmatrix(mod_2state_ARD_Umbral_hrm_flat)

colnames(BinaryState_Umbral_flat_qMat) <- c("Diapause", "no_diapause", "Diapause*", "no_diapause*")
rownames(BinaryState_Umbral_flat_qMat) <- c("Diapause", "no_diapause", "Diapause*", "no_diapause*")

# Fitzjohn root prior
BinaryState_Umbral_fitz_qMat <- as.Qmatrix(mod_2state_ARD_Umbral_hrm_fitz)

colnames(BinaryState_Umbral_fitz_qMat) <- c("Diapause", "no_diapause", "Diapause*", "no_diapause*")
rownames(BinaryState_Umbral_fitz_qMat) <- c("Diapause", "no_diapause", "Diapause*", "no_diapause*")


# Figure S5C: Plotting rates
pdf("plotMK_Binary_Umbral_flat.pdf", height=7, width=7)
plot(BinaryState_Umbral_flat_qMat, show.zeros=T, color=T, spacer=0.2, cex.rates=0.9, 
     lwd=3, main= "hrm_BinaryState_Umbral_flat_prior")

dev.off()


# Figure S5F: Plotting rates
pdf("plotMK_Binary_Umbral_fitz.pdf", height=7, width=7)
plot(BinaryState_Umbral_fitz_qMat, show.zeros=T, color=T, spacer=0.2, cex.rates=0.9, 
     lwd=3, main= "hrm_BinaryState_Umbral_fitz_prior")

dev.off()

```


# Ancestral state estimation usin the 'umbral' hidden rates model

##* Extracting marginal reconstructions and then using hide_hidden_state function (see: http://blog.phytools.org/2023/03/ancestral-character-estimation-under.html.)

```{r}
# Five states- flat prior
FiveASR_flat <- mod_5state_ARD_Umbral_hrm_flat$states

colnames(FiveASR_flat) <- c("Adult", "Egg", "Larva", "no_diapause", "Pupa", 
                            "Adult*", "Egg*", "Larva*", "no_diapause*", "Pupa*")
rownames(FiveASR_flat) <- 1:nrow(FiveASR_flat)


# Five states- fitz prior
FiveASR_fitz <- mod_5state_ARD_Umbral_hrm_fitz$states

colnames(FiveASR_fitz) <- c("Adult", "Egg", "Larva", "no_diapause", "Pupa", 
                            "Adult*", "Egg*", "Larva*", "no_diapause*", "Pupa*")
rownames(FiveASR_fitz) <- 1:nrow(FiveASR_fitz)


# Three states- flat prior
ThreeASR_flat <- mod_3state_ARD_Umbral_hrm_flat$states

colnames(ThreeASR_flat) <- c("Adult", "Juvenile", "no_diapause",
                              "Adult*", "Juvenile*", "no_diapause*")
rownames(ThreeASR_flat) <- 1:nrow(ThreeASR_flat)


# Three states- fitz prior
ThreeASR_fitz <- mod_3state_ARD_Umbral_hrm_fitz$states

colnames(ThreeASR_fitz) <- c("Adult", "Juvenile", "no_diapause",
                              "Adult*", "Juvenile*", "no_diapause*")
rownames(ThreeASR_fitz) <- 1:nrow(ThreeASR_fitz)


# Binary states- flat prior
BinaryASR_flat <- mod_2state_ARD_Umbral_hrm_flat$states
colnames(BinaryASR_flat) <- c("Diapause", "no_diapause", "Diapause*", "no_diapause*")
rownames(BinaryASR_flat) <- 1:nrow(BinaryASR_flat)

# Binary states- fitz prior
BinaryASR_fitz <- mod_2state_ARD_Umbral_hrm_fitz$states
colnames(BinaryASR_fitz) <- c("Diapause", "no_diapause", "Diapause*", "no_diapause*")
rownames(BinaryASR_fitz) <- 1:nrow(BinaryASR_fitz)



#=========== Using function to hide the hidden states =====================#
# see here (this function has been slightly modified corHMM output): http://blog.phytools.org/2023/03/ancestral-character-estimation-under.html

# function to hide hidden states 
hide_hidden<-function(object,...){
  ss<-colnames(object) # getting column names of ancr object
  ss<-ss[-grep("*",ss,fixed=TRUE)] # removing columns containing '*'
  
  # Makes an empty matrix corresponding to column and row names of the ancr object
  anc<-matrix(0,nrow(object),length(ss), dimnames=list(rownames(object),ss)) 
  
  # Loop to get rowsums   
  for(i in 1:length(ss)){
    anc[,ss[i]]<-rowSums(object[,grep(ss[i],
                                          colnames(object))])
  }
  anc
}


# Applying the function
FiveASR_flat_hidden <- hide_hidden(FiveASR_flat)
FiveASR_fitz_hidden <- hide_hidden(FiveASR_fitz)

ThreeASR_flat_hidden <- hide_hidden(ThreeASR_flat)
ThreeASR_fitz_hidden <- hide_hidden(ThreeASR_fitz)

BinaryASR_flat_hidden <- hide_hidden(BinaryASR_flat)
BinaryASR_fitz_hidden <- hide_hidden(BinaryASR_fitz)
```


# Plotting ancestral state estimates

##* Setting plotting parameters
```{r}
# Function for adding concentric showing timepoints on the phylogeny
add_rings <- function(phylo){
  
  T<-max(nodeHeights(phylo))
  tick.spacing<-10
  min.tick<-10
  obj<-axis(1, pos=-3,at=seq(T,min.tick,by=-tick.spacing),cex.axis=2, labels=FALSE)
  
  for(i in 1:length(obj)){
    a1<-0
    a2<-2*pi
    draw.arc(0,0,radius=obj[i],a1,a2,lwd=3,
             col=make.transparent("dodgerblue",0.1))
  }
  
  axis(1,pos=-3,at=seq(T,min.tick,by=-tick.spacing),cex.axis=0.7, labels=FALSE)
  text(obj,rep(-8.5,length(obj)),T-obj,cex=0.8)
  text(mean(obj),-12,"time (Mya)",cex=0.8)
}


# Setting colors for states (it seems like the trick is to set the level factor in the same order as the ancr object. or else it seems like the color for tip label and node labels becomes different)

# Five states
head(FiveASR_flat_hidden)

DiapauseFive <- factor(DiapauseFive, levels= c("Adult", "Egg", "Larva", "no_diapause", "Pupa"))

cols_five <- setNames(c("black", "#7EB0DF", "salmon", "grey", "#5D9979"),
                      c("Adult", "Egg", "Larva", "no_diapause", "Pupa"))
                      
                       
# Three states
head(ThreeASR_flat_hidden)

DiapauseThree <- factor(DiapauseThree, levels = c("Adult", "Juvenile", "no_diapause"))

cols_three <- setNames(c("black", "salmon", "grey"), c("Adult", "Juvenile", "no_diapause"))


# Binary states
head(BinaryASR_flat_hidden)

DiapauseBinary <- factor(DiapauseBinary, levels= c("Diapause", "no_diapause"))

cols_binary <- setNames(c("salmon", "grey"), c("Diapause", "no_diapause"))
                        


# Setting colors to butterfly families
library(viridis)

dat_diapause$family <- as.factor(dat_diapause$family)
family <- setNames(dat_diapause$family, rownames(dat_diapause))

cols_family <- viridis_pal()(length(levels(family)))

```


##* Plotting reconstructions
```{r}
#============= Figure S8A: Five states: flat prior =============#

pdf("corhmm_FiveStates_flat_7Feb24.pdf", height=7, width=7)

par(fg="transparent")
pp_five_flat_ML <- FiveASR_flat_hidden[1:ButterflyPhylo$Nnode,]

plotTree(ButterflyPhylo, ftype="off", type="fan", lwd=1, tips=seq(1, 953, by=919/953),maxY=953)
         
obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
n<-Ntip(ButterflyPhylo)
h<-max(nodeHeights(ButterflyPhylo))
ii<-order(rowSums(pp_five_flat_ML^2), decreasing=TRUE)

par(lend=3)
for(i in 1:Ntip(ButterflyPhylo)){
    cc<-if(obj$xx[i]>0) 2 else -2
    th<-atan(obj$yy[i]/obj$xx[i])
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]-cc*cos(th),
        obj$yy[i]-cc*sin(th),lwd=4,
        col=cols_five[DiapauseFive[ButterflyPhylo$tip.label[i]]])
    
    # Adding rings according to families
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]+cc*cos(th),
        obj$yy[i]+cc*sin(th),lwd=4,
        col=cols_family[family[ButterflyPhylo$tip.label[i]]])
}

for(i in ii){
    plotrix::floating.pie(obj$xx[i+Ntip(ButterflyPhylo)],
        obj$yy[i+Ntip(ButterflyPhylo)],
        radius=if(max(pp_five_flat_ML[i,])<=0.80) 0.03*h else 0.01*h,
        x=pp_five_flat_ML[i,],col=cols_five,border="transparent")
}

par(fg="black")
add_rings(ButterflyPhylo)

legend(x=0.98*par()$usr[1],y=0.9*par()$usr[4],
    levels(DiapauseFive),pch=16, col=cols_five,
    pt.cex=2,bty="n",cex=0.8)

dev.off()



#============= Figure 4D: Five states: fitz prior =============#

pdf("corhmm_FiveStates_fitz_7Feb24.pdf", height=7, width=7)

par(fg="transparent")
pp_five_fitz_ML <- FiveASR_fitz_hidden[1:ButterflyPhylo$Nnode,]

plotTree(ButterflyPhylo, ftype="off", type="fan", lwd=1, tips=seq(1, 953, by=919/953),maxY=953)
         
obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
n<-Ntip(ButterflyPhylo)
h<-max(nodeHeights(ButterflyPhylo))
ii<-order(rowSums(pp_five_fitz_ML^2), decreasing=TRUE)

par(lend=3)
for(i in 1:Ntip(ButterflyPhylo)){
    cc<-if(obj$xx[i]>0) 2 else -2
    th<-atan(obj$yy[i]/obj$xx[i])
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]-cc*cos(th),
        obj$yy[i]-cc*sin(th),lwd=4,
        col=cols_five[DiapauseFive[ButterflyPhylo$tip.label[i]]])
    
    # Adding rings according to families
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]+cc*cos(th),
        obj$yy[i]+cc*sin(th),lwd=4,
        col=cols_family[family[ButterflyPhylo$tip.label[i]]])
}

for(i in ii){
    plotrix::floating.pie(obj$xx[i+Ntip(ButterflyPhylo)],
        obj$yy[i+Ntip(ButterflyPhylo)],
        radius=if(max(pp_five_flat_ML[i,])<=0.80) 0.03*h else 0.01*h,
        x=pp_five_fitz_ML[i,],col=cols_five,border="transparent")
}

par(fg="black")
add_rings(ButterflyPhylo)

legend(x=0.98*par()$usr[1],y=0.9*par()$usr[4],
    levels(DiapauseFive),pch=16, col=cols_five,
    pt.cex=2,bty="n",cex=0.8)

dev.off()


#============= Figure S8B: Three states: flat prior =============#

pdf("corhmm_ThreeStates_flat_7Feb24.pdf", height=7, width=7)

par(fg="transparent")
pp_three_flat_ML <- ThreeASR_flat_hidden[1:ButterflyPhylo$Nnode,]

plotTree(ButterflyPhylo, ftype="off", type="fan", lwd=1, tips=seq(1, 953, by=919/953),maxY=953)
         
obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
n<-Ntip(ButterflyPhylo)
h<-max(nodeHeights(ButterflyPhylo))
ii<-order(rowSums(pp_three_flat_ML^2), decreasing=TRUE)


par(lend=3)
for(i in 1:Ntip(ButterflyPhylo)){
    cc<-if(obj$xx[i]>0) 2 else -2
    th<-atan(obj$yy[i]/obj$xx[i])
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]-cc*cos(th),
        obj$yy[i]-cc*sin(th),lwd=4,
        col=cols_three[DiapauseThree[ButterflyPhylo$tip.label[i]]])
    
    # Adding rings according to families
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]+cc*cos(th),
        obj$yy[i]+cc*sin(th),lwd=4,
        col=cols_family[family[ButterflyPhylo$tip.label[i]]])
}

for(i in ii){
    plotrix::floating.pie(obj$xx[i+Ntip(ButterflyPhylo)],
        obj$yy[i+Ntip(ButterflyPhylo)],
        radius=if(max(pp_three_flat_ML[i,])<=0.80) 0.03*h else 0.01*h,
        x=pp_three_flat_ML[i,],col=cols_three,border="transparent")
}

par(fg="black")
add_rings(ButterflyPhylo)

legend(x=0.98*par()$usr[1],y=0.9*par()$usr[4],
    levels(DiapauseThree),pch=16, col=cols_three,
    pt.cex=2,bty="n",cex=0.8)

dev.off()


#============= Figure 4E: Three states: fitzjohn prior =============#

pdf("corhmm_ThreeStates_fitz_7Feb24.pdf", height=7, width=7)

par(fg="transparent")
pp_three_fitz_ML <- ThreeASR_fitz_hidden[1:ButterflyPhylo$Nnode,]

plotTree(ButterflyPhylo, ftype="off", type="fan", lwd=1, tips=seq(1, 953, by=919/953),maxY=953)
         
obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
n<-Ntip(ButterflyPhylo)
h<-max(nodeHeights(ButterflyPhylo))
ii<-order(rowSums(pp_three_fitz_ML^2), decreasing=TRUE)


par(lend=3)
for(i in 1:Ntip(ButterflyPhylo)){
    cc<-if(obj$xx[i]>0) 2 else -2
    th<-atan(obj$yy[i]/obj$xx[i])
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]-cc*cos(th),
        obj$yy[i]-cc*sin(th),lwd=4,
        col=cols_three[DiapauseThree[ButterflyPhylo$tip.label[i]]])
    
    # Adding rings according to families
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]+cc*cos(th),
        obj$yy[i]+cc*sin(th),lwd=4,
        col=cols_family[family[ButterflyPhylo$tip.label[i]]])
}

for(i in ii){
    plotrix::floating.pie(obj$xx[i+Ntip(ButterflyPhylo)],
        obj$yy[i+Ntip(ButterflyPhylo)],
        radius=if(max(pp_three_fitz_ML[i,])<=0.80) 0.03*h else 0.01*h,
        x=pp_three_fitz_ML[i,],col=cols_three,border="transparent")
}

par(fg="black")
add_rings(ButterflyPhylo)

legend(x=0.98*par()$usr[1],y=0.9*par()$usr[4],
    levels(DiapauseThree),pch=16, col=cols_three,
    pt.cex=2,bty="n",cex=0.8)

dev.off()


#============= Figure S8C: Binary states: flat prior =============#

pdf("corhmm_BinaryStates_flat_7Feb24.pdf", height=7, width=7)

par(fg="transparent")
pp_binary_flat_ML <- BinaryASR_flat_hidden[1:ButterflyPhylo$Nnode,]

plotTree(ButterflyPhylo, ftype="off", type="fan", lwd=1, tips=seq(1, 953, by=919/953),maxY=953)
         
obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
n<-Ntip(ButterflyPhylo)
h<-max(nodeHeights(ButterflyPhylo))
ii<-order(rowSums(pp_binary_flat_ML^2), decreasing=TRUE)

par(lend=3)
for(i in 1:Ntip(ButterflyPhylo)){
    cc<-if(obj$xx[i]>0) 2 else -2
    th<-atan(obj$yy[i]/obj$xx[i])
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]-cc*cos(th),
        obj$yy[i]-cc*sin(th),lwd=4,
        col=cols_binary[DiapauseBinary[ButterflyPhylo$tip.label[i]]])
    
    # Adding rings according to families
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]+cc*cos(th),
        obj$yy[i]+cc*sin(th),lwd=4,
        col=cols_family[family[ButterflyPhylo$tip.label[i]]])
}

for(i in ii){
    plotrix::floating.pie(obj$xx[i+Ntip(ButterflyPhylo)],
        obj$yy[i+Ntip(ButterflyPhylo)],
        radius=if(max(pp_binary_flat_ML[i,])<=0.80) 0.03*h else 0.01*h,
        x=pp_binary_flat_ML[i,],col=cols_binary,border="transparent")
}

par(fg="black")
add_rings(ButterflyPhylo)

legend(x=0.98*par()$usr[1],y=0.9*par()$usr[4],
    levels(DiapauseBinary),pch=16, col=cols_binary,
    pt.cex=2,bty="n",cex=0.8)

dev.off()



#============= Figure 4F: Binary states: fitzjohn prior =============#

pdf("corhmm_BinaryStates_fitz_7Feb24.pdf", height=7, width=7)

par(fg="transparent")
pp_binary_fitz_ML <- BinaryASR_fitz_hidden[1:ButterflyPhylo$Nnode,]

plotTree(ButterflyPhylo, ftype="off", type="fan", lwd=1, tips=seq(1, 953, by=919/953),maxY=953)
         
obj<-get("last_plot.phylo",envir=.PlotPhyloEnv)
n<-Ntip(ButterflyPhylo)
h<-max(nodeHeights(ButterflyPhylo))
ii<-order(rowSums(pp_binary_fitz_ML^2), decreasing=TRUE)

par(lend=3)
for(i in 1:Ntip(ButterflyPhylo)){
    cc<-if(obj$xx[i]>0) 2 else -2
    th<-atan(obj$yy[i]/obj$xx[i])
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]-cc*cos(th),
        obj$yy[i]-cc*sin(th),lwd=4,
        col=cols_binary[DiapauseBinary[ButterflyPhylo$tip.label[i]]])
    
    # Adding rings according to families
    segments(obj$xx[i],obj$yy[i],
        obj$xx[i]+cc*cos(th),
        obj$yy[i]+cc*sin(th),lwd=4,
        col=cols_family[family[ButterflyPhylo$tip.label[i]]])
}

for(i in ii){
    plotrix::floating.pie(obj$xx[i+Ntip(ButterflyPhylo)],
        obj$yy[i+Ntip(ButterflyPhylo)],
        radius=if(max(pp_binary_fitz_ML[i,])<=0.80) 0.03*h else 0.01*h,
        x=pp_binary_fitz_ML[i,],col=cols_binary,border="transparent")
}

par(fg="black")
add_rings(ButterflyPhylo)

legend(x=0.98*par()$usr[1],y=0.9*par()$usr[4],
    levels(DiapauseBinary),pch=16, col=cols_binary,
    pt.cex=2,bty="n",cex=0.8)

dev.off()

```