#Metropolis-Hastings on the space of step functions

f.log.lq <- function(t1,t2,t3,g,h1,h2,y,tobs)
{
  ## Purpose: Computation of log likelihood ratio  in the model
  ##    y[i]=mean(tobs[i]) + standard normal errors where mean is
  ##    one of two step functions which are equal outside (t1,t3].
  ##    Either the mean is =g on (t1,t3] or =h1 on (t1,t2] and =h2 on (t2,t3].
  ## -------------------------------------------------------------------------
  ## Arguments: as described above 
  ## -------------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 31 Jan 2002, 10:25
  n <- length(tobs)
  erg <- 0.0
  ind <- (t1<tobs & tobs<=t2)
  if (any(ind)) {
    erg <- erg + sum((y[ind]-g)^2)-sum((y[ind]-h1)^2)
  }
  ind <- (t2<tobs & tobs<=t3)
  if (any(ind)) {
    erg <- erg + sum((y[ind]-g)^2)-sum((y[ind]-h2)^2)
  }
  0.5*erg
}


f.jump.metr <- function(tt,g,y,tobs,lambda,tau)
{
  ## Purpose: One step in the Metropolis algorithm for jump functions. At each
  ##    iteration, one proposes to add or remove one jump. Proposal density
  ##    for heights equals the prior density.
  ## -------------------------------------------------------------------------
  ## Arguments: tt = jump times, g = heights
  ##            (modified by the function if move is accepted),
  ##            y, tobs = data; lambda,tau= parameters of the prior.
  ## -------------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 31 Jan 2002, 11:01
  k <- length(g)
  if (runif(1) < 1/k) {   # propose a move with one more jump 
    j <- sample((1:k),1,prob=(tt[-1]-tt[-(k+1)])) #select interval
    r <- tt[j] + runif(1)*(tt[j+1]-tt[j]) #propose new jump time
    h1 <- tau*rnorm(1) #propose new heights
    h2 <- tau*rnorm(1)
    acc <- lambda*k*exp(f.log.lq(tt[j],r,tt[j+1],g[j],h1,h2,y,tobs))/(k+1)
    if (runif(1) < acc) { #proposed move accepted 
      tt <- c(tt[1:j],r,tt[(j+1):(k+1)]) #add new jump to the function
      g[j] <- h1
      if (j==k) g <- c(g,h2)
      else g <- c(g[1:j],h2,g[(j+1):k])
    }
  }
  else { #propose a move with one jump less
    j <- sample((1:(k-1)),1) #select jump to be removed
    h <- tau*rnorm(1) #propose new height for the combined interval
    acc <- exp(-f.log.lq(tt[j],tt[j+1],tt[j+2],h,g[j],g[j+1],y,tobs))
    acc <- acc*k/(lambda*(k-1))
    if (runif(1) < acc) { #proposed move accepted
      tt <- tt[-(j+1)]  #remove the jump
      g[j] <- h
      g <- g[-(j+1)]
    }
  }
  list(tt=tt,g=g)
}

#setup: choosing the true mean function, generating data
n <- 100 
tobs <- ((1:n)-0.5)/n
y.t <- c(25,40,35) # lengths of constant subintervals
y.m <- c(2,-1,1)   # heights 
y <- rep(y.m,y.t) + rnorm(n)
plot(tobs,y,col=2)
lines(tobs,rep(y.m,y.t))
tt <- c(0,1) #Initialization: no jump
g <- mean(y)
lambda <- 3
tau <- 2
erg <- list(tt=tt,g=g)

#single steps of the method (for didactic purposes)
k <- length(g)
#plot the current step function
if (k > 1) plot(stepfun(tt[2:k],g),xlim=c(0,1),add=T) else abline(h=g) 
#add a jump
j <- sample((1:k),1,prob=(tt[-1]-tt[-(k+1)])) #select interval
r <- tt[j] + runif(1)*(tt[j+1]-tt[j]) #propose new jump time
h1 <- tau*rnorm(1) #propose new heights
h2 <- tau*rnorm(1)
tr <- c(tt[1:j],r,tt[(j+1):(k+1)]) #add new jump to the function
if (j==k) 
  h <- c(g,h2) else h <- c(g[1:j],h2,g[(j+1):k])
h[j] <- h1
plot(stepfun(tr[2:(k+1)],h),xlim=c(0,1),add=T) #plot proposed function
#compute acceptance probability
#lambda*k/(k+1) #if no data
acc <- lambda*k*exp(f.log.lq(tt[j],r,tt[j+1],g[j],h1,h2,y,tobs))/(k+1) #with data
acc
if (runif(1)<acc) {
  tt <- tr
  g <- h
  k <- k+1
}
#remove a jump
j <- sample((1:(k-1)),1) #select jump to be removed
h1 <- tau*rnorm(1) #propose new height for the combined interval
tr <- tt[-(j+1)]  #remove the jump
h <- g[-(j+1)]
h[j] <- h1
plot(stepfun(tr[2:(k-1)],h),add=T) #plot proposed function if k>2
#abline(h=h[1]) #plot proposed function if k=2
#compute acceptance probability
#k/(lambda*(k-1)) #if no data
acc <-  exp(-f.log.lq(tt[j],tt[j+1],tt[j+2],h1,g[j],g[j+1],y,tobs))*
  k/(lambda*(k-1))
if (runif(1) < acc) {
  tt <- tr
  g <- h
  k <- k-1
}






# Many steps 
m.est <- matrix(0,nrow=1000,ncol=n) #rows = estimated function at tobs,
                                   #for every 10th replicate.
k.est <- rep(0,10000) #number of jumps of estimated function
for (i in (1:10000)) {
  erg <- f.jump.metr(erg$tt,erg$g,y,tobs,lambda,tau)
  k.est[i] <- length(erg$g)-1
  if ((i %% 10) == 0){
    anz <- hist(tobs,breaks=erg$tt,plot=F)$counts
    m.est[i/10,] = rep(erg$g,anz)
  }
}
plot(1:10000,k.est,type="l")
sum(abs(k.est[-1]-k.est[-10000])>0.2) #number of accepted proposals
plot(tobs,y) #plotting the estimated functions and their quantiles
lines(tobs,rep(y.m,y.t))
for (i in (1:100)) lines(tobs,m.est[10*i,],col=2)
lines(tobs,apply(m.est,2,mean),col=3)
lines(tobs,apply(m.est,2,quantile,probs=0.1),col=2)
lines(tobs,apply(m.est,2,quantile,probs=0.5),col=3)
lines(tobs,apply(m.est,2,quantile,probs=0.9),col=4)

#Fuer Zuekost-Vortrag
ps.latex("zuekost-5-1.eps",height=8)
plot(tobs,y,xlab="x")
ps.end()

ps.latex("zuekost-5-2.eps",height=8)
plot(1:10000,k.est,type="l",xlab="iteration",ylab="k")
ps.end()

ps.latex("zuekost-5-3.eps",height=8)
plot(tobs,y,xlab="x") 
lines(tobs,rep(y.m,y.t))
for (i in (1:100)) lines(tobs,m.est[10*i,],col=2)
ps.end()

ps.latex("zuekost-5-4.eps",height=8)
plot(tobs,y,xlab="x")
lines(tobs,rep(y.m,y.t))
lines(tobs,apply(m.est,2,mean),col=3)
lines(tobs,apply(m.est,2,quantile,probs=0.1),col=2)
lines(tobs,apply(m.est,2,quantile,probs=0.5),col=3)
lines(tobs,apply(m.est,2,quantile,probs=0.9),col=4)
lines(tobs,rep(y.m,y.t))
ps.end()