#---------------------------------------
# 1000 Genomes project data
# Downloaded from https://www.cog-genomics.org/plink/2.0/resources
#---------------------------------------
plink2 --pgen all_phase3_ns.pgen --pvar all_phase3_ns.pvar.zst --psam phase3_corrected.psam --remove deg2.king.cutoff.out.id --snps-only just-acgt --autosome --maf 0.01 --rm-dup exclude-all --exclude-snp . --exclude bed1 longLD.GRCh37.range --make-pgen --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD --threads 48

plink2 --pfile all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD --missing --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD --threads 24

awk '{if(($4=="A" && $5=="T") || ($4=="T" && $5=="A") || ($4=="C" && $5=="G") || ($4=="G" && $5=="C")) print $3}' all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD.pvar > all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD.ambiguous.exclude

plink2 --pfile all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD --exclude all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD.ambiguous.exclude --make-pgen --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig --threads 24

#-------------------
for i in {1..22}
do set -- $i
    cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=48
#$ -l h_vmem=48G
#$ -N chr$1.ukbb_extract_snp
#$ -cwd
#$ -V
#$ -e /MYPATH/chr$1.ukbb_extract_snp.err
#$ -o /MYPATH/chr$1.ukbb_extract_snp.out
plink2 --bfile /MYPATH/ukb_imp_chr"$1"_info08maf001 --rm-dup exclude-all --extract /MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig.pvar --make-bed --out /MYPATH/ukb_imp_chr$1_info08maf001_downsample1KG --threads 48
EOS
done

cat ukb_imp_chr1_info08maf001_downsample1KG.bim ukb_imp_chr2_info08maf001_downsample1KG.bim ukb_imp_chr3_info08maf001_downsample1KG.bim ukb_imp_chr4_info08maf001_downsample1KG.bim ukb_imp_chr5_info08maf001_downsample1KG.bim ukb_imp_chr6_info08maf001_downsample1KG.bim ukb_imp_chr7_info08maf001_downsample1KG.bim ukb_imp_chr8_info08maf001_downsample1KG.bim ukb_imp_chr9_info08maf001_downsample1KG.bim ukb_imp_chr10_info08maf001_downsample1KG.bim ukb_imp_chr11_info08maf001_downsample1KG.bim ukb_imp_chr12_info08maf001_downsample1KG.bim ukb_imp_chr13_info08maf001_downsample1KG.bim ukb_imp_chr14_info08maf001_downsample1KG.bim ukb_imp_chr15_info08maf001_downsample1KG.bim ukb_imp_chr16_info08maf001_downsample1KG.bim ukb_imp_chr17_info08maf001_downsample1KG.bim ukb_imp_chr18_info08maf001_downsample1KG.bim ukb_imp_chr19_info08maf001_downsample1KG.bim ukb_imp_chr20_info08maf001_downsample1KG.bim ukb_imp_chr21_info08maf001_downsample1KG.bim ukb_imp_chr22_info08maf001_downsample1KG.bim > ukb_imp_chrauto_info08maf001_downsample1KG.bim 

awk '{if(!($1 ~ /^#/)) print $3}' all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig.pvar > all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig.snplist

#-----------------------------
# R script
#-----------------------------
library(data.table)
ukbb <- fread("/MYPATH/ukb_imp_chrauto_info08maf001_downsample1KG.bim")
kg <- fread("/MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig.snplist")
length(intersect(ukbb$V2, kg$V1))
write.table(intersect(ukbb$V2, kg$V1), "/MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig.ukb_imp_chrauto_info08maf001_downsample1KG.extract", col.names=FALSE, row.names=FALSE, quote=FALSE, sep=" ")
#-----------------------------

plink2 --pfile all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig --extract all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig.ukb_imp_chrauto_info08maf001_downsample1KG.extract --make-pgen --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb 

plink2 --pfile all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb --indep-pairwise 500kb 0.2 --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb --threads 48

plink2 --pfile all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb --extract all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb.prune.in --make-pgen --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned --threads 48

plink2 --pfile all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned --pca 20 allele-wts --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned --threads 48

plink2 --pfile all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned --freq counts --out all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned --threads 48

#---------------------------------------
# Population assignment in UKBB using 1000 Genomes project data
#---------------------------------------
#-------------------
for i in {1..22}
do set -- $i
    cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=48
#$ -l h_vmem=48G
#$ -N chr$1.pcproj
#$ -cwd
#$ -V
#$ -e MYPATH/chr$1.pcproj.err
#$ -o MYPATH/chr$1.pcproj.out
plink2 --bfile /MYPATH/ukb_imp_chr"$1"_info08maf001 --read-freq /MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned.acount --score /MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned.eigenvec.allele 2 5 header-read no-mean-imputation variance-standardize cols=+scoresums --score-col-nums 6-25 --out /MYPATH/ukb_imp_chr$1_info08maf001.1kgpc --rm-dup exclude-all --threads 48
EOS
done

#-------------------
plink2 --pfile /MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned --read-freq /MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned.acount --score /MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned.eigenvec.allele 2 5 header-read no-mean-imputation variance-standardize --score-col-nums 6-25 --out /MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_noambig_ukbb_pruned.1kgpc --threads 48

#-----------------------------
# R script
#-----------------------------
library(data.table)
for (i in 1:22){
  tmp_dat_name = paste0("ukb_imp_chr",i,"_info08maf001.1kgpc.sscore")
  tmp_dat <- fread(tmp_dat_name)
  colnames(tmp_dat) <- c(paste0("FID","_chr",i), "IID", paste0(c("ALLELE_CT", "NAMED_ALLELE_DOSAGE_SUM", "PC1_AVG", "PC2_AVG", "PC3_AVG", "PC4_AVG", "PC5_AVG", "PC6_AVG", "PC7_AVG", "PC8_AVG", "PC9_AVG", "PC10_AVG", "PC11_AVG", "PC12_AVG", "PC13_AVG", "PC14_AVG", "PC15_AVG", "PC16_AVG", "PC17_AVG", "PC18_AVG", "PC19_AVG", "PC20_AVG", "PC1_SUM", "PC2_SUM", "PC3_SUM", "PC4_SUM", "PC5_SUM", "PC6_SUM", "PC7_SUM", "PC8_SUM", "PC9_SUM", "PC10_SUM", "PC11_SUM", "PC12_SUM", "PC13_SUM", "PC14_SUM", "PC15_SUM", "PC16_SUM", "PC17_SUM", "PC18_SUM", "PC19_SUM", "PC20_SUM"),"_chr",i))
  assign(paste0("dat_chr",i),tmp_dat)
}

dat_raw <- Reduce(function(x,y) merge(x,y,by="IID",all=FALSE), list(dat_chr1, dat_chr2, dat_chr3, dat_chr4, dat_chr5, dat_chr6, dat_chr7, dat_chr8, dat_chr9, dat_chr11, dat_chr12, dat_chr13, dat_chr14, dat_chr15, dat_chr16, dat_chr17, dat_chr18, dat_chr19, dat_chr20, dat_chr21, dat_chr22))

dat_raw_back <- dat_raw

dat_raw <- cbind(dat_raw, rowSums(dat_raw[,c("ALLELE_CT_chr1", "ALLELE_CT_chr2", "ALLELE_CT_chr3", "ALLELE_CT_chr4", "ALLELE_CT_chr5", "ALLELE_CT_chr6", "ALLELE_CT_chr7", "ALLELE_CT_chr8", "ALLELE_CT_chr9", "ALLELE_CT_chr11", "ALLELE_CT_chr12", "ALLELE_CT_chr13", "ALLELE_CT_chr14", "ALLELE_CT_chr15", "ALLELE_CT_chr16", "ALLELE_CT_chr17", "ALLELE_CT_chr18", "ALLELE_CT_chr19", "ALLELE_CT_chr20", "ALLELE_CT_chr21", "ALLELE_CT_chr22")]))
colnames(dat_raw)[length(names(dat_raw))] <- "ALLELE_CT_all"

for (i in 1:20){
dat_raw <- cbind(dat_raw, rowSums(dat_raw[,c(paste0("PC",i,"_SUM_chr1"), paste0("PC",i,"_SUM_chr2"), paste0("PC",i,"_SUM_chr3"), paste0("PC",i,"_SUM_chr4"), paste0("PC",i,"_SUM_chr5"), paste0("PC",i,"_SUM_chr6"), paste0("PC",i,"_SUM_chr7"), paste0("PC",i,"_SUM_chr8"), paste0("PC",i,"_SUM_chr9"), paste0("PC",i,"_SUM_chr11"), paste0("PC",i,"_SUM_chr12"), paste0("PC",i,"_SUM_chr13"), paste0("PC",i,"_SUM_chr14"), paste0("PC",i,"_SUM_chr15"), paste0("PC",i,"_SUM_chr16"), paste0("PC",i,"_SUM_chr17"), paste0("PC",i,"_SUM_chr18"), paste0("PC",i,"_SUM_chr19"), paste0("PC",i,"_SUM_chr20"), paste0("PC",i,"_SUM_chr21"), paste0("PC",i,"_SUM_chr22"))])) 
colnames(dat_raw)[length(names(dat_raw))] <- paste0("projPC_sum",i)
tmp_avg  <- dat_raw[,paste0("projPC_sum",i)]/dat_raw[,"ALLELE_CT_all"]
dat_raw <- cbind(dat_raw, tmp_avg)
colnames(dat_raw)[length(names(dat_raw))] <- paste0("projPC_avg",i)
}

output <- dat_raw[,c("IID", "ALLELE_CT_all", "projPC_sum1", "projPC_avg1", "projPC_sum2", "projPC_avg2", "projPC_sum3", "projPC_avg3", "projPC_sum4", "projPC_avg4", "projPC_sum5", "projPC_avg5", "projPC_sum6", "projPC_avg6", "projPC_sum7", "projPC_avg7", "projPC_sum8", "projPC_avg8", "projPC_sum9", "projPC_avg9", "projPC_sum10", "projPC_avg10", "projPC_sum11", "projPC_avg11", "projPC_sum12", "projPC_avg12", "projPC_sum13", "projPC_avg13", "projPC_sum14", "projPC_avg14", "projPC_sum15", "projPC_avg15", "projPC_sum16", "projPC_avg16", "projPC_sum17", "projPC_avg17", "projPC_sum18", "projPC_avg18", "projPC_sum19", "projPC_avg19", "projPC_sum20", "projPC_avg20")]

write.table(output, "/MYPATH/ukb_imp_chrauto_info08maf001.1kgpc.projpc", col.names=TRUE, row.names=FALSE, sep =" ", quote=FALSE)

#-------------------
library(ggplot2)
library(ggpubr)

pcproj_1kg <- read.table("/MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_pruned.1kgpc.sscore", header=FALSE, stringsAsFactors=FALSE)
colnames(pcproj_1kg) <- c("IID", "superpop", "pop", "ALLELE_CT", "NAMED_ALLELE_DOSAGE_SUM", paste0("pc",seq(1,20,1)))
pcproj_1kg$IID <- factor(pcproj_1kg$IID)
pcproj_1kg$proj <- "YES"
pcproj_1kg <- pcproj_1kg[,c("IID","superpop","pop","proj",paste0("pc",seq(1,20,1)))]
pcproj_1kg <- pcproj_1kg[-which(pcproj_1kg$IID %in% c("NA20314", "HG00731", "HG00732", "HG01880", "HG01882", "HG01944", "HG02497", "NA20320", "NA20321")),]
pcproj_1kg_nocorrect <- pcproj_1kg

eval_1kg <- read.table("/MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_pruned.eigenval", header=FALSE)
colnames(eval_1kg) <- "eval"

pc_1kg <- read.table("/MYPATH/all_phase3_ns_unrelated_snponly_auto_maf01_nolongLD_pruned.eigenvec", header=FALSE, stringsAsFactors=FALSE)
colnames(pc_1kg) <- c("IID", paste0("pc",seq(1,20,1)))
pc_1kg$IID <- factor(pc_1kg$IID)
pc_1kg$superpop[pcproj_1kg$IID] <- pcproj_1kg$superpop[pcproj_1kg$IID]
pc_1kg$pop[pcproj_1kg$IID] <- pcproj_1kg$pop[pcproj_1kg$IID]
pc_1kg$proj <- "NO"
pc_1kg <- pc_1kg[,c("IID","superpop","pop","proj",paste0("pc",seq(1,20,1)))]
pc_1kg <- pc_1kg[-which(pc_1kg$IID %in% c("NA20314", "HG00731", "HG00732", "HG01880", "HG01882", "HG01944", "HG02497", "NA20320", "NA20321")),]

# correct pcproj_1kg
pcproj_1kg$pc1 <- -1*pcproj_1kg$pc1/(eval_1kg$eval[1]**0.5/2*0.94)
pcproj_1kg$pc2 <- -1*pcproj_1kg$pc2/(eval_1kg$eval[2]**0.5/2*0.94)
pcproj_1kg$pc3 <- -1*pcproj_1kg$pc3/(eval_1kg$eval[3]**0.5/2*0.94)
pcproj_1kg$pc4 <- -1*pcproj_1kg$pc4/(eval_1kg$eval[4]**0.5/2*0.94)
pcproj_1kg$pc5 <- -1*pcproj_1kg$pc5/(eval_1kg$eval[5]**0.5/2*0.94)
pcproj_1kg$pc6 <- -1*pcproj_1kg$pc6/(eval_1kg$eval[6]**0.5/2*0.94)
pcproj_1kg$pc7 <- -1*pcproj_1kg$pc7/(eval_1kg$eval[7]**0.5/2*0.94)
pcproj_1kg$pc8 <- -1*pcproj_1kg$pc8/(eval_1kg$eval[8]**0.5/2*0.94)
pcproj_1kg$pc9 <- -1*pcproj_1kg$pc9/(eval_1kg$eval[9]**0.5/2*0.94)
pcproj_1kg$pc10 <- -1*pcproj_1kg$pc10/(eval_1kg$eval[10]**0.5/2*0.94)

dat_1kg_wide <- merge(pc_1kg,pcproj_1kg,by="IID",all=FALSE)

ggplot(data=dat_1kg_wide, aes(x=pc1.x, y=pc1.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc2.x, y=pc2.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc3.x, y=pc3.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc4.x, y=pc4.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc5.x, y=pc5.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc6.x, y=pc6.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc7.x, y=pc7.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc8.x, y=pc8.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc9.x, y=pc9.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

ggplot(data=dat_1kg_wide, aes(x=pc10.x, y=pc10.y)) +
geom_point() + 
geom_smooth(method="lm", se=TRUE, fullrange=FALSE, level=0.95) + 
stat_regline_equation() 

#-------------------
ggplot(data=dat_1kg_wide_qc, aes(x=pc1.y, y=pc2.y, color=superpop.y)) +
geom_point() 

ggplot(data=dat_1kg_wide_qc, aes(x=pc3.y, y=pc4.y, color=superpop.x)) +
geom_point() 

ggplot(data=dat_1kg_wide_qc, aes(x=pc5.y, y=pc6.y, color=superpop.x)) +
geom_point() 

ggplot(data=dat_1kg_wide_qc, aes(x=pc7.y, y=pc8.y, color=superpop.x)) +
geom_point() 

ggplot(data=dat_1kg_wide_qc, aes(x=pc9.y, y=pc10.y, color=superpop.x)) +
geom_point() 

#-----------------------------
# Random forest population assignment
#-----------------------------
pcproj_ukb <- read.table("/MYPATH/ukb_imp_chrauto_info08maf001_downsample1KGpruned.1kgpc.sscore", header=FALSE)
colnames(pcproj_ukb) <- c("FID", "IID", "ALLELE_CT", "NAMED_ALLELE_DOSAGE_SUM", paste0("pc",seq(1,20,1)))
pcproj_ukb$proj <- "YES"
pcproj_ukb$superpop <- "UKBB"
pcproj_ukb$pop <- "UKBB"
pcproj_ukb <- pcproj_ukb[,c("IID","superpop","pop","proj",paste0("pc",seq(1,20,1)))]
pcproj_ukb <- pcproj_ukb[which(pcproj_ukb$IID > 0),]

pcproj_1kg_nocorrect$IID <- as.character(pcproj_1kg_nocorrect$IID)
dat_long <- rbind(pcproj_1kg_nocorrect,pcproj_ukb)

library(randomForest)
library(tidyverse)
set.seed(0)
forest.fit <- randomForest(as.factor(superpop) ~ pc1+pc2+pc3+pc4+pc5+pc6, data = dat_long[which(dat_long$superpop!="UKBB"),], importance = TRUE, ntree = 100) 
forest.fit
data_long_forest <- cbind.data.frame(dat_long,predict(forest.fit, dat_long, type='prob'))
data_long_forest <- cbind.data.frame(data_long_forest,predict(forest.fit, dat_long, type='response'))
colnames(data_long_forest)[ncol(data_long_forest)] <- "predsuperpop"

data_long_forest$predsuperpop_qc <- NA
data_long_forest$predsuperpop_qc[which(data_long_forest$EUR > 0.8 & data_long_forest$predsuperpop=="EUR")] <- "EUR"
data_long_forest$predsuperpop_qc[which(data_long_forest$AFR > 0.8 & data_long_forest$predsuperpop=="AFR")] <- "AFR"
data_long_forest$predsuperpop_qc[which(data_long_forest$SAS > 0.8 & data_long_forest$predsuperpop=="SAS")] <- "SAS"
data_long_forest$predsuperpop_qc[which(data_long_forest$EAS > 0.8 & data_long_forest$predsuperpop=="EAS")] <- "EAS"
data_long_forest$predsuperpop_qc[which(data_long_forest$AMR > 0.8 & data_long_forest$predsuperpop=="AMR")] <- "AMR"

# reclassification
table(data_long_forest$predsuperpop, data_long_forest$superpop, useNA="always")

table(data_long_forest$predsuperpop_qc, data_long_forest$superpop, useNA="always")

ggplot(data=data_long_forest[which(data_long_forest$superpop != 'UKBB'),], aes(x=EUR,color=superpop,fill=superpop)) + 
geom_histogram(position="dodge", bins = 10, alpha=0.5)

ggplot(data=data_long_forest[which(data_long_forest$superpop == 'UKBB'),], aes(x=EUR,color=superpop,fill=superpop)) + 
geom_histogram(position="dodge", bins = 10, alpha=0.5)

ggplot(data=data_long_forest[which(data_long_forest$superpop != 'UKBB'),], aes(x=AFR,color=superpop,fill=superpop)) + 
geom_histogram(position="dodge", bins = 10, alpha=0.5)

ggplot(data=data_long_forest[which(data_long_forest$superpop == 'UKBB'),], aes(x=AFR,color=superpop,fill=superpop)) + 
geom_histogram(position="dodge", bins = 10, alpha=0.5)

ggplot(data=data_long_forest[which(data_long_forest$superpop != 'UKBB'),], aes(x=EAS,color=superpop,fill=superpop)) + 
geom_histogram(position="dodge", bins = 10, alpha=0.5)

ggplot(data=data_long_forest[which(data_long_forest$superpop == 'UKBB'),], aes(x=EAS,color=superpop,fill=superpop)) + 
geom_histogram(position="dodge", bins = 10, alpha=0.5)

#-------------------
ggplot(data=data_long_forest[which(data_long_forest$superpop=="UKBB"),], aes(x=pc1, y=pc2, color=predsuperpop_qc)) +
geom_point()

ggplot(data=data_long_forest[which(data_long_forest$superpop=="UKBB" & is.na(data_long_forest$predsuperpop_qc)==FALSE),], aes(x=pc1, y=pc2, color=predsuperpop_qc)) +
geom_point()

ggplot(data=data_long_forest[which(data_long_forest$superpop=="UKBB" & is.na(data_long_forest$predsuperpop_qc)==FALSE),], aes(x=pc3, y=pc4, color=predsuperpop_qc)) +
geom_point()

ggplot(data=data_long_forest[which(data_long_forest$superpop=="UKBB" & is.na(data_long_forest$predsuperpop_qc)==FALSE),], aes(x=pc5, y=pc6, color=predsuperpop_qc)) +
geom_point()

write.table(data_long_forest[which(data_long_forest$superpop=="UKBB"),c("IID", "predsuperpop_qc")],"ukb_imp_superpop.tsv", row.names=FALSE, col.names=TRUE, quote=FALSE, sep="\t")

write.table(data_long_forest[which(data_long_forest$superpop=="UKBB" & data_long_forest$predsuperpop_qc=="EUR"),c("IID","IID")],"ukb_imp_superpop_EUR.tsv", row.names=FALSE, col.names=FALSE, quote=FALSE, sep="\t")
write.table(data_long_forest[which(data_long_forest$superpop=="UKBB" & data_long_forest$predsuperpop_qc=="AFR"),c("IID","IID")],"ukb_imp_superpop_AFR.tsv", row.names=FALSE, col.names=FALSE, quote=FALSE, sep="\t")
write.table(data_long_forest[which(data_long_forest$superpop=="UKBB" & data_long_forest$predsuperpop_qc=="AMR"),c("IID","IID")],"ukb_imp_superpop_AMR.tsv", row.names=FALSE, col.names=FALSE, quote=FALSE, sep="\t")
write.table(data_long_forest[which(data_long_forest$superpop=="UKBB" & data_long_forest$predsuperpop_qc=="SAS"),c("IID","IID")],"ukb_imp_superpop_SAS.tsv", row.names=FALSE, col.names=FALSE, quote=FALSE, sep="\t")
write.table(data_long_forest[which(data_long_forest$superpop=="UKBB" & data_long_forest$predsuperpop_qc=="EAS"),c("IID","IID")],"ukb_imp_superpop_EAS.tsv", row.names=FALSE, col.names=FALSE, quote=FALSE, sep="\t")
#-----------------------------


#-----------------------------
# Run PCA in each population - run1
#-----------------------------
for i in EUR AFR SAS EAS AMR
do set -- $i
    cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=48
#$ -l h_vmem=48G
#$ -N $1.pca
#$ -cwd
#$ -V
#$ -e /MYPATH/$1.pca.err
#$ -o /MYPATH/$1.pca.out
plink2 --bfile /MYPATH/ukb_imp_chrauto_info08maf001_downsample1KGpruned --keep /MYPATH/ukb_imp_superpop_$1.tsv --pca 20 approx --out /MYPATH/ukb_imp_$1_run1 --threads 48
EOS
done

#-----------------------------
# PCA outlier removal - run1
# 5 SD outlier
# REPEAT FOR EACH POP
#-----------------------------
pc <- read.table("ukb_imp_POP_run1.eigenvec",header=FALSE)
colnames(pc) <- c("FID","IID",paste0("pc",seq(1,20,1)))
str(pc)

range = 5
pc$outlier <- "NO"
pc$outlier[which((pc$pc1 > mean(pc$pc1)+range*sd(pc$pc1)) | (pc$pc1 < mean(pc$pc1)-range*sd(pc$pc1)))] <- "YES"
pc$outlier[which((pc$pc2 > mean(pc$pc2)+range*sd(pc$pc2)) | (pc$pc2 < mean(pc$pc2)-range*sd(pc$pc2)))] <- "YES"
pc$outlier[which((pc$pc3 > mean(pc$pc3)+range*sd(pc$pc3)) | (pc$pc3 < mean(pc$pc3)-range*sd(pc$pc3)))] <- "YES"
pc$outlier[which((pc$pc4 > mean(pc$pc4)+range*sd(pc$pc4)) | (pc$pc4 < mean(pc$pc4)-range*sd(pc$pc4)))] <- "YES"
pc$outlier[which((pc$pc5 > mean(pc$pc5)+range*sd(pc$pc5)) | (pc$pc5 < mean(pc$pc5)-range*sd(pc$pc5)))] <- "YES"
pc$outlier[which((pc$pc6 > mean(pc$pc6)+range*sd(pc$pc6)) | (pc$pc6 < mean(pc$pc6)-range*sd(pc$pc6)))] <- "YES"
pc$outlier[which((pc$pc7 > mean(pc$pc7)+range*sd(pc$pc7)) | (pc$pc7 < mean(pc$pc7)-range*sd(pc$pc7)))] <- "YES"
pc$outlier[which((pc$pc8 > mean(pc$pc8)+range*sd(pc$pc8)) | (pc$pc8 < mean(pc$pc8)-range*sd(pc$pc8)))] <- "YES"
pc$outlier[which((pc$pc9 > mean(pc$pc9)+range*sd(pc$pc9)) | (pc$pc9 < mean(pc$pc9)-range*sd(pc$pc9)))] <- "YES"
pc$outlier[which((pc$pc10 > mean(pc$pc10)+range*sd(pc$pc10)) | (pc$pc10 < mean(pc$pc10)-range*sd(pc$pc10)))] <- "YES"

summary(factor(pc$outlier))

ggplot(data=pc, aes(x=pc1, y=pc2, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc3, y=pc4, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc5, y=pc6, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc7, y=pc8, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc9, y=pc10, color=factor(outlier))) +
    geom_point()

write.table(pc[which(pc$outlier=="YES"),c("FID","IID")],"ukb_imp_POP_run1.SDoutlier.remove", col.names=FALSE, row.names=FALSE, sep= "\t", quote=FALSE)

#-----------------------------
# Run PCA in each population - run2
# REPEAT FOR EACH POP
#-----------------------------
for i in EUR AFR SAS EAS AMR
do set -- $i
cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=24
#$ -l h_vmem=96G
#$ -N ukb_imp_$1_run3.pca
#$ -cwd
#$ -V
#$ -e /MYPATH/ukb_imp_$1_run2.pca.err
#$ -o /MYPATH/ukb_imp_$1_run2.pca.out
plink2 \
--bfile /MYPATH/ukb_imp_chrauto_info08maf001_downsample1KGpruned \
--keep /MYPATH/ukb_imp_superpop_$1.tsv \
--remove /MYPATH/ukb_imp_$1_run1.SDoutlier.remove \
--pca 20 approx \
--out /MYPATH/ukb_imp_$1_run2 \
--threads 24
EOS
done

#-----------------------------
# PCA outlier removal - run2
# 5 SD outlier
# REPEAT FOR EACH POP
#-----------------------------
pc <- read.table("ukb_imp_POP_run2.eigenvec",header=FALSE)
colnames(pc) <- c("FID","IID",paste0("pc",seq(1,20,1)))
str(pc)

range = 5
pc$outlier <- "NO"
pc$outlier[which((pc$pc1 > mean(pc$pc1)+range*sd(pc$pc1)) | (pc$pc1 < mean(pc$pc1)-range*sd(pc$pc1)))] <- "YES"
pc$outlier[which((pc$pc2 > mean(pc$pc2)+range*sd(pc$pc2)) | (pc$pc2 < mean(pc$pc2)-range*sd(pc$pc2)))] <- "YES"
pc$outlier[which((pc$pc3 > mean(pc$pc3)+range*sd(pc$pc3)) | (pc$pc3 < mean(pc$pc3)-range*sd(pc$pc3)))] <- "YES"
pc$outlier[which((pc$pc4 > mean(pc$pc4)+range*sd(pc$pc4)) | (pc$pc4 < mean(pc$pc4)-range*sd(pc$pc4)))] <- "YES"
pc$outlier[which((pc$pc5 > mean(pc$pc5)+range*sd(pc$pc5)) | (pc$pc5 < mean(pc$pc5)-range*sd(pc$pc5)))] <- "YES"
pc$outlier[which((pc$pc6 > mean(pc$pc6)+range*sd(pc$pc6)) | (pc$pc6 < mean(pc$pc6)-range*sd(pc$pc6)))] <- "YES"
pc$outlier[which((pc$pc7 > mean(pc$pc7)+range*sd(pc$pc7)) | (pc$pc7 < mean(pc$pc7)-range*sd(pc$pc7)))] <- "YES"
pc$outlier[which((pc$pc8 > mean(pc$pc8)+range*sd(pc$pc8)) | (pc$pc8 < mean(pc$pc8)-range*sd(pc$pc8)))] <- "YES"
pc$outlier[which((pc$pc9 > mean(pc$pc9)+range*sd(pc$pc9)) | (pc$pc9 < mean(pc$pc9)-range*sd(pc$pc9)))] <- "YES"
pc$outlier[which((pc$pc10 > mean(pc$pc10)+range*sd(pc$pc10)) | (pc$pc10 < mean(pc$pc10)-range*sd(pc$pc10)))] <- "YES"

summary(factor(pc$outlier))

ggplot(data=pc, aes(x=pc1, y=pc2, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc3, y=pc4, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc5, y=pc6, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc7, y=pc8, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc9, y=pc10, color=factor(outlier))) +
    geom_point()

write.table(pc[which(pc$outlier=="YES"),c("FID","IID")],"ukb_imp_POP_run2.SDoutlier.remove", col.names=FALSE, row.names=FALSE, sep= "\t", quote=FALSE)

#-----------------------------
# Run PCA in each population - run3
# REPEAT FOR EACH POP
#-----------------------------
for i in EUR AFR SAS EAS AMR
do set -- $i
cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=24
#$ -l h_vmem=96G
#$ -N ukb_imp_$1_run3.pca
#$ -cwd
#$ -V
#$ -e /MYPATH/ukb_imp_$1_run3.pca.err
#$ -o /MYPATH/ukb_imp_$1_run3.pca.out
plink2 \
--bfile /MYPATH/ukb_imp_chrauto_info08maf001_downsample1KGpruned \
--keep /MYPATH/ukb_imp_superpop_$1.tsv \
--remove /MYPATH/ukb_imp_$1_run1.SDoutlier.remove /MYPATH/ukb_imp_$1_run2.SDoutlier.remove \
--pca 20 approx \
--out /MYPATH/ukb_imp_$1_run3 \
--threads 24
EOS
done

#-----------------------------
# PCA outlier removal - run3
# 5 SD outlier
# REPEAT FOR EACH POP
#-----------------------------
pc <- read.table("ukb_imp_POP_run3.eigenvec",header=FALSE)
colnames(pc) <- c("FID","IID",paste0("pc",seq(1,20,1)))
str(pc)

range = 5
pc$outlier <- "NO"
pc$outlier[which((pc$pc1 > mean(pc$pc1)+range*sd(pc$pc1)) | (pc$pc1 < mean(pc$pc1)-range*sd(pc$pc1)))] <- "YES"
pc$outlier[which((pc$pc2 > mean(pc$pc2)+range*sd(pc$pc2)) | (pc$pc2 < mean(pc$pc2)-range*sd(pc$pc2)))] <- "YES"
pc$outlier[which((pc$pc3 > mean(pc$pc3)+range*sd(pc$pc3)) | (pc$pc3 < mean(pc$pc3)-range*sd(pc$pc3)))] <- "YES"
pc$outlier[which((pc$pc4 > mean(pc$pc4)+range*sd(pc$pc4)) | (pc$pc4 < mean(pc$pc4)-range*sd(pc$pc4)))] <- "YES"
pc$outlier[which((pc$pc5 > mean(pc$pc5)+range*sd(pc$pc5)) | (pc$pc5 < mean(pc$pc5)-range*sd(pc$pc5)))] <- "YES"
pc$outlier[which((pc$pc6 > mean(pc$pc6)+range*sd(pc$pc6)) | (pc$pc6 < mean(pc$pc6)-range*sd(pc$pc6)))] <- "YES"
pc$outlier[which((pc$pc7 > mean(pc$pc7)+range*sd(pc$pc7)) | (pc$pc7 < mean(pc$pc7)-range*sd(pc$pc7)))] <- "YES"
pc$outlier[which((pc$pc8 > mean(pc$pc8)+range*sd(pc$pc8)) | (pc$pc8 < mean(pc$pc8)-range*sd(pc$pc8)))] <- "YES"
pc$outlier[which((pc$pc9 > mean(pc$pc9)+range*sd(pc$pc9)) | (pc$pc9 < mean(pc$pc9)-range*sd(pc$pc9)))] <- "YES"
pc$outlier[which((pc$pc10 > mean(pc$pc10)+range*sd(pc$pc10)) | (pc$pc10 < mean(pc$pc10)-range*sd(pc$pc10)))] <- "YES"

summary(factor(pc$outlier))

ggplot(data=pc, aes(x=pc1, y=pc2, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc3, y=pc4, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc5, y=pc6, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc7, y=pc8, color=factor(outlier))) +
    geom_point()
ggplot(data=pc, aes(x=pc9, y=pc10, color=factor(outlier))) +
    geom_point()

my_cols <- c("#FC4E07", "#00AFBB")
pairs(pc[,paste0("pc",seq(1,10,1))], pch = 19,  cex = 0.5, col = my_cols[factor(pc$outlier)], lower.panel=NULL)

write.table(pc[which(pc$outlier=="YES"),c("FID","IID")],"ukb_imp_EUR_run3.SDoutlier.remove", col.names=FALSE, row.names=FALSE, sep= "\t", quote=FALSE)

#-----------------------------
# Run PCA in each population - run4
#-----------------------------
for i in EUR AFR SAS EAS AMR
do set -- $i
cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=24
#$ -l h_vmem=96G
#$ -N ukb_imp_$1_run4.pca
#$ -cwd
#$ -V
#$ -e /MYPATH/ukb_imp_$1_run4.pca.err
#$ -o /MYPATH/ukb_imp_$1_run4.pca.out
plink2 \
--bfile /MYPATH/ukb_imp_chrauto_info08maf001_downsample1KGpruned \
--keep /MYPATH/ukb_imp_superpop_$1.tsv \
--remove /MYPATH/ukb_imp_$1_run1.SDoutlier.remove /MYPATH/ukb_imp_$1_run2.SDoutlier.remove /MYPATH/ukb_imp_$1_run3.SDoutlier.remove \
--pca 20 approx \
--out /MYPATH/ukb_imp_$1_run4 \
--threads 24
EOS
done

#-----------------------------
# PCA final check
#-----------------------------
library(ggplot2)
pc <- read.table("ukb_imp_EUR_run4.eigenvec",header=FALSE)
colnames(pc) <- c("FID","IID",paste0("pc",seq(1,20,1)))
str(pc)

whitebirtish <- read.table("/MYPATH/whiteBritish_IID.keep")

length(whitebirtish$V1) 
length(pc$IID) 
length(intersect(whitebirtish$V1, pc$IID)) 
length(setdiff(whitebirtish$V1, pc$IID))   
length(setdiff(pc$IID, whitebirtish$V1))   

ggplot(data=pc, aes(x=pc1, y=pc2)) + geom_point()
ggplot(data=pc, aes(x=pc3, y=pc4)) + geom_point()
ggplot(data=pc, aes(x=pc5, y=pc6)) + geom_point()
ggplot(data=pc, aes(x=pc7, y=pc8)) + geom_point()
ggplot(data=pc, aes(x=pc9, y=pc10)) + geom_point()
ggplot(data=pc, aes(x=pc11, y=pc12)) + geom_point()
ggplot(data=pc, aes(x=pc13, y=pc14)) + geom_point()
ggplot(data=pc, aes(x=pc15, y=pc16)) + geom_point()
ggplot(data=pc, aes(x=pc17, y=pc18)) + geom_point()
ggplot(data=pc, aes(x=pc19, y=pc20)) + geom_point()

pc <- read.table("ukb_imp_EAS_run4.eigenvec",header=FALSE)
pc <- read.table("ukb_imp_AMR_run4.eigenvec",header=FALSE)
pc <- read.table("ukb_imp_SAS_run4.eigenvec",header=FALSE)
pc <- read.table("ukb_imp_AFR_run4.eigenvec",header=FALSE)
colnames(pc) <- c("FID","IID",paste0("pc",seq(1,20,1)))
str(pc)

pairs(pc[,paste0("pc",seq(1,10,1))], pch = 19,  cex = 0.5, lower.panel=NULL)

# REPEAT FOR EACH POP
ggplot(data=pc, aes(x=pc1, y=pc2)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc1pc2.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc3, y=pc4)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc3pc4.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc5, y=pc6)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc5pc6.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc7, y=pc8)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc7pc8.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc9, y=pc10)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc9pc10.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc11, y=pc12)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc11pc12.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc13, y=pc14)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc13pc14.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc15, y=pc16)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc15pc16.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc17, y=pc18)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc17pc18.pdf", width=6, height=6,units="in")
ggplot(data=pc, aes(x=pc19, y=pc20)) + geom_point()
ggsave("ukb_imp_AMR_run4.pc19pc20.pdf", width=6, height=6,units="in")

#-----------------------------
# Run PCA in each population - final with 40 PCs
#-----------------------------
for i in EUR AFR SAS EAS AMR
do set -- $i
cat <<EOS | qsub
#!/usr/bin/env bash
#$ -q long.q
#$ -l thread=24
#$ -l h_vmem=96G
#$ -N ukb_imp_$1_final.pca
#$ -cwd
#$ -V
#$ -e /MYPATH/ukb_imp_$1_final.pca.err
#$ -o /MYPATH/ukb_imp_$1_final.pca.out
plink2 \
--bfile /MYPATH/ukb_imp_chrauto_info08maf001_downsample1KGpruned \
--keep /MYPATH/ukb_imp_superpop_$1.tsv \
--remove /MYPATH/ukb_imp_$1_run1.SDoutlier.remove /MYPATH/ukb_imp_$1_run2.SDoutlier.remove /MYPATH/ukb_imp_$1_run3.SDoutlier.remove \
--pca 40 approx \
--out /MYPATH/ukb_imp_$1_final \
--threads 24
EOS
done
