### PCA computation using a panel of 32 AIMs (Huerta-Chagoya et al, 2018)
## Adapted from: Jorsboe et al, Bioinfomatics, 2017:33(19):3148-3150. 


# Read parental genotypes
parental.aims <- read.table(file="Additional_Parental_AIMs.txt", he=T)

# Read your genotypes
# You must previously create a dataframe where rows are IIDs and cols are SNPs. In col=POP, identify your samples (e.g. MEX).
# Genotypes are coded as 0,1,2. Reference allele must be the same of parental file. Headers must be equal. See example file.
mymestizos.aims <- read.table(header=TRUE, text='IID	POP	rs3843249_G	rs9659240_T	rs3755095_T	rs3827760_A	rs10510511_G	rs12495357_G	rs67929453_A	rs10016699_C	rs4833808_T	rs35407_G	rs12521018_C	rs12529753_T	rs9406333_T	rs1858892_A	rs61097563_C	rs12549875_T	rs10116041_A	rs57432666_G	rs734241_G	rs1533224_A	rs11612312_T	rs1409264_C	rs1243370_T	rs4904274_A	rs1426654_A	rs10794640_G	rs59021505_T	rs11657785_T	rs7259453_T	rs1418029_G	rs9975044_C	rs743832_T																																		
                      MEX1	MEX	0	1	1	2	2	2	1	2	2	2	2	2	2	0	0	2	2	2	2	1	2	1	2	0	2	2	2	1	2	2	2	0																																		
                      MEX2	MEX	0	0	1	0	0	0	0	0	0	0	0	0	0	0	2	0	0	1	0	1	0	0	1	2	0	0	0	1	0	1	0	2																																		
                      MEX3	MEX	0	1	0	0	0	0	0	0	0	0	0	0	0	1	2	0	0	0	1	0	0	0	0	0	1	1	0	1	0	0	1	2																																		
                      MEX4	MEX	0	0	0	1	0	1	1	0	0	1	0	0	0	0	0	0	0	0	0	0	0	0	1	1	0	0	1	0	0	0	1	2																																		
                      MEX5	MEX	1	1	0	0	0	0	1	0	0	0	0	1	0	0	1	0	2	1	1	0	0	0	0	2	0	0	1	1	0	1	1	1																																		
                      MEX6	MEX	0	0	1	1	0	0	0	2	0	1	0	1	1	1	2	1	2	2	1	1	0	1	0	0	1	0	1	2	0	1	0	1																																		
                      MEX7	MEX	1	0	0	0	1	1	0	1	0	2	1	1	1	0	0	1	1	2	1	0	0	1	0	1	2	0	0	1	1	2	1	1																																		
                      MEX8	MEX	0	0	1	0	0	0	0	0	0	0	0	0	0	0	2	1	0	0	0	0	0	0	0	1	0	0	1	0	0	1	0	2																																		
                      MEX9	MEX	1	1	1	1	2	1	0	1	0	1	2	1	2	2	2	1	1	0	2	0	2	0	1	2	2	1	2	2	2	2	0	1																																		
                      MEX10	MEX	2	0	1	0	2	1	2	2	0	0	2	0	1	0	1	0	0	2	1	1	0	2	2	1	0	0	1	0	1	0	0	0')

# Merge genotypes
aims <- rbind(parental.aims, mymestizos.aims)
iid <- aims[c(1:2)]; aims <- aims[c(1,3:34)]
aims <- setNames(data.frame(t(aims[,-1])), aims[,1])

# Compute PCs
snp<-nrow(aims)
freq <- rowMeans(aims,na.rm=T)/2    #get allele frequency
M <- (aims-freq*2)/sqrt(freq*(1-freq))    #normalize the genotype matrix
M[is.na(M)]<-0
X<-t(M)%*%as.matrix(M)  #get the (almost) covariance matrix
X<-X/(sum(diag(X))/(snp-1))
E<-eigen(X)

# Extract the top 2 PCs
# Change the interval in case you want to extract more than the top 2 PCs.
# Now, you can use them as covariates in your association analyses.
E.pcs <- data.frame(E$vectors[,1:2]); names(E.pcs) <- c("PC1", "PC2")
E.pcs <- cbind(iid, E.pcs)
write.table(E.pcs, file="myPCs.txt", sep=" ", quote=F, row.names=F)

# Plot
library(ggplot2)
ggplot(E.pcs, aes(x=-PC1, y=PC2, color=POP, shape=POP)) + 
  geom_point()+
  scale_color_manual(values=c("#999999", "#E69F00", "#56B4E9")) +
  theme_bw()