#Script for Exercise set #3, problem 1
#by R. Gomulkiewicz 14 Oct 2013

#Set up in-common parameter values and define functions

#alpha = effect of A locus on phenotype & beta = effect of B locus genotype on phenotype
alpha <- 1
beta <- 2

#Phenotypes
#store in a 4 x 4 matrix phi; each row and column position corresponds to gametes in diploid genotype
#gamete subscripts: 1 <-> A1B1, 2 <-> A1B2, 3 <-> A2B1, 4 <-> A2B2
phi <- matrix(c(alpha+beta, alpha, beta, 0,alpha, alpha-beta, 0, -beta, beta, 0, -alpha+beta, -alpha, 0, -beta, -alpha, -alpha-beta), nrow=4, ncol=4)

#Strength of optimizing selection 
s <- 0.01

#Generate 4x4 matrix with fitnesses for each diploid genotype based on optimizing selection
w <- 1-s*(phi^2) #"^2" squares each component of phi

#recombination rate
r <- 0.1

#Functions of gamete frequencies
pA1 <- function(x){x[1]+x[2]}
pB1 <- function(x){x[1]+x[3]}
D <- function(x){x[1]*x[4] - x[2]*x[3]}
#use matrix-vector multiplication to compute the mean fitness, mean phenotype, and phenotypic/genotypic variance
wbar <- function(x){x %*% w %*% x} 
phibar <- function(x){x %*% phi %*% x}
phivar <- function(x){(x %*% (phi^2) %*% x) - (phibar(x))^2 }

#Recursions for gamete frequencies
x.next <- function(x){wstar<-w %*% x; c(x[1]*wstar[1] - r*D(x)*w[1,4], x[2]*wstar[2] + r*D(x)*w[1,4], x[3]*wstar[3] + r*D(x)*w[1,4], x[4]*wstar[4] - r*D(x)*w[1,4])/wbar(x)}

#Function to iterate and plot genetic and phenotypic dynamics for T generations with initial gamete frequency vector x0
GenoPhenoDyn <- function(T, x0){
x.out <- matrix(0, nrow = T+1, ncol=4)
p.out <- matrix(0,nrow=T+1, ncol=2)
phibar.out <- matrix(0,nrow=T+1, ncol=1)
phivar.out <- matrix(0,nrow=T+1, ncol=1)
D.out <- matrix(0,nrow=T+1, ncol=1)
wbar.out <- matrix(0,nrow=T+1,ncol=1)

#Initializations
x.out[1,] <- x0
p.out[1,] <- c(pA1(x.out[1,]), pB1(x.out[1,]))
D.out[1,] <- D(x.out[1,])
phibar.out[1] <- phibar(x.out[1,])
phivar.out[1] <- phivar(x.out[1,])
wbar.out[1] <- wbar(x.out[1,])

#Iterate dynamics
for(i in 1:T){
	x.out[i+1,] <- x.next(x.out[i,])
	p.out[i+1,] <- c(pA1(x.out[i+1,]), pB1(x.out[i+1,]))
	D.out[i+1] <- D(x.out[i+1,])
	phibar.out[i+1] <- phibar(x.out[i+1,])
	phivar.out[i+1] <- phivar(x.out[i+1,])
	wbar.out[i+1] <- wbar(x.out[i+1,])
	}

#Plot genetic dynamics
quartz()
par(mfrow=c(3,1))
matplot(0:T,x.out,type="l",xlab="generation", ylab="Gamete Freq")  #plot haplotype frequencies
legend("topright",c(expression(x[11]),expression(x[12]),expression(x[21]),expression(x[22])),lty=1:4,bty="n",col=1:4)

matplot(0:T,p.out,type="l",xlab="generation", ylab="Allele Freq")  #plot allele frequencies at both loci
legend("topright",c("p","q"),lty=1:2,bty="n",col=1:2)

plot(0:T, D.out, type="l",xlab="generation", ylab="D") #plot disequilibria

#Plot phenotypic dynamics
quartz()
par(mfrow=c(2,1))
plot(0:T, phibar.out, type="l",xlab="generation", ylab=expression(bar(phi))) #plot disequilibria
plot(0:T, phivar.out, type="l",xlab="generation", ylab=expression(sigma[phi]^2)) #plot disequilibria
}

###################################################
#SCENARIO 1: equal haplotype frequencies
x0 <- rep(0.25,4)
GenoPhenoDyn(400,x0)

###################################################
#SCENARIO 2
x0 <- c(0.2, 0.2, 0.3, 0.3)
GenoPhenoDyn(400,x0)