# Various couplings of markov chains 

# Choosing a transition matrix
n <- 4
P <- 0.1*matrix(c(7,1,1,1,2,6,1,1,2,1,5,2,1,1,1,7),nrow=4,ncol=4,byrow=T)
P

# transition function using the quantile transform
C <- P
for (i in (1:n)) C[i,] <- cumsum(P[i,])
f.ueberg <- function(x,u,n,C)  min((1:n)[C[x,]>u]) 

# transition function with higher coupling probability
C2 <- apply(P,2,min) 
C2 <- matrix(C2,nrow=n,ncol=n,byrow=T)
C2 <- cbind(C2,P-C2)
C2
for (i in (1:nrow(C2))) C2[i,] <- cumsum(C2[i,])
f.ueberg.k <- function(x,u,n,C) {
  y <- min((1:(2*n))[C[x,]>u]) 
  if (y>n) y <- y-n
  y}

#20 steps with all initial values
nstep <- 20
par(mfrow=c(2,2))
x <- f.mc(n,nstep,f.ueberg,C) # independent transitions
plot(1:nstep,x[1,],type="b",ylim=c(1,n),xlab="time",ylab="state")
for (i in (2:n)) lines(1:nstep,x[i,]+0.02*i,type="b",col=i)
x <- f.mc.2(n,nstep,f.ueberg,C) # independent transitions until 1st meeting
plot(1:nstep,x[1,],type="b",ylim=c(1,n),xlab="time",ylab="state")
for (i in (2:n)) lines(1:nstep,x[i,]+0.02*i,type="b",col=i)
x <- f.mc.coupling(n,nstep,f.ueberg,C) # dependent transitions 1
plot(1:nstep,x[1,],type="b",ylim=c(1,n),xlab="time",ylab="state")
for (i in (2:n)) lines(1:nstep,x[i,]+0.02*i,type="b",col=i)
x <- f.mc.coupling(n,nstep,f.ueberg.k,C2) # dependent transitions 2
plot(1:nstep,x[1,],type="b",ylim=c(1,n),xlab="time",ylab="state")
for (i in (2:n)) lines(1:nstep,x[i,]+0.02*i,type="b",col=i)

# coupling from the past (cfp), visualisation
par(mfrow=c(2,2)) 
f.mc.cfp(n,plot=T,f.ueberg.k,C)

# generating 100 values of the stationary distribution using cfp
p0 <- eigen(t(P))$vectors[,1]
p0 <- p0/sum(p0) #stationary distribution
x <- rep(0,100) 
for (i in (1:100)) x[i] <- f.mc.cfp(n,plot=F,f.ueberg.k,C)
par(mfrow=c(1,1))
plot(1:4,table(x),type="h",ylim=c(0,max(table(x))))
points((1:4)+0.05,length(x)*p0,type="h",col=2)
chisq.test(table(x),p=p0,correct = F) # Testing if distribution is correct
p0 <- eigen(t(P))$vectors[,1]
p0 <- p0/sum(p0) #stationary distribution

# Example 2: Random walk on {1,2,...,n} with reflection at the boundary
                                        
n <- 5
P <- matrix(0,nrow=n,ncol=n) 
P[1,1] <- 0.5
P[n,n] <- 0.5
for (i in (2:n)){
  P[i,i-1] <- 0.5
  P[i-1,i] <- 0.5
}
P
f.irrfahrt <- function(x,u,n) {max(c(1,min(c(n,x+2*(u>0.5)-1))))}

# 30 steps with all initial values
par(mfrow=c(2,2))
nstep <- 30
x <- f.mc(n,nstep,f.irrfahrt) # independent transitions 
plot(1:nstep,x[1,],type="b",ylim=c(1,n),xlab="time",ylab="state")
for (i in (2:n)) lines(1:nstep,x[i,]+i*0.02,type="b",col=i)
x <- f.mc.2(n,nstep,f.irrfahrt) # independent transitions until coupling
plot(1:nstep,x[1,],type="b",ylim=c(1,n),xlab="time",ylab="state")
for (i in (2:n)) lines(1:nstep,x[i,]+i*0.02,type="b",col=i)
x <- f.mc.coupling(n,nstep,f.irrfahrt) # dependent transitions 
plot(1:nstep,x[1,],type="b",ylim=c(1,n),xlab="time",ylab="state")
for (i in (2:n)) lines(1:nstep,x[i,]+i*0.02,type="b",col=i)

# coupling from the past
par(mfrow=c(2,3)) 
f.mc.cfp(5,plot=T,f.irrfahrt)

#--------- Functions used ---------------------------------------

f.mc <- function(n,nstep,f.ueberg,...)
{
  ## Purpose: Generating a Markov chain for all initial values using
  ##          independent transitions
  ## -------------------------------------------------------------------------
  ## Arguments: f.ueberg=function to generate a transiton, nstep=number of steps.
  ## -------------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 11 Jan 2002, 11:40
  x <- matrix(1:n,nrow=n,ncol=nstep)
  for (j in (2:nstep)){
    for (i in (1:n)) x[i,j] <- f.ueberg(x[i,j-1],runif(1),n,...)
  }
  x
}

f.mc.2 <- function(n,nstep,f.ueberg,...)
{
  ## Purpose: Generating a Markov chain for all initial values using
  ##  independent transitions until first meeting, and coupling afterwards
  ## -------------------------------------------------------------------------
  ## Arguments: f.ueberg=function to generate a transiton, nstep=number of steps.
  ## -------------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 11 Jan 2002, 11:40
  x <- matrix(1:n,nrow=n,ncol=nstep)
  y <- 1:n
  for (j in (2:nstep)) {
    for (i in (1:n)) y[i] <- f.ueberg(i,runif(1),n,...)
    # value of the chain at time j+1 if value at time j equal i
    x[,j] <- y[x[,j-1]] # choose the next value according to the current value
  }
  x
}

f.mc.coupling <- function(n,nstep,f.ueberg,...)
{
  ## Purpose: Generating a Markov chain for all initial values using
  ##  the same random numbers 
  ## -------------------------------------------------------------------------
  ## Arguments: f.ueberg=function to generate a transiton, nstep=number of steps
  ## -------------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date: 11 Jan 2002, 11:40
  x <- matrix(1:n,nrow=n,ncol=nstep)
  for (j in (2:nstep)){
    u <- runif(1)
    for (i in (1:n)) x[i,j] <- f.ueberg(x[i,j-1],u,n,...)
  }
  x
}

f.mc.cfp <- function(n,plot=F,f.ueberg,...)
{
  ## Purpose: Coupling from the past of Markov chains
  ## -------------------------------------------------------------------------
  ## Arguments: n=number of states, f.ueberg=transition function
  ## -------------------------------------------------------------------------
  ## Author: Hans-Ruedi Kuensch, Date:  4 Feb 2002, 18:42
  x <- matrix(1:n,nrow=n,ncol=2) # 1 transition for intialisation
  u <- runif(1)
  for (i in (1:n)) x[i,2] <- f.ueberg(x[i,1],u,n,...)
  nstep <- 2
  if (plot) { #Plotting
    plot((1-nstep):0,x[1,],type="b",xlim=c(-15,0),ylim=c(1,n),
       xlab="time",ylab="state") 
    for (i in (2:n)) lines((1-nstep):0,x[i,],type="b",col=i)
  }
  while (any(x[-1,nstep] !=x[1,nstep])){ # Testing if all couplings occured 
    y <- matrix(1:n,nrow=n,ncol=nstep+1) # additional nstep transitions
    for (j in (1:nstep)){
      u <- runif(1)
      for (i in (1:n)) y[i,j+1] <- f.ueberg(y[i,j],u,n,...)
    }
    z <- x   # Composition with previously generated paths 
    x <- matrix(0,nrow=n,ncol=2*nstep)
    for (i in (1:n)) x[i,] <- c(y[i,],z[y[i,nstep+1],-1])
    nstep=nstep*2
    if (plot) { #Plotting
      plot((1-nstep):0,x[1,],type="b",xlim=c(1-max(16,nstep),0),ylim=c(1,n),
           xlab="time",ylab="state") 
      for (i in (2:n)) lines((1-nstep):0,x[i,],type="b",col=i)
    }
    }
x[1,nstep]}

# --------------------------------------------------------------------
# Discrete version of random walk Metropolis-Hastings

#defining the target distribution (mixture of 2 discretized Beta-distributions)
a1 <- 0.3
n1 <- 500
a2 <- 0.5
n2 <- 500

k <- 0.01*(1:99)
p1 <- k^(a1*n1)*(1-k)^((1-a1)*n1)
p2 <- k^(a2*n2)*(1-k)^((1-a2)*n2)
# work with logs to avoid underflow in target probabilities
lp1 <- a1*n1*log(k) + (1-a1)*n1*log(1-k) +log(0.4) - log(sum(p1))
lp2 <- a2*n2*log(k) + (1-a2)*n2*log(1-k) +log(0.6) - log(sum(p2))
lmn <- pmin(lp1,lp2)
lmx <- pmax(lp1,lp2)
lp0 <- lmx + log(1+exp(lmn-lmx))
p0 <- exp(lp0) # same as p0 <- 0.4*p1/sum(p1) + 0.6*p2/sum(p2)
plot(k,p0,type="h")

#defining the transition matrix
m <- 15 # range of the random walk
q <- outer(1:99,1:99,FUN=function(x,y) {1<= abs(x-y) & abs(x-y) <= m})
q <- q + outer(1:99,1:99,FUN=function(x,y) {abs(x-y) >=99-m})
q <- q/(2*m) # proposal
a <- outer(lp0,lp0, FUN="-")
a <- exp(pmin(0,t(a))) #acceptance 
p.t <- a*q
rej <- 1-apply(p.t,1,sum) # rejection probabilities
sum(rej*p0) # rejection probability in stationary state
p.t <- p.t + diag(rej) # adjust diagonal

sum(abs(p0-as.vector(t(p.t)%*%p0))) #check stationarity

ev <- eigen(p.t)
plot(abs(ev$values),type="h")
ev$value[2] #geometric convergence rate to the stationary distribution

mimage(q)
mimage(a*q)
p.p <- p.t
mimage(p.p)
p.p <- p.p%*%p.p #P^2, P^4, P^8 etc.


mimage <- function(mat, ncol = 40, zlim=range(mat))
{
  ## Purpose: Color representation of a matrix 
  ## ----------------------------------------------------------------------
  ## Arguments:
  ## ----------------------------------------------------------------------
  ## Author: Hans-Rudolf Kuensch, Date:  3 Dec 2007, 16:47

  ## Interpolating a 'sequential' ColorBrewer palette
  YlOrBr <- c("#FFFFD4", "#FED98E", "#FE9929", "#D95F0E", "#993404")

  stopifnot(length(d <- dim(mat)) == 2)
  image(x=1:d[2], y=1:d[1], z=t(mat)[,d[1]:1], axes = FALSE,
        xlab = "j", ylab  = "i",zlim=zlim,
        col = colorRampPalette(YlOrBr, space = "Lab")(ncol))
  axis(1)
  axis(2, at= 1:d[2], labels= as.character(d[2]:1))
}