---
title: "Lab4_impute"
author: "Zhou"
date: "2/8/2021"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
```

## 1.Stochastic imputation with allele frequency

### 1.1 Data prepare

Missing value simulation

```{r}
#Import data
myGD=read.table(file="http://zzlab.net/GAPIT/data/mdp_numeric.txt",head=T)

X.raw=myGD[,-1] #keep as original for comparison 
X=X.raw # Create a new variable than may be changed later

#Set missing values
n=nrow(X)
m=ncol(X)
dp=m*n #total data points
uv=runif(dp)
hist(uv)

mr=.2 #missing rate
missing=uv<mr
length(missing)
missing[1:10]

#Format indicator as matrix
index.m=matrix(missing,n,m)
dim(index.m)   # total number of TRUE: 173844

#Set missing values as NA
X[index.m]=NA
X.raw[1:5,1:5]
X[1:5,1:5]

```

### 1.2 implication of stochastic imputation

```{r}
#Define Stochastic Impute function
StochasticImpute = function(X){
# X is the genotype data with n samples and m markers
n=nrow(X)
m=ncol(X)

# Calculatet the frequency of allele A
fn=colSums(X, na.rm=T) # sum of genotypes for all individuals to get the number of A.
fc=colSums(floor(X/3+1),na.rm=T) #count number of non missing individuals
fa=fn/(2*fc) #Frequency of allele A.

for(i in 1:m){
  index.na=is.na(X[,i])  # where is the missing value
	index.a=runif(n)<fa[i] # randomly sign positions where is less than the frequency of allele "2"
	index.m2 = index.na & index.a  # For those missing value, if it is less than the frequency of allele "2"
	X[index.m2,i]=2                # then, sign the 2
	index.m0 = index.na & !index.a # For those missing value, if it is not less than the frequency of allele "2"
	X[index.m0,i]=0                # then, sign the 0
}
return(X)} 
```

### 1.3 Accuracy calculation

```{r}
#Impute 
XI= StochasticImpute(X)
#Correlation
accuracy.r=cor(X.raw[index.m], XI[index.m])
#Proportion of match
index.match=X.raw==XI # get the index of SNPs that are correct
index.mm=index.match&index.m # get the index of imputed SNPs that are correct
accuracy.m=length(X[index.mm])/length(X[index.m]) #Proportion of match
accuracy.r
accuracy.m

```

### 1.4 the two type accuracy are correlated

```{r}
nrep=100
myimp=replicate(nrep,{
uv=runif(dp)
#hist(uv)
missing=uv<mr
length(missing)
missing[1:10]
index.m=matrix(missing,n,m)
dim(index.m)
X[index.m]=NA
X.raw[1:5,1:5]
X[1:5,1:5]

#=======================================
#Impute with StochasticImpute
XI= StochasticImpute(X)

#Calcuate accuracy
accuracy.r=cor(X.raw[index.m], XI[index.m])
index.match=X.raw==XI
index.mm=index.match&index.m
accuracy.m=length(X[index.mm])/length(X[index.m])
accuracy.r
accuracy.m
acc=c(accuracy.r, accuracy.m)
})
plot(myimp[1,],myimp[2,])
```

## 2. KNN

### 2.1 KNN Example, SNPs as featues

```{r}
library(impute)
#Impute using marker as features
X.knn= impute.knn(as.matrix((X)), k=10)
# Check the imputed SNPs
plot(X.knn$data[index.m],pch=".")
hist(X.knn$data[index.m])

#change to integer
X.knn.I= X.knn$data
Index1= X.knn$data<1
Index2= X.knn$data>1
X.knn.I[Index1]= 0
X.knn.I[Index2]= 2

#Correlation
accuracy.r=cor(X.raw[index.m], (X.knn.I)[index.m])

#Match proportion
index.match=X.raw== X.knn.I
index.mm=index.match&index.m
accuracy.m=length(X[index.mm])/length(X[index.m])
accuracy.r
accuracy.m
```

### 2.2 Individual as features

```{r}
#Impute using individuals as features
X.knn= impute.knn(as.matrix(t(X)), k=10)

#change to integer
X.knn.I= X.knn$data
Index1= X.knn$data<1
Index2= X.knn$data>1
X.knn.I[Index1]= 0
X.knn.I[Index2]= 2

#Correlation
accuracy.r=cor(X.raw[index.m], t(X.knn.I)[index.m])

#Match proportion
index.match=X.raw== t(X.knn.I)
index.mm=index.match&index.m
accuracy.m=length(X[index.mm])/length(X[index.m])
accuracy.r
accuracy.m

```

## 3. BEAGLE

### 3.1 Install the BEAGLE
Firstly, install the JRE in computer

Secondly, Download the BEAGLE from zzlab website

### 3.2 Run the BEAGLE

```{r}

#Convert to BEAGLE input format
index0=X==0
index1=X==1
index2=X==2
indexna=is.na(X)
X2=X
X2[index0]="A\tA"
X2[index1]="A\tB"
X2[index2]="B\tB"
X2[indexna]="?\t?"
myGD2=cbind("M",myGD[,1],X2)
setwd("/Users/ZhouTang/Documents/WSU_courses/2021_spring/545/2021/Lab/Lab4/lab4_imputation")
write.table(myGD2,file="test.bgl",quote=F,sep="\t",col.name=F,row.name=F)

#Impute with BEAGLE
system("java -Xmx12g -jar /Users/ZhouTang/Documents/WSU_courses/2021_spring/545/2021/Lab/Lab4/lab4_imputation/beagle.jar unphased=test.bgl missing=? out=test1" )

#Convert output format
genotype.full <- read.delim("test1.test.bgl.phased.gz",sep=" ",head=T)
genotype.c=as.matrix(genotype.full[,-(1:2)])
index.A=genotype.c=="A"
index.B=genotype.c=="B"
nr=nrow(genotype.c)
nc=ncol(genotype.c)
genotype.n=matrix(0,nr,nc)
genotype.n[index.A]=0
genotype.n[index.B]=1
n2=ncol(genotype.n)
odd=seq(1,n2-1,2)
even=seq(2,n2,2)
g0=genotype.n[,odd]
g1=genotype.n[,even]
X.bgl=g0+g1

#Impute and calculate correlation
accuracy.r=cor(X.raw[index.m], X.bgl[index.m])
index.match=X.raw==X.bgl
index.mm=index.match&index.m
accuracy.m=length(X[index.mm])/length(X[index.m])
accuracy.r
accuracy.m

# Result comparison
Method=c("Stochastic", "KNN(Ind)", "KNN(MK)", "BEAGLE")
Cor=c(.39,.63,.60,.61)
Match=c(.67,.77,.75,.75)
par(mfrow=c(2,1),mar = c(3,4,1,1))
barplot(Cor, names.arg=Method, ylab="Correlation")
barplot(Match, names.arg=Method, ylab="Match")

```