# Homework assignment 2, CROPS545 Statistical Genomics by Ryan Oliveira

#Before anything we set the working directory (for my computer)
setwd("~/R/RyanDirectory")

#Part 1: Imputing missing values with stochastic imputation

#First, we create a formula for stochastic imputation
StochasticImpute = function(X){
  n=nrow(X)
  m=ncol(X)
  fn=colSums(X, na.rm=T) #Sum of genotypes for all individuals 
  fc=colSums(floor(X/3+1),na.rm=T) #count number of non missing individuals
  fa=fn/(2*fc) #Frequency of allele "2"
  for(i in 1:m){
    index.a=runif(n)<fa[i]
    index.na=is.na(X[,i])
    index.m2=index.a  &  index.na
    index.m0=!index.a  &  index.na
    X[index.m2,i]=2
    X[index.m0,i]=0
  }
  return(X)} 

#Next, import data
myGD=read.table(file="~/R/RyanDirectory/GAPIT_Tutorial_Data/mdp_numeric.txt")
X.raw=myGD[,-1] #truncates first column
#Note: for the following to work, one of two of the following must be true:
#1) You are using a version of R that automatically makes the first row into column names
#2) You modify your input data to not include column titles as they are non-numeric

#New 5 variables that can be chanegd, X.raw remains original for reference
X1=X.raw
X2=X.raw 
X3=X.raw 
X4=X.raw 
X5=X.raw 

#Set variables to be used to randomize missing values
n=nrow(X.raw)
m=ncol(X.raw)
dp=m*n #dp gives the total # of data points - use for % calculation later
uv=runif(dp) #uniform variable

#We are now ready to set each missing rate
mr1=0.05
mr2=0.10
mr3=0.25
mr4=0.50
mr5=0.75

#Converts each uniform vector into a series of TRUE/FALSE where each TRUE is a missing variable
missing1=uv<mr1
missing2=uv<mr2
missing3=uv<mr3
missing4=uv<mr4
missing5=uv<mr5

#Returns the true rate of missing generated to ensure function worked
length(grep("TRUE",missing1))/length(missing1)
length(grep("TRUE",missing2))/length(missing2)
length(grep("TRUE",missing3))/length(missing3)
length(grep("TRUE",missing4))/length(missing4)
length(grep("TRUE",missing5))/length(missing5)
#It works!

#Creates a matrix of true/false values corresponding to dataframe X in layout
index.m1=matrix(missing1,n,m)
index.m2=matrix(missing2,n,m)
index.m3=matrix(missing3,n,m)
index.m4=matrix(missing4,n,m)
index.m5=matrix(missing5,n,m)

#Changes "missing" (TRUE) values to read "NA" and puts them in matrix/dataframe X
X1[index.m1]=NA
X2[index.m2]=NA
X3[index.m3]=NA
X4[index.m4]=NA
X5[index.m5]=NA

#Generates comparisons of X1 and X5 (lowest and highest missing rates) to raw data
X1[1:10,1:10]
X5[1:10,1:10]

#We are ready to impute each data set

#Gives XI1, 2, etc. as a vector of accuracies for 30 repetitions of the stochastic imputation
XI1=replicate(30,{
  x=StochasticImpute(X1)
  y=cor(X.raw[index.m1],x[index.m1])
  return(y)
})

XI2=replicate(30,{
  x=StochasticImpute(X2)
  y=cor(X.raw[index.m2],x[index.m2])
  return(y)
})

XI3=replicate(30,{
  x=StochasticImpute(X3)
  y=cor(X.raw[index.m3],x[index.m3])
  return(y)
})

XI4=replicate(30,{
  x=StochasticImpute(X4)
  y=cor(X.raw[index.m4],x[index.m4])
  return(y)
})

XI5=replicate(30,{
  x=StochasticImpute(X5)
  y=cor(X.raw[index.m5],x[index.m5])
  return(y)
})

#Averages and standard deviation on each accuracy vector
avg.1=mean(XI1)
sd.1=sd(XI1)

avg.2=mean(XI2)
sd.2=sd(XI2)

avg.3=mean(XI3)
sd.3=sd(XI3)

avg.4=mean(XI4)
sd.4=sd(XI4)

avg.5=mean(XI5)
sd.5=sd(XI5)

#Part 2: Fixed missing rate with variable sampling for stochastic imputation

#Creates modifiable version of X.raw we will work with for #2
X=X.raw

#Fixes missing rate at 25%
mr=0.25
#Convert uniform vector to T/F series
missing=uv<mr
#Make matrix of these T/F values with dimensions of data
index.m=matrix(missing,n,m)
#Replaces data with "NA" where missing index has "TRUE"
X[index.m]=NA
#Making sure it worked correctly
X[1:10,1:10]
length(grep("TRUE",missing))/length(missing)

#We randomly samples approx 75%, 50%, and 25% of rows aka individuals (rounded to nearest whole)
#Then we perform stochastic imputation and report the accuracy as a vector
#This is repeated 30 times for each %
XI.75=replicate(30,{
  sel.rows=sample(nrow(X),round(0.75*nrow(X)))
  subpop=X[sel.rows,]
  subpop.raw=X.raw[sel.rows,]
  x=StochasticImpute(subpop)
  y=cor(subpop.raw[index.m[sel.rows,]],x[index.m[sel.rows,]])
  return(y)
})

XI.50=replicate(30,{
  sel.rows=sample(nrow(X),round(0.50*nrow(X)))
  subpop=X[sel.rows,]
  subpop.raw=X.raw[sel.rows,]
  x=StochasticImpute(subpop)
  y=cor(subpop.raw[index.m[sel.rows,]],x[index.m[sel.rows,]])
  return(y)
})

XI.25=replicate(30,{
  sel.rows=sample(nrow(X),round(0.25*nrow(X)))
  subpop=X[sel.rows,]
  subpop.raw=X.raw[sel.rows,]
  x=StochasticImpute(subpop)
  y=cor(subpop.raw[index.m[sel.rows,]],x[index.m[sel.rows,]])
  return(y)
})

#Reports out average and stand
avg.75=mean(XI.75)
sd.75=sd(XI.75)
avg.50=mean(XI.50)
sd.50=sd(XI.50)
avg.25=mean(XI.25)
sd.25=sd(XI.25)

#Part 3: Imputation with KNN, differing levels of K

#Sets the constants we will need
X=X.raw
n=nrow(X.raw)
m=ncol(X.raw)
dp=m*n 
mr=0.25

#Use KNN to impute at k=2, 5, 10, & 20, repeated 20 times, then reports average and SD of accuracy
#Since KNN is non-random, it will perform the operation exactly the same way each time
#To make each iteration of a loop different, the missing values are assigned in each loop
X.knn2=replicate(10,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X)),k=2)
  y=cor(X.raw[index.m],t(x$data)[index.m])
    return(y)
})
avg.knn2=mean(X.knn2)
sd.knn2=sd(X.knn2)

X.knn5=replicate(20,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X)),k=5)
  y=cor(X.raw[index.m],t(x$data)[index.m])
  return(y)
})
avg.knn5=mean(X.knn5)
sd.knn5=sd(X.knn5)

X.knn10=replicate(20,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X)),k=10)
  y=cor(X.raw[index.m],t(x$data)[index.m])
  return(y)
})
avg.knn10=mean(X.knn10)
sd.knn10=sd(X.knn10)

X.knn20=replicate(20,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X)),k=20)
  y=cor(X.raw[index.m],t(x$data)[index.m])
  return(y)
})
avg.knn20=mean(X.knn20)
sd.knn20=sd(X.knn20)

#Part 4: Comparing stochastic imputation, KNN, and BEAGLE

#Sets constants
X=X.raw
n=nrow(X.raw)
m=ncol(X.raw)
dp=m*n 
mr=0.25

#Imputation with Stochastic Imputation
X.si=replicate(10,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X[index.m]=NA
  x=StochasticImpute(X)
  y=cor(X.raw[index.m],x[index.m])
  return(y)
  X=X.raw
})
avg.si=mean(X.si)
sd.si=sd(X.si)

#Resets X
X=X.raw

#Imputation with KNN
X.knn=replicate(10,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X)),k=10)
  y=cor(X.raw[index.m],t(x$data)[index.m])
  return(y)
  X=X.raw
})
avg.knn=mean(X.knn)
sd.knn=sd(X.knn)

#Re-sets X 
X=X.raw

#Imputation with BEAGLE

#Make running BEAGLE function so the process can be repeated anew on different missing indices
run.beagle=function(X){
#Convert data 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)
write.table(myGD2,file="test.bgl",quote=F,sep="\t",col.name=F,row.name=F)
#Run BEAGLE
system("java -Xmx1g -jar /Users/Ryan/Documents/R/RyanDirectory/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
}

#Looping function for BEAGLE

beagle.accuracy=replicate(10,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X[index.m]=NA
  run.beagle(X)
  y=cor(X.raw[index.m],X.bgl[index.m])
  return(y)
})
avg.bg=mean(beagle.accuracy)
sd.bg=sd(beagle.accuracy)

#Part 5: Re-doing KNN imputation with markers as neighbors and individuals as attributes

#Transposes rows and columns in our dataset
#This ensures markers become "neighbors" and sample animals/plants become "attributes"
X.rawrev=t(X.raw)
X.rev=X.rawrev

#Sets the constants we will need (which are different than 3!)
n=nrow(X.rawrev)
m=ncol(X.rawrev)
dp=m*n 
mr=0.25

#Use KNN to impute similar to problem 3
X.knn2=replicate(20,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X.rev[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X.rev)),k=2)
  y=cor(X.rawrev[index.m],t(x$data)[index.m])
  return(y)
})
avg.knn2=mean(X.knn2)
sd.knn2=sd(X.knn2)

X.rev=X.rawrev

X.knn5=replicate(20,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X.rev[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X.rev)),k=5)
  y=cor(X.rawrev[index.m],t(x$data)[index.m])
  return(y)
})
avg.knn5=mean(X.knn5)
sd.knn5=sd(X.knn5)

X.rev=X.rawrev

X.knn10=replicate(20,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X.rev[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X.rev)),k=10)
  y=cor(X.rawrev[index.m],t(x$data)[index.m])
  return(y)
})
avg.knn10=mean(X.knn10)
sd.knn10=sd(X.knn10)

X.rev=X.rawrev

X.knn20=replicate(20,{
  uv=runif(dp) 
  missing=uv<mr
  index.m=matrix(missing,n,m)
  X.rev[index.m]=NA
  x=impute::impute.knn(as.matrix(t(X.rev)),k=20)
  y=cor(X.rawrev[index.m],t(x$data)[index.m])
  return(y)
})
avg.knn20=mean(X.knn20)
sd.knn20=sd(X.knn20)

#Conducting a paired t-test on the results
#Data is entered manually, but the hard numbers can also be replaced by output variables from #3 & #5
#"a" stores the accuracy averages from #3
a=c(0.538,0.607,0.626,0.630)
#"b" stores the accuracy averages from #5
b=c(0.652,0.680,0.692,0.690)
p=t.test(a,b,paired=TRUE)