# hw09.R
# Bret Larget
# Partial assignment for Statistics 310, Spring 2010

# Data
s=c(753.0,1396.9,1528.2,1646.9,717.6,446.5,848.0,1222.0,748.5,610.1)

# prior distribution
#
# This function is written to accept vectors instead of single values
# for alpha and or lambda.  If one is longer than the other, the shorter
# vector is repeated to match the length of the longer vector.
# The purpose of "vectorizing" the function in this way is so that
# it is easy to compute the value of the prior distribution for many values in one call
# which allows simpler code for graphing the function, for example.
#
# This prior density arises from:
# 1/lambda | alpha ~ Gamma(2,alpha/400)
# alpha/4 ~ Beta(2,2)
# With this prior density and the likelihood X | alpha,lambda ~ Gamma(alpha,lambda),
# we have:
#   E(X) = E( E(X | alpha,lambda) )
#        = E( alpha/lambda )
#        = E( E(alpha/lambda | alpha) )
#        = E( alpha * E(1/lambda | alpha) )
#        = E( alpha * (800/alpha) )
#        = 800

prior = function(alpha,lambda) {
  n.alpha = length(alpha)
  n.lambda = length(lambda)
  # The shorter argument is repeated until it matches length with the longer argument
  if(n.alpha < n.lambda)
    alpha = rep(alpha,length.out=n.lambda)
  if(n.lambda < n.alpha)
    lambda = rep(lambda,length.out=n.alpha)
  # Compute the answer ignoring the range
  out = ( 3*alpha^3*(4-alpha) / (32*400^2*lambda^3) ) * exp(-alpha/(400*lambda))
  # Now set to zero arguments out of range
  zero = (alpha<0) | (alpha>4) | (lambda<0)
  if(any(zero))
    out[zero] = 0
  return( out )
}

# Likelihood
#
# The likelihood model is X[i] | alpha,lambda ~ Gamma(alpha,lambda)
# This function is also written to be calculable for vector arguments
# for alpha and lambda.

likelihood = function(alpha,lambda,s) {
  # Lines as in the previous function to make the arguments the same length.
  n.alpha = length(alpha)
  n.lambda = length(lambda)
  if(n.alpha < n.lambda)
    alpha = rep(alpha,length.out=n.lambda)
  if(n.lambda < n.alpha)
    lambda = rep(lambda,length.out=n.alpha)
  # Create space for the answer
  out = rep(NA,length(alpha))
  for(i in 1:length(alpha))
    out[i] = prod(dgamma(s,alpha[i],lambda[i]))
  return(out)
}

# Unnormalized posterior distribution is prior()*likelihood()

h = function(alpha,lambda,s) { return( prior(alpha,lambda) * likelihood(alpha,lambda,s) ) }

# MCMC
#
# propose() can be changed for other proposals, such as normal.
#
propose = function(x,w=1) {
  return( x + runif(1,-w,w) )
}

# mcmc()
#
# Description of arguments:
#   n.mcmc is the desired sample size
#   alpha.0 is the initial value of alpha
#   lambda.0 is the inital value of lambda
#   s is a vector with the data
#   w.alpha is a tuning parameter for the proposal for a new alpha
#   w.lambda is a tuning parameter for a new lambda
#
# Output is a matrix with n.mcmc rows and 4 columns.
#   The first column is alpha
#   The second column is lambda
#   The third column is h
#   The fourth column is the acceptance probability.
#
# Example:
# # run mcmc()
# > out = mcmc(10000,alpha.0=2,lambda.0=0.001,s,w.alpha=0.5,w.lambda=0.0005)
# 
# #  find the posterior means by applying mean() to each column (dimension 2)
# > apply(out,2,mean)
#
# #  find the posterior sds by applying sd() to each column (dimension 2)
# #  (Maybe use this information for a longer run with better starting values and tuning parameters
# > apply(out,2,mean)
#
# #  find 95% confidence intervals by finding quantiles from first two columns
# > apply(out[,1:2],2,quantile,probs=c(0.025,0.975))

mcmc = function(n.mcmc=10000,alpha.0,lambda.0,s,w.alpha=0.4,w.lambda=0.0004) {
  out = matrix(NA,n.mcmc,4)
  alpha.curr = alpha.0
  lambda.curr = lambda.0
  h.curr = h(alpha.curr,lambda.curr,s)
  for(i in 1:n.mcmc) {
    alpha.prop = propose(alpha.curr,w.alpha)
    lambda.prop = propose(lambda.curr,w.lambda)
    if( (alpha.prop > 0) && (alpha.prop < 4) && (lambda.prop > 0) ) {
      h.prop = h(alpha.prop,lambda.prop,s)
      accept.prob = min(c(1,h.prop/h.curr))
      if(runif(1) < accept.prob) { # accept
        alpha.curr = alpha.prop
        lambda.curr = lambda.prop
        h.curr = h.prop
      }
    }
    else
      accept.prob=0
    out[i,] = c(alpha.curr,lambda.curr,h.curr,accept.prob)
  }
  # Add names to the columns, but not the rows
  dimnames(out) = list(NULL,c("alpha","lambda","h","accept.prob"))
  return( out )
}