---
title: "Generate F_ST heatmap, F_ST vs geographical distance, and additional analyses including Mantel and Pearson correlation tests"
author: "Mary T. Yohannes"
date: "11/21/2021"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
library(pheatmap)
library(dplyr)
library(RColorBrewer)
library(geosphere)
library(ggplot2)
library(dendextend)
library(tibble)
library(ggpubr)
library(cowplot)
library(ade4)
library(ggtext) 
```

# In order to run the code below, you first need to download the mean F_ST table, produced in HGDP+1kGP tutorial nb4, locally. Make sure to update the input and output file paths in the code below accordingly.  
# This markdown shows how to make a population-level F_ST heatmap, a scatter plot of F_ST vs geographical distance (with and without waypoints), and run additional analyses based off of the heatmap  

# 1a) Set up matrix for heatmap
```{r}
# read-in file with fst values 
ds <- read.table('mean_fst.txt')

# get population labels - 80 populations 
vars <- unique(c(ds$V1, ds$V2)) 

# create all possible pairs including reverse pairs and pairs between same population labels 
# 3160 (original pairs) + 3160 (pairs reversed) + 80 (pairs between the same popn) = 6400 pairs in total
cor.ds <- expand.grid(vars, vars) 

# add the fst value for each pair
# for pairs between same populations, set fst value to 0
cor.ds <- rbind(merge(cor.ds, ds, by.x=c('Var1', 'Var2'), by.y=c('V1', 'V2')), 
                merge(cor.ds, ds, by.x=c('Var2', 'Var1'), by.y=c('V1', 'V2')),
                data.frame(Var1=vars, Var2=vars, V3=0))

# convert table to a correlation data frame (same population pairs on the diagonal)
cor.mat <- reshape2::dcast(cor.ds, Var1~Var2, value.var='V3')

# set first column values (population names) as row names 
rownames(cor.mat) <- cor.mat[,1]

# remove the first column and convert the correlation data frame into a matrix 
cor.mat <- as.matrix(cor.mat[-1])
```

# 1b) Set up annotation data frame and colors for heatmap - region for row
```{r}
# read-in the HGDP+1kGP metadata from gnomAD - one with the updated population labels. The version that we read-in here is after subsetting the updated metadata to the columns we need for plotting  
hgdp_tgp_unfil <- read.delim('gnomad_meta_for_plots.tsv', header=T, sep='\t') 

# set up df for annotation 
hgdp_tgp <- unique(hgdp_tgp_unfil[c("population", "genetic.region")]) # select only population and region columns 
rownames(hgdp_tgp) <- hgdp_tgp[,1] # set populations as rownames 
hgdp_tgp <- hgdp_tgp[-1] # remove the population column 
colnames(hgdp_tgp) <- c("Region") # rename column 

# color for heatmap annotation - obtain from metadata itself for consistency across all plots
reg_colr_df <- unique(hgdp_tgp_unfil[c("genetic.region", "continent.colors")])
reg_colr <- reg_colr_df %>% pull(continent.colors, genetic.region)
annot_reg_colr <- list(Region = reg_colr) # convert into a list  
```

# 1c) Generate heatmap and save as a png 
```{r}
# show plot - heatmap color scale within the same color family - Greens
pheatmap(cor.mat, annotation_col = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, color = brewer.pal(9,"Greens"))

# save as png
pheatmap(cor.mat, filename='fst_heatmap.png', annotation_col = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, color = brewer.pal(9,"Greens"))

# save as pdf 
pheatmap(cor.mat, filename='fst_heatmap.pdf', annotation_col = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, color = brewer.pal(9,"Greens"))
```

# 1d) Plot only the lower triangle of the heatmap and save as a png 
```{r}
# order the rows and columns of the correlation matrix according to the clustering from the heatmap and set the upper triangle to NA values (keeping the original dendogram)
rnames <- rownames(cor.mat)
hc <-  hclust(dist(cor.mat))
cor.mat2 <- cor.mat[hc$order, hc$order]
cor.mat2[upper.tri(cor.mat2)] <- NA
cor.mat2 <- cor.mat2[rnames, rnames]

myColor <- colorRampPalette(c("darkgreen", "white"))(15) # color palette for color bar 

# replot just the lower triangle of the heatmap without reordering the clutser
pheatmap(cor.mat2, cluster_rows = hc, cluster_cols = hc, annotation_row = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, na_col="white", show_rownames = F, treeheight_col=0, color = myColor)

# save plot as png 
pheatmap(cor.mat2, filename='fst_heatmap_lowertri.png', cluster_rows = hc, cluster_cols = hc, annotation_row = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, na_col="white", show_rownames = F, treeheight_col=0, color = myColor)
```

# 2a) Figure out which populations are interspersed in the heatmap and save as a tsv 
```{r}
# generate whole heatmap and assign to a variable 
p_heatmap <- pheatmap(cor.mat, annotation_col = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, color = brewer.pal(9,"Greens"))

# plot heatmap dendrogram 
plot(p_heatmap$tree_row) 

# re-order original data (populations) to match ordering in heatmap
pop_names <- rownames(cor.mat[p_heatmap$tree_row[["order"]],]) 

# convert vector in to a dataframe
pop_df <- data.frame(matrix(unlist(pop_names), nrow=length(pop_names), byrow=TRUE)) 

# add column name
colnames(pop_df) <- "Population" 

# obtain region info for each population using the metadata df used for heatmap annotation - found in code chunk 1b above
hgdp_tgp2 <- hgdp_tgp
hgdp_tgp2$Population <- rownames(hgdp_tgp2) # convert rownames into a column

# merge the two dfs while keeping the order 
# by going through the region column in the merged df, we can figure out which populations are interspersed 
df_merge <- inner_join(pop_df,hgdp_tgp2)
df_merge

# write out table 
write.csv(df_merge, "FST_interspersed_populations_tbl.csv", row.names=FALSE)
```

# 2b) Color dendrogram branches and nodes to see the clusters and interspersed samples more clearly  
```{r}
# generate whole heatmap and assign to a variable 
p_heatmap <- pheatmap(cor.mat, annotation_col = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, color = brewer.pal(9,"Greens"))

# obtain heatmap dendrogram 
dend <- as.dendrogram(p_heatmap$tree_row)

# obtain color info from the metadata df used for heatmap annotation - found in code chunk 1b above  
pop_reg_col <- unique(hgdp_tgp_unfil[c("population", "continent.colors")])
colnames(pop_reg_col)[1] <- "Population" 

# put the color code in the same order as the dendrogram 
col_ordered <- inner_join(df_merge, pop_reg_col)$continent.colors
col_ordered <- factor(col_ordered, unique(col_ordered))

# make a new dendrogram with colored branches and nodes 
dend2 <- dend %>% 
  color_branches(clusters = as.numeric(col_ordered), col = levels(col_ordered)) %>% # color the branches 
  set("labels_colors", as.character(col_ordered)) # color the nodes 
  
# plot dend and save as a png 
png("FST_dend_COLORED.png", width=15,height=15,units="cm",res=300)
par(cex=0.4) 
plot(dend2)
region_col <- unique(hgdp_tgp_unfil[c("continent.colors", "genetic.region")])
legend("topright", legend = region_col$genetic.region,
       fill = region_col$continent.colors, border = region_col$continent.colors, bty = "n", title = "Region", cex=1.8)
dev.off()
```

# 2c) Figure out node lengths from heatmap dendrogram
```{r}
# generate whole heatmap and assign to a variable 
p_heatmap <- pheatmap(cor.mat, annotation_col = hgdp_tgp, annotation_colors = annot_reg_colr, fontsize_row = 5, fontsize_col= 5, angle_col = 90, border_color=NA, height=7, width=9, color = brewer.pal(9,"Greens"))

# obtain heatmap dendrogram 
dend <- as.dendrogram(p_heatmap$tree_row)

# get node heights 
heights <- get_branches_heights(dend, sort=FALSE)

# covert heights in to a df, assign rownames, and save as a csv 
height_df <- data.frame(heights)
rownames(height_df)<-c(1:length(heights))
write.csv(height_df, "FST_dend_heights_tbl.csv")

height_df # show table 
```

# 3) F_ST vs geographical distance between population pairs - used 6371 km for the Earth's radius
# 3a) Set up datasets 
```{r}
# read-in the HGDP+1kGP metadata with updated population labels and geographical coordinates
hgdp_tgp_coor <- read.delim('gnomad_meta_for_plots.tsv', header=T, sep='\t') 

# select only needed columns - order matters here so that the next line of code executes without any issues 
hgdp_tgp_coor <- select(hgdp_tgp_coor, c("population", "genetic.region", "latitude", "longitude", "project"))    

# since the samples within each population have the same coordinates, only keep one for each  
hgdp_tgp_coor <- unique(hgdp_tgp_coor) # the default depends on the first column which is "population" in this case

# read-in file with fst values 
ds <- read.table('mean_fst.txt')
colnames(ds) <- c("Pop1","Pop2", "F_ST") # rename column names 

# join the metadata table with the fst values using population names
# so for each population (in the pair), we will have the coordinates, genetic region and which project they're from 
merged_tbl <-  left_join(ds, hgdp_tgp_coor, by=c("Pop1" = "population")) %>% left_join(hgdp_tgp_coor, by=c("Pop2" = "population")) %>% rename(GR1 = genetic.region.x, Lat1 = latitude.x, Long1 = longitude.x, Project1 = project.x, GR2 = genetic.region.y, Lat2 = latitude.y, Long2 = longitude.y, Project2 = project.y)

## what are each columns in the merged table?
# col 1 = pop1 name
# col 2 = pop2 name
# col 3 = mean F_ST
# col 4 = pop1 genetic region 
# col 5 =  pop1 latitudinal coordinate 
# col 6 = pop1 longitudinal coordinate
# col 7 = pop1 project 
# col 8 = pop2 genetic region 
# col 9 =  pop2 latitudinal coordinate 
# col 10 = pop2 longitudinal coordinate
# col 11 = pop2 project
```

# 3a1) F_ST vs geographical distance between population pairs "without" waypoints 
```{r}
# number of rows in merged table - 3160 population pairs
n_rows <- dim(merged_tbl)[1]

# create an empty matrix with 8 columns to hold final result
mat_wo <- array(0, dim=c(n_rows,8))
colnames(mat_wo) <- c("Pop1", "GR1", "Project1", "Pop2", "GR2", "Project2", "Mean.F_ST", "Distance")

# for each population pair, run the following:
for(i in 1:n_rows){
  
  # pops and coordinates 
  Pop1 <- merged_tbl[i, 1]
  Pop2 <- merged_tbl[i, 2]
  Pop1_coor <- merged_tbl[i, 6:5]
  Pop2_coor <- merged_tbl[i, 10:9]
    
  # get distance between pop1 and pop2
  pop1_pop2 <- distHaversine(Pop1_coor, Pop2_coor, r=6371) 
  
  # Balochi & Brahui, and Palestinian & Druze have geographical distance of 0 because they've got the same coordinates - located on top of each other 
  
  # add final results and other additional information (from the original merged table) to the final matrix
  mat_wo[i,1] = Pop1 # beginning pop
  mat_wo[i,2] = merged_tbl[i, 4] # GR for beginning pop
  mat_wo[i,3] = merged_tbl[i, 7] # project beginning pop came from  
  mat_wo[i,4] = Pop2 # destination pop
  mat_wo[i,5] = merged_tbl[i, 8] # GR for destination pop  
  mat_wo[i,6] = merged_tbl[i,11] # project destination pop came from  
  mat_wo[i,7] = merged_tbl[i,3] # mean F_ST between the population pair 
  mat_wo[i,8] = pop1_pop2 # calculated distance between the population pair
}
```

# 3a2) Plot
```{r}
# convert matrix to df 
df_wo <- as.data.frame(mat_wo)
class(df_wo$Distance) <- "numeric"
class(df_wo$Mean.F_ST) <- "numeric"

# set a True or False column for coloring 
df_wo$Same.GR <- ifelse(df_wo$GR1 == df_wo$GR2, 'Yes', 'No')
df_wo$Project.Status <- ifelse(df_wo$Project1 == df_wo$Project2, df_wo$Project1, 'Different Projects')

# generate plot 
p <- ggplot(df_wo, aes(x = Distance, y = Mean.F_ST, color = Project.Status)) + 
  geom_point(size = 1, alpha = 0.4) + 
  geom_smooth(method=lm) + 
  theme_classic() + 
  labs(color = "Projects", x = "Geographical Distance (km)", y = expression("Mean F"["ST"])) +  
  theme(
    axis.title = element_text(size = 9), 
    legend.title = element_text(size = 9),
    legend.text = element_text(size = 7.5)) +
  guides(color=guide_legend(override.aes=list(size=3)), shape=guide_legend(override.aes=list(size=3)))

# show plot 
p

# save plot as png 
ggsave('fst_vs_distanceinkm_NOWAYPOINTS.png', p, width=7, height=5) 
```

# 3b1) F_ST vs geographical distance between population pairs using "5 waypoints"
```{r}
# create a data frame of all waypoints 
istanbul <- c(41,28) # city in Turkey
cairo <- c(30,31) # capital of Egypt
phnom.penh <- c(11,104) # capital of Cambodia
anadyr <- c(64,177) # town in Russia
prince.rupert <- c(54,-130) # city in British Columbia, Canada

# create a df
waypoints <- as.data.frame(rbind(istanbul, cairo, phnom.penh, anadyr, prince.rupert), stringsAsFactors = FALSE)  

# rename col names
waypoints <- waypoints %>% rownames_to_column("City") %>% rename(Lat=V1, Long=V2) 
```

# 3b2) Function to go through waypoints and calculate distance between pairs of pop
```{r}
pop_waypoint_distance <- function (tbl, waypoints){
  
  # number of rows in the subset of the merged table 
  n_rows = dim(tbl)[1]
  
  # create an empty matrix with 9 columns to hold final result
  mat <- array(0, dim=c(n_rows,9))
  colnames(mat) <- c("Pop1", "GR1", "Project1", "Pop2", "GR2", "Project2", "Mean.F_ST", "Distance", "Path")
  
  # for each population pair, run the following:
  for(i in 1:n_rows){
  
    # pop and coordinates 
    Pop1 <- tbl[i, 1]
    Pop2 <- tbl[i, 2]
    Pop1_coor <- tbl[i, 6:5]
    Pop2_coor <- tbl[i, 10:9]
    
    # no waypoint
    if (nrow(waypoints) == 0) {
    pop1_pop2 <- distHaversine(Pop1_coor, Pop2_coor, r=6371) # get distance between pop1 and pop2
    TotalDistance <- pop1_pop2 
    MyPath <- paste(Pop1, Pop2, sep=" -> ")
      
    # only one waypoint 
    } else if (nrow(waypoints) == 1) {
    pop1_waypoint <- distHaversine(Pop1_coor, waypoints[, 3:2], r=6371) # get distance between pop1 and the one waypoint 
    waypoint_pop2 <- distHaversine(waypoints[, 3:2], Pop2_coor, r=6371) # get distance between the one waypoint and pop2
    TotalDistance <- pop1_waypoint + waypoint_pop2 # add up the distances 
    MyPath <- paste(Pop1, waypoints[,1], Pop2, sep=" -> ") 
    
    # two waypoints 
    } else if (nrow(waypoints) == 2) {
    pop1_waypoint1 <- distHaversine(Pop1_coor, waypoints[1, 3:2], r=6371) # get distance between pop1 and the first waypoint
    waypoint1_waypoint2 <- distHaversine(waypoints[1, 3:2], waypoints[2, 3:2], r=6371) # get distance between the first and second waypoints
    waypoint2_pop2 <- distHaversine(waypoints[2, 3:2], Pop2_coor, r=6371) # get distance between the second waypoint and pop2
    TotalDistance <- pop1_waypoint1 + waypoint1_waypoint2 + waypoint2_pop2 # add up the distances 
    MyPath <- paste(Pop1, waypoints[1,1], waypoints[2,1], Pop2, sep=" -> ") 
    
    # three waypoints
    } else if (nrow(waypoints) == 3) {
    pop1_waypoint1 <- distHaversine(Pop1_coor, waypoints[1, 3:2], r=6371) # get distance between pop1 and the first waypoint
    waypoint1_waypoint2 <- distHaversine(waypoints[1, 3:2], waypoints[2, 3:2], r=6371) # get distance between the first and second waypoints
    waypoint2_waypoint3 <- distHaversine(waypoints[2, 3:2], waypoints[3, 3:2], r=6371) # get distance between the second and third waypoints
    waypoint3_pop2 <- distHaversine(waypoints[3, 3:2], Pop2_coor, r=6371) # get distance between the third waypoint and pop2
    TotalDistance <- pop1_waypoint1 + waypoint1_waypoint2 + waypoint2_waypoint3 + waypoint3_pop2 # add up the distances 
    MyPath <- paste(Pop1, waypoints[1,1], waypoints[2,1], waypoints[3,1], Pop2, sep=" -> ")   
    }
    
    # add the final results and other additional information (from the original merged table) to the final matrix
    mat[i,1] = Pop1 # beginning pop
    mat[i,2] = tbl[i, 4] # GR for beginning pop
    mat[i,3] = tbl[i, 7] # project beginning pop came from  
    mat[i,4] = Pop2 # destination pop
    mat[i,5] = tbl[i, 8] # GR for destination pop  
    mat[i,6] = tbl[i,11] # project destination pop came from  
    mat[i,7] = tbl[i,3] # mean F_ST between the population pair 
    mat[i,8] = TotalDistance # calculated total distance between the population pair
    mat[i,9] = MyPath # path that was taken when calculating distance (through which waypoints?)
    }
  return(mat)
}
```

# 3b3) Set up a matrix with all possible region pairs and each of the waypoints between them 
```{r}
# set up predetermined paths between pairs of regions going through waypoints 
GR <- unique(merged_tbl$GR1) # get all regions

# all possible combinations including pairs between same regions and reverse pairs 
comb <- expand.grid(Pop1 = GR, Pop2 = GR)

for (i in 1:nrow(comb)){ # for each pair, assign waypoints accordingly 
  
  # 1) ONE WAYPOINT 
  
  # 1.a) through phnom.penh (Cambodia)
  if (((comb[i,1] == 'CSA' | comb[i,1] == 'EAS' | comb[i,1] == 'EUR' | comb[i,1] == 'MID') & comb[i,2] == 'OCE') | (comb[i,1] == 'OCE' & (comb[i,2] == 'CSA' | comb[i,2] == 'EAS' | comb[i,2] == 'EUR' | comb[i,2] == 'MID'))){
    comb[i,3] <- 'phnom.penh'
  }
  
  # 1.b) through cairo (Egypt)
  else if (((comb[i,1] == 'CSA' | comb[i,1] == 'EAS') & comb[i,2] == 'AFR') | (comb[i,1] == 'AFR' & (comb[i,2] == 'CSA' | comb[i,2] == 'EAS'))){
    comb[i,3] <- 'cairo'
  }
  
  # 1.c) through istanbul (Turkey)
  else if ((comb[i,1] == 'EUR' & comb[i,2] == 'MID') | (comb[i,1] == 'MID' & comb[i,2] == 'EUR')){
    comb[i,3] <- 'istanbul'
  }
  
  # 2) TWO WAYPOINTS (have to also account for waypoint order for the reverse pairs)
  
  # 2.a.1) through anadyr (Russia) and prince.rupert (Canada) 
  else if ((comb[i,1] == 'CSA' | comb[i,1] == 'EAS' | comb[i,1] == 'EUR' | comb[i,1] == 'MID') & comb[i,2] == 'AMR'){
    comb[i,3] <- 'anadyr' 
    comb[i,4] <- 'prince.rupert'
  }
  
  # 2.a.2) reverse - through prince.rupert (Canada) and anadyr (Russia) 
  else if (comb[i,1] == 'AMR' & (comb[i,2] == 'CSA' | comb[i,2] == 'EAS' | comb[i,2] == 'EUR' | comb[i,2] == 'MID')){
    comb[i,3] <- 'prince.rupert' 
    comb[i,4] <- 'anadyr' 
  }
  
  # 2.b.1) through cairo (Egypt) and istanbul (Turkey)
  else if (comb[i,1] == 'AFR' & comb[i,2] == 'EUR'){
    comb[i,3] <- 'cairo'
    comb[i,4] <- 'istanbul'
  }
  
  # 2.b.2) reverse - through istanbul (Turkey) and cairo (Egypt) 
  else if (comb[i,1] == 'EUR' & comb[i,2] == 'AFR'){
    comb[i,3] <- 'istanbul'
    comb[i,4] <- 'cairo'
  }
  
  
  # 2.c.1) through cairo (Egypt) and phnom.penh (Cambodia)
  else if (comb[i,1] == 'AFR' & comb[i,2] == 'OCE'){
    comb[i,3] <- 'cairo'
    comb[i,4] <- 'phnom.penh'
  }
  
  # 2.c.2) reverse through phnom.penh (Cambodia) and cairo (Egypt) 
  else if (comb[i,1] == 'OCE' & comb[i,2] == 'AFR'){
    comb[i,3] <- 'phnom.penh'
    comb[i,4] <- 'cairo'
  }
  
  
  # 3) THREE WAYPOINTS (have to also account for waypoint order for the reverse pairs)
  
  # 3.a.1) go through prince.rupert (Canada), anadyr (Russia) and cairo (Egypt)
  else if (comb[i,1] == 'AMR' & comb[i,2] == 'AFR'){
    comb[i,3] <- 'prince.rupert'
    comb[i,4] <- 'anadyr'
    comb[i,5] <- 'cairo'
  }
  
  # 3.a.2) reverse - go through cairo (Egypt), anadyr (Russia) and prince.rupert (Canada)
  else if (comb[i,1] == 'AFR' & comb[i,2] == 'AMR'){
    comb[i,3] <- 'cairo'
    comb[i,4] <- 'anadyr'
    comb[i,5] <- 'prince.rupert'
  }

  
  # 3.b.1) go through go through prince.rupert (Canada), anadyr (Russia) and phnom.penh (Cambodia) 
  else if (comb[i,1] == 'AMR' & comb[i,2] == 'OCE'){
    comb[i,3] <- 'prince.rupert'
    comb[i,4] <- 'anadyr'
    comb[i,5] <- 'phnom.penh'
  }
  
  # 3.b.2) reverse - go through go through phnom.penh (Cambodia), anadyr (Russia) and prince.rupert (Canada)
  else if (comb[i,1] == 'OCE' & comb[i,2] == 'AMR'){
    comb[i,3] <- 'phnom.penh' 
    comb[i,4] <- 'anadyr'
    comb[i,5] <- 'prince.rupert'
    
  }
  
  # all other pairs have no waypoints in-between them 
}
```

# 3b4) Calculate distance between each pop pair by going through the specified waypoints path 
```{r}
mat_wth <- array() # empty matrix to hold final result (with waypoints)
t <- 0
for (i in 1:nrow(comb)){
  # subset merged tbl to desired population pairs 
  subset_tbl <- merged_tbl %>% filter(GR1 == comb[i,1] & GR2 == comb[i,2]) 
  
  # pick the corresponding waypoint/s - order matters here for the distance function to follow the right waypoint path when calculating 
  which_waypoints <- c(comb[i,3], comb[i,4], comb[i,5]) 
  use_waypoints <- na.omit(waypoints[match(which_waypoints, waypoints$City), ])
  
  subset_mat <- pop_waypoint_distance(subset_tbl, use_waypoints) 
  t <- t+dim(subset_tbl)
  mat_wth <- rbind(mat_wth, subset_mat)
}

# remove NA row 
mat_wth <- mat_wth[-1,]
```

# 3b5) Annotate final table with additional info and plot 
```{r}
# convert matrix to df 
df_wth <- as.data.frame(mat_wth)
class(df_wth$Distance) <- "numeric"
class(df_wth$Mean.F_ST) <- "numeric"

# set a True or False column for coloring 
df_wth$Same.GR <- ifelse(df_wth$GR1 == df_wth$GR2, 'Yes', 'No')
df_wth$Project.Status <- ifelse(df_wth$Project1 == df_wth$Project2, df_wth$Project1, 'Different Projects')

# remove unnecessary rownames
rownames(df_wth) <- NULL

# generate plot
p_waypoints <- ggplot(df_wth, aes(x = Distance, y = Mean.F_ST, color = Project.Status)) +
  geom_point(size = 1, alpha = 0.4) +
  geom_smooth(method=lm) +
  theme_classic() +
  labs(color = "Projects", x = "Geographical Distance with Waypoints (km)", y = expression("Mean " ~italic(F)[ST])) + 
  theme(
    axis.title = element_text(size = 9),
    legend.title = element_text(size = 9),
    legend.text = element_text(size = 7.5)) +
  guides(color=guide_legend(override.aes=list(size=3)), shape=guide_legend(override.aes=list(size=3)))

# show plot 
p_waypoints

# save plot as png 
ggsave('fst_vs_distanceinkm_WAYPOINTS.png', p_waypoints, width=7, height=5)  
```

# 3c) Arrange the F_ST vs distance plots (with and without waypoints) together for comparison 
```{r}
# put plots together 
p_together <- plot_grid(p + ggtitle("A") + theme(legend.position="none"), p_waypoints + ggtitle("B") + theme(legend.position="none"), nrow=1)

legend <- get_legend(p)

p_together_legend <- plot_grid(p_together, legend, rel_widths = c(4.5, 1)) # add legend

p_together_legend # show plot

# save plots in grid as png 
ggsave('fst_vs_distanceinkm_TOGETHER.png', p_together_legend, width=13, height=4)
```

# 3d) Mantel and Pearson correlation tests (Mantel test - correlation between two matrices typically containing measures of distance - takes a few min to run)

# 3d1) Set up functions 
```{r}
# function to run Mantel test 
calc.mantel <- function (tbl){
  
  # convert into distance matrices
  fst.dists <- dist(tbl$Mean.F_ST)
  dist.dists <- dist(tbl$Distance)
  
  # run test 
  mantel.result <- mantel.rtest(fst.dists, dist.dists, nrepet = 99)
  
  return(mantel.result)
}

# function to run Pearson correlation test 
calc.pearson <- function (tbl){
  
  cor.result <- cor.test(tbl$Mean.F_ST, tbl$Distance)
  
  return(cor.result) 
}

# function to subset tables and compare correlations among projects 
subset.compare <- function (tbl, corr=c("mantel", "pearson")){
  
  # subset tables 
  kgp <- tbl %>% filter(Project1 == "1000 Genomes Project" & Project2 == "1000 Genomes Project")
  hgdp <- tbl %>% filter(Project1 == "Human Genome Diversity Project" & Project2 == "Human Genome Diversity Project")
  cross.proj <- tbl %>% filter((Project1 == "Human Genome Diversity Project" & Project2 == "1000 Genomes Project") | (Project1 == "1000 Genomes Project" & Project2 == "Human Genome Diversity Project"))
  
  if (corr=="mantel"){
    # run Mantel correlation 
    everything.corr <- calc.mantel(tbl)
    kgp.corr <- calc.mantel(kgp)
    hgdp.corr <- calc.mantel(hgdp)
    cross.corr <- calc.mantel(cross.proj)
  } 
  else if (corr=="pearson") {
    # run Pearson correlation 
    everything.corr <- calc.pearson(tbl)
    kgp.corr <- calc.pearson(kgp)
    hgdp.corr <- calc.pearson(hgdp)
    cross.corr <- calc.pearson(cross.proj)
    
  } else{
    print("wrong function input")
  }
  
  return (list("everything" = everything.corr, "kgp" = kgp.corr, "hgdp" = hgdp.corr, "cross" = cross.corr))
}
```

# 3d2) run correlation tests on dataset WITHOUT waypoints 
```{r}
# Pearson 
wo_pearson <- subset.compare(df_wo, "pearson") #11:20

# Mantel 
wo_mantel <- subset.compare(df_wo, "mantel")
```

# 3d3) run correlation tests on dataset WITH waypoints 
```{r}
# Pearson 
wth_pearson <- subset.compare(df_wth, "pearson") 

# Mantel 
wth_mantel <- subset.compare(df_wth, "mantel")
```



















