# by Cecile Ane, December 2012.
#
# Purpose: standalone R code to visualize the prior distribution on the 
# number of distinct trees from the Dirichlet process. Replicates the graph
# from this website: http://www.stat.wisc.edu/~ane/bucky/prior.html
# To use this R code (especially when the website is down...),
# download this file and do this:
#
# 1. open an R session
# 2. change the working directory to where this file is located,
# 3. source the file using the following command
#    source("prior_standalone.r")
#    (or adapt with this file's current name)
# 4. get the graph using the following command:
#    priorPlot(ntax=8,ngenes=106,alpha=2,las=1,lwd=3)
#    where the number of taxa (ntax), number of genes (ngenes)
#    and alpha values can be adapted.
#    Other graphical parameters can be added.
#    las=1 is to have horizontal labels on y axis,
#    lwd=3 is to have a line width of 3.
#
# Note that the graph can take a long time if the number of genes is large.
# Hence, first try a fast example (like the one above with ngenes=106) 
# to check that the function is working in your environment.

logUB = function(ntax){
# log number of unrooted binary trees on ntax taxa
 if (ntax<4) {return(0)} else {
  sum(log(2*(4:ntax)-5))
 }
}
getlogA = function(n,x){
# calculates the log of x(x+1)...(x+n-1)
 sum(log(x+ 0:(n-1)))
}
logProbAdd = function(x,y){
# computes log( exp(x) + exp(y) ) carefully
# By computing max + log(1 + exp(min-max))
 if (x<y){ y+log(1+exp(x-y))} else {x+log(1+exp(y-x))}
}
logProbSubtract = function(x,y){
# computes log( exp(x) - exp(y) ). assumes x>=y
 x+log(1-exp(y-x))
}

logStirling1 = function(n,k){
# log of |stirling number| of the first kind: x_[i,j] = log( A[i,j] )
# returns n by k matrix.
 if (k>n){ k=n }
 x = matrix(NA,n,k) 
 # actually, x[i,] has non-zero elements for j<=i (or k if smaller)
 diag(x) = 0; # this is a partial diagonal if k different from n
 for(i in 2:n) {
  li = log(i-1)
  x[i,1] = li + x[i-1,1]
  if (i>2){ for (j in 2:min(i-1,k)){
   x[i,j] = logProbAdd(li + x[i-1,j], x[i-1,j-1])
  }}
 }
 return(x)
}
# exp(logStirling1(4,2))

logStirling2 = function(n,k){
# log of |stirling numbers| of the second kind. They are found
# by the recursive formula: S(n,k) = S(n-1,k-1) + k S(n-1,k).
# output: n by k matrix.
 x = matrix(NA,n,k) 
 # actually, x[i,] has non-zero elements for j<=i (or k if smaller)
 diag(x) = 0;
 x[,1] = 0;
 for(i in 2:n) {
  if (i>2){ for (j in 2:min(i-1,k)){
   x[i,j] = logProbAdd(log(j) + x[i-1,j],x[i-1,j-1])
  }}
 }
 return(x)
}
# exp(logStirling2(4,2))

dpprior.base = function(n,logalpha,logT,maxK,s1,s2){
# Dirichlet prior probability of # groups
# n: number genes, alpha: dirichlet parameter, T: total number trees
# s1,s2: log |stirling #| of first,second kind.
 p = rep(NA,min(n,maxK)) # p[i]=prob of i groups
 lA = getlogA(n,exp(logalpha))
 logSum = s1[1] + s2[1,1] + (logalpha - logT);
 for(i in 2:n){
  logSum = logProbAdd(logSum, s1[i] + s2[i,1] + i*(logalpha - logT))
 }
 p[1] = exp(logT - lA + logSum);
 c1 = logT
 for(k in 2:min(n,maxK)){
  c1 = c1 + logProbSubtract(logT,log(k-1))
  logSum = s1[k] + s2[k,k] + k*(logalpha - logT);
  if (k<n){ for(i in (k+1):n){
   logSum = logProbAdd(logSum, s1[i] + s2[i,k] + i*(logalpha - logT))
  }}
  p[k] = exp(c1 - lA + logSum);
 }
 return(p)
}

dpprior = function(n,alpha,ntax){
 logT = logUB(ntax)
 # default choice of maxK: based on the expected number of jokers: 
 # formula in (1974) paper on Dirichlet process.
 meanK = sum(1/(alpha+0:(n-1))) * alpha
 sdK = sqrt(sum(0:(n-1) /(alpha+0:(n-1))) * alpha)
 maxK = round(meanK + 3.5*sdK)
 maxK = max(maxK,1) # we don't want maxK = 0
 maxK = min(maxK,n) # we don't want maxK > number of genes
 maxK = round(min(maxK,exp(logT))) # we don't want maxK > # distinct trees
 #cat("mean=",meanK,"sd=",sdK,"maxK=",maxK,"\n")
 if (ntax<4){ p=c(1,rep(0,maxK-1)) } else {
  s1 = logStirling1(n,n)[n,]
  s2 = logStirling2(n,maxK)
  #print(exp(s1)); print(exp(s2))
  p = dpprior.base(n,log(alpha),logT,maxK,s1,s2)
 }
 return(p)
}

#myp = dpprior(106,0.1,8) # dpprior(10,1,6)
#round(myp,3)  # returns the prior prob. of 1,2,..,106 groups
#sum(myp*1:length(myp), na.rm=T)# returns the expected number of groups.

priorPlot = function(ntax,ngenes,alpha,...) {
 par(mar=c(4.1,4.1,4.1,1), mgp=c(2.5,.5,0), tck=-.02)
 p = dpprior(ngenes,alpha,ntax)
 x = 1:length(p)
 plot(x,p,type="h",xlab="Number of Distinct Trees",
      ylab="probability",main="",col="blue",xlim=c(0.5,max(x)),ylim=c(0,max(p)),
      cex.lab=1.2,cex.axis=1.2,...)
 title("Prior Distribution on Number of Distinct Trees",font.main=1,line=3,cex.main=1.5)
 mtext(paste("#genes =",ngenes,", #taxa =",ntax),line=1.8,cex=1.2)
 mtext(bquote(alpha == .(alpha)),line=1,cex=1.2)
 abline(h=0)
}

# Example of use: 
# priorPlot(ntax=8,ngenes=106,alpha=2,las=1,lwd=3)