#Implement equation 2.10 (for all haplotypes) on p. 45 of Rice's book for selection on two loci with two alleles
#Compute quasi-linkage equilibrium approximation for the dynamics of D (
#use "AB", "Ab", "aB" and "ab" in place of "11", "12", or "22", resp.
#By R. Gomulkiewicz (16 Sep 2013)

##### parameters and initial values############
#number of time steps to evolve
steps <- 200

#initial allele and haplotype frequencies
p0 <- 0.9  #initial frequency of A allele
q0 <- 0.1  #initial frequency of B allele
D0 <- min(p0*(1-q0), (1-p0)*q0)  #positive initial disequilibrium (maximum possible value)
#D0 <- (-1)*min(p0*q0, (1-p0)*(1-q0))  #negative initial disequilibrium (minimum possible value)

#initial haplotype frequencies
x0 <- list(AB = p0*q0 + D0, Ab = p0*(1-q0) - D0, aB = (1-p0)*q0 - D0, ab = (1-p0)*(1-q0) + D0)

#genotypic fitnesses
w <- list(ABAB = 0.95, ABAb = 1, ABaB = 0.95, ABab = 1.05, AbAb = 0.95, AbaB = 1, Abab = 0.95, aBaB = 0.95, aBab = 0.95, abab = 0.95)

#recombination rate
r <- 0.2


######functions of haplotype frequencies ########

#haplotype marginal fitnesses
wAB <- function(x, w){x$AB*w$ABAB + x$Ab*w$ABAb + x$aB*w$ABaB + x$ab*w$ABab}
wAb <- function(x, w){x$AB*w$ABAb + x$Ab*w$AbAb + x$aB*w$AbaB + x$ab*w$Abab}
waB <- function(x, w){x$AB*w$ABaB + x$Ab*w$AbaB + x$aB*w$aBaB + x$ab*w$aBab}
wab <- function(x, w){x$AB*w$ABab + x$Ab*w$Abab + x$aB*w$aBab + x$ab*w$abab}

#mean fitness given genotype frequencies and fertilities
wbar <- function(x, w){x$AB*wAB(x,w) + x$Ab*wAb(x,w) + x$aB*waB(x,w) + x$ab*wab(x,w)}

#linkage disequilibrium
ld <- function(x){x$AB*x$ab - x$Ab*x$aB}


########## evolutionary dynamics ###########

#create a matrix that holds the 4 haplotype frequencies at each time step
#separate columns for x$AB, x$Ab, x$aB, and x$ab
x.out <-matrix(0,nrow=steps+1,ncol=4)

#initial haplotype frequencies
x.out[1,]<- c(x0$AB, x0$Ab, x0$aB, x0$ab)

for(i in 1:steps){
	x <- list(AB = x.out[i,1],Ab = x.out[i,2],aB = x.out[i,3],ab = x.out[i,4])
	w.AB <- wAB(x,w)
	w.Ab <- wAb(x,w)
	w.aB <- waB(x,w)
	w.ab <- wab(x,w)
	w.bar <- wbar(x, w)
	D <- ld(x)
	x.out[i+1,1] <- (x$AB*w.AB - r*D*w$ABab)/w.bar #new freq of AB
	x.out[i+1,2] <- (x$Ab*w.Ab + r*D*w$ABab)/w.bar #new freq of Ab
	x.out[i+1,3] <- (x$aB*w.aB + r*D*w$ABab)/w.bar #new freq of aB
	x.out[i+1,4] <- (x$ab*w.ab - r*D*w$ABab)/w.bar #new freq of ab
	}

#evolutionary dynamics of D, p, and q
D.out <- x.out[,1]*x.out[,4] - x.out[,2]*x.out[,3]
freqs <- cbind((x.out[,1]+x.out[,2]),(x.out[,1]+x.out[,3])) 

#plot dynamics of haplotype frequencies, allele frequencies, and D

quartz() #opens a new graph window on macs
#windows()  #opens a new graph window in Windows---delete or comment out the command above
#X11()  #opens a new graph window in Unix---delete or comment out the two commands above

plot(0:steps,D.out,type = "l", xlab = "generation", ylab = "D",ylim=c(0,D0))

quartz() #opens a new graph window on macs
#windows()  #opens a new graph window in Windows---delete or comment out the command above
#X11()  #opens a new graph window in Unix---delete or comment out the two commands above


matplot(0:steps,x.out,type="l",xlab="generation", ylab="frequency",ylim=c(0,1))  #plot haplotype frequencies
legend("topright",c("x11","x12","x21","x22"),lty=1:4,bty="n",col=1:4)

quartz() #opens a new graph window on macs
#windows()  #opens a new graph window in Windows---delete or comment out the command above
#X11()  #opens a new graph window in Unix---delete or comment out the two commands above

matplot(0:steps,freqs,type="l",xlab="generation", ylab="frequency",ylim=c(0,1))  #plot allele frequencies at both loci
legend("topright",c("p","q"),lty=1:2,bty="n",col=1:2)


####### QLE approximation ##########
p <- freqs[,1]
q <- freqs[,2]
#QLE marginal and mean fitnesses
wAB.qle <- p*q*w$ABAB + p*(1-q)*w$ABAb + (1-p)*q*w$ABaB + (1-p)*(1-q)*w$ABab
wAb.qle <- p*q*w$ABAb + p*(1-q)*w$AbAb + (1-p)*q*w$AbaB + (1-p)*(1-q)*w$Abab
waB.qle <- p*q*w$ABaB + p*(1-q)*w$AbaB + (1-p)*q*w$aBaB + (1-p)*(1-q)*w$aBab
wab.qle <- p*q*w$ABab + p*(1-q)*w$Abab + (1-p)*q*w$aBab + (1-p)*(1-q)*w$abab
wbar.qle <- p*q*wAB.qle + p*(1-q)*wAb.qle + (1-p)*q*waB.qle + (1-p)*(1-q)*wab.qle

#QLE approximation for D (Rice eq. 2.22, p. 51)
d.qle <- p*q*(1-p)*(1-q)*(wAB.qle*wab.qle - wAb.qle*waB.qle)/(r*wbar.qle)

quartz() #opens a new graph window on macs
#windows()  #opens a new graph window in Windows---delete or comment out the command above
#X11()  #opens a new graph window in Unix---delete or comment out the two commands above

matplot(0:steps,cbind(D.out,d.qle),type="l",xlab="generation", ylab="D")  #plot D and its qle approximation
legend("topright",c("D","Dqle"),lty=1:2,bty="n",col=1:2)


############ plots like fig. 2.2 on p. 49 of Rice   ############
quartz() #opens a new graph window on macs
#windows()  #opens a new graph window in Windows---delete or comment out the command above
#X11()  #opens a new graph window in Unix---delete or comment out the two commands above

par(mfrow=c(3,1))

matplot(0:steps,cbind(D.out,d.qle),type="l",xlab="generation", ylab="D")  #plot D and its qle approximation
legend("topright",c("D","Dqle"),lty=1:2,bty="n",col=1:2)

matplot(0:steps,x.out,type="l",xlab="generation", ylab="frequency",ylim=c(0,1))  #plot haplotype frequencies
legend("topright",c("x11","x12","x21","x22"),lty=1:4,bty="n",col=1:4)

matplot(0:steps,freqs,type="l",xlab="generation", ylab="frequency",ylim=c(0,1))  #plot allele frequencies at both loci
legend("topright",c("p","q"),lty=1:2,bty="n",col=1:2)