# program to simulate random mating with mutation and selection at a diallelic locus. To plot results save output gc from rmsel and run one of the plotting functions given below.

rmsel <- function(ngen=10^3,nind=100,p=0.8,mu=0,s1=0,s2=0){
  gc <- matrix(0,ngen,3)	# initialize matrix gc to hold adult genotype counts
  for(i in 1:ngen){		#start evolution
    p <- p*(1-mu)+(1-p)*mu  # mutation changes allele fraction

#sample genotypes in the next generation
    geno <- sample(2:0,nind,repl=T,prob=c(p^2,2*p*(1-p)*(1+s1),(1-p)^2*(1+s2)))

# obtain the counts for the three genotypes

    gc[i,] <- hist(geno,br=seq(-0.5,2.5,1),plot=F)$c

    p <- mean(geno)/2	#update the value of p
  }
  gc		#return the gc matrix
}

# function to plot the frequencies of the three genotypes over time.  To eliminate the first say 100 generations from gc (when the simulation is far from equilibrium) you can replace gc with gc[-(1:100),] when calling these functions.

plotgen <- function(gc,ngen=10^3,nind=100,p=0.5,mu=0,s1=0,s2=0){
  matplot(gc,type="l",xlim=c(0,ngen*1.2),ylim=c(0,nind),lty=1,lwd=2,xlab="generation",ylab="genotype count",main=paste("nind=",nind, ", mu=",mu,", s1=",s1,", s2=",s2))
  legend(4.5*ngen/5,nind,leg=c("aa","Aa","AA"),lty=1,lwd=2,col=1:3,cex=1.8)
}

# function to plot allele frequency distribution over all generations, and return the mean count of one allele.

plotall <- function(gc,ngen=10^3,nind=100,p=0.5,mu=0,s1=0,s2=0){
	hist(gc %*% 2:0,n=50,xlab = "Allele count",main=paste("ngen=",ngen,", nind=",nind,", mu=",mu,", s1=",s1,", s2=",s2))
	sum(apply(t(gc[-(1:round(nrow(gc)/5)),])*(2:0),1,mean))
}
	
# function to plot the observed homozygosity and its expected value given the observed allele fractions
	  
plothom <- function(gc,ngen=10^3,nind=100,p=0.8,mu=0,s1=0,s2=0){
  homoz <- (gc[,1]+gc[,3])/nind
  plot(homoz,type="l",lwd=2,ylim=c(0,1),xlab="generation",ylab="homozygosity",main=paste("nind=",nind,", mu=",mu,", s1=",s1,", s2=",s2))
  p <- (2*gc[,1]+gc[,2])/(2*nind)
  lines(p^2+(1-p)^2,col=2,lwd=2)
  legend(0,0.26,leg=c("observed","HWE expected"),lty=1,lwd=2,col=1:3,cex=1.8)
}