# TEST OF A METHOD FOR DRAGGING FAST VARIABLES.
#
# Radford Neal, October 2004.
#
# These functions perform the tests in my technical report on "Taking
# Bigger Metropolis Steps by Dragging Fast Variables".  There are two
# sets of tests.  In the first, the state is (x,y).  In the second, the
# state is (x,y,z), with the marginal distribution for (x,y) being the
# same as for the first set of tests, and with the distribution of z
# given y being N(y,0.2^2).  The second set of tests uses the functions
# below ending with "2".


# The energy function E(x,y) defining the joint distribution.

Ejoint <- function (x,y)
{ 
  x^2 + 0.5 * (10*(1+x^2))^2 * (y[1]-sin(x))^2
}


# As above, but for the second set of tests.

Ejoint2 <- function (x,y,z)
{ 
  x^2 + 0.5 * (10*(1+x^2))^2 * (y-sin(x))^2 + (z-y)^2 / (2*0.2^2) 
}


# An energy function that produces the marginal distribution for x.

Emarg <- function (x)
{
  x^2 + log (1+x^2) 
}


# A Metropolis method that samples for (x,y) using proposals that change
# x and y simultaneously.  Arguments are as follows:
#
#   x0, y0    the initial values for the chain
#   sx, sy    the standard deviations for the normal proposal distributions
#             for the random walk metropolis updates
#   k         the number of groups of Metropolis updates to do
#   rep       the number of updates in each group (default 1)
#
# The value returned is a list with components x and y, containing the k+1
# values from the chain (including the initial value), and component r
# containing the rejection rate for the Metropolis updates.

met.joint <- function (x0,y0,sx,sy,k,rep=1)
{
  x <- c(x0,rep(0,k))
  y <- c(y0,rep(0,k))

  cx <- x0
  cy <- y0

  r <- 0

  for (i in 1:k)
  { 
    for (j in 1:rep)
    { 
      px <- rnorm(1,cx,sx)
      py <- rnorm(1,cy,sy)

      if (runif(1) < exp (Ejoint(cx,cy) - Ejoint(px,py)))
      { cx <- px
        cy <- py
      }
      else
      { r <- r+1
      }

    }

    x[i+1] <- cx
    y[i+1] <- cy

  }

  list (x=x, y=y, r=r/(k*rep))
}


# A Metropolis method that samples for (x,y,z) using proposals that change
# x, y, and z simultaneously.  Arguments are as follows:
#
#   x0, y0, z0  the initial values for the chain
#   sx, sy, sz  the standard deviations for the normal proposal distributions
#               for the random walk metropolis updates
#   k           the number of groups of Metropolis updates to do
#   rep         the number of updates in each group (default 1)
#
# The value returned is a list with components x, y, and z, containing the 
# k+1 values from the chain (including the initial value), and component r
# containing the rejection rate for the Metropolis updates.

met.joint2 <- function (x0,y0,z0,sx,sy,sz,k,rep=1)
{
  x <- c(x0,rep(0,k))
  y <- c(y0,rep(0,k))
  z <- c(z0,rep(0,k))

  cx <- x0
  cy <- y0
  cz <- z0

  r <- 0

  for (i in 1:k)
  { 
    for (j in 1:rep)
    { 
      px <- rnorm(1,cx,sx)
      py <- rnorm(1,cy,sy)
      pz <- rnorm(1,cz,sz)

      if (runif(1) < exp (Ejoint2(cx,cy,cz) - Ejoint2(px,py,pz)))
      { cx <- px
        cy <- py
        cz <- pz
      }
      else
      { r <- r+1
      }

    }

    x[i+1] <- cx
    y[i+1] <- cy
    z[i+1] <- cz

  }

  list (x=x, y=y, z=z, r=r/(k*rep))
}


# A Metropolis method that samples for (x,y) using proposals that change
# x alone followed by proposals that change y alone.  Arguments are as follows:
#
#   x0, y0    the initial values for the chain
#   sx, sy    the standard deviations for the normal proposal distributions
#             for the random walk metropolis updates
#   k         the number of groups of Metropolis updates to do
#   rep       the number of updates in each group (default 1)
#
# The value returned is a list with components x and y, containing the k+1
# values from the chain (including the initial value), and components rx and
# ry containing the rejection rates for the Metropolis updates of x and y.

met.single <- function (x0,y0,sx,sy,k,rep=1)
{
  x <- c(x0,rep(0,k))
  y <- c(y0,rep(0,k))

  cx <- x0
  cy <- y0

  rx <- 0
  ry <- 0

  for (i in 1:k)
  { 
    for (j in 1:rep)
    { 
      px <- rnorm(1,cx,sx)
      if (runif(1) < exp (Ejoint(cx,cy) - Ejoint(px,cy)))
      { cx <- px
      }
      else
      { rx <- rx+1
      }

      py <- rnorm(1,cy,sy)
      if (runif(1) < exp (Ejoint(cx,cy) - Ejoint(cx,py)))
      { cy <- py
      }
      else
      { ry <- ry+1
      }

    }

    x[i+1] <- cx
    y[i+1] <- cy

  }

  list (x=x, y=y, rx=rx/(k*rep), ry=ry/(k*rep))
}


# A Metropolis method that samples for (x,y) using proposals that change
# x, y, and z one at a time.  Arguments are as follows:
#
#   x0, y0, z0  the initial values for the chain
#   sx, sy, sz  the standard deviations for the normal proposal distributions
#               for the random walk metropolis updates
#   k           the number of groups of Metropolis updates to do
#   rep         the number of updates in each group (default 1)
#
# The value returned is a list with components x, y, and z, containing the k+1
# values from the chain (including the initial value), and components rx, ry,
# and rz containing rejection rates for the Metropolis updates of x, y, and z

met.single2 <- function (x0,y0,z0,sx,sy,sz,k,rep=1)
{
  x <- c(x0,rep(0,k))
  y <- c(y0,rep(0,k))
  z <- c(z0,rep(0,k))

  cx <- x0
  cy <- y0
  cz <- z0

  rx <- 0
  ry <- 0
  rz <- 0

  for (i in 1:k)
  { 
    for (j in 1:rep)
    { 
      px <- rnorm(1,cx,sx)
      if (runif(1) < exp (Ejoint2(cx,cy,cz) - Ejoint2(px,cy,cz)))
      { cx <- px
      }
      else
      { rx <- rx+1
      }

      py <- rnorm(1,cy,sy)
      if (runif(1) < exp (Ejoint2(cx,cy,cz) - Ejoint2(cx,py,cz)))
      { cy <- py
      }
      else
      { ry <- ry+1
      }

      pz <- rnorm(1,cz,sz)
      if (runif(1) < exp (Ejoint2(cx,cy,cz) - Ejoint2(cx,cy,pz)))
      { cz <- pz
      }
      else
      { rz <- rz+1
      }

    }

    x[i+1] <- cx
    y[i+1] <- cy
    z[i+1] <- cz

  }

  list (x=x, y=y, z=z, rx=rx/(k*rep), ry=ry/(k*rep), rz=rz/(k*rep))
}


# A Metropolis method that samples for x using the energy function that 
# produces its marginal distribution.  Arguments are as follows:
#
#   x0        the initial value for the chain
#   sx        the standard deviation for the normal proposal distributions
#             for the random walk metropolis updates
#   k         the number of groups of Metropolis updates to do
#   rep       the number of updates in each group (default 1)
#
# The value returned is a list with component x containing the k+1 values 
# from the chain (including the initial value) and component r containing 
# the rejection rate for the Metropolis updates.

met.marg <- function (x0,sx,k,rep=1)
{
  x <- c(x0,rep(0,k))

  cx <- x0

  r <- 0

  for (i in 1:k)
  { 
    for (j in 1:rep)
    { 
      px <- rnorm(1,cx,sx)
      if (runif(1) < exp (Emarg(cx) - Emarg(px)))
      { cx <- px
      }
      else
      { r <- r+1
      }

    }

    x[i+1] <- cx
  }

  list (x=x, r=r/(k*rep))
}


# A function to sample y values to go with x values.

yfill <- function (x) rnorm (length(x), sin(x), 0.1/(1+x^2))


# A function to sample z values to go with y values.

zfill <- function (x) rnorm (length(y), y, 0.2)


# The function for performing updates to x using transitions that drag along y.
# Arguments are as follows:
#
#   x0, y0    the initial values for the chain
#   sx, sy    the standard deviations for the normal proposal distributions
#             for the random walk metropolis updates
#   n         the number of interpolating distributions
#   k         the number of groups of updates to do
#   rep       the number of updates in each group (default 1)
#   repi      the number of inner Metropolis updates to do at each step
#
# The value returned is a list with components x and y, containing the k+1
# values from the chain (including the initial value), components r containing
# the rejection rate for the updates, and component ri containing the rejection
# rate for the inner Metropolis updates.

drag <- function (x0,y0,sx,sy,n,k,rep=1,repi=1)
{
  x <- c(x0,rep(0,k))
  y <- c(y0,rep(0,k))

  cx <- x0
  cy <- y0

  r <- 0
  ri <- 0

  for (i in 1:k)
  { 
    for (j in 1:rep)
    { 
      # Propose a change to x from its current value, cx.

      px <- rnorm(1,cx,sx)

      # Begin the series of dragging updates with the dragged state, yh,
      # set to the current value, cy.  Start accumulating the sum needed
      # to compute the acceptance probability.

      yh <- cy
      s <- Ejoint(cx,yh)-Ejoint(px,yh)

      # Perform n-1 groups of inner Metropolis updates for yh, keeping track 
      # of the sum needed to compute the acceptance probability.

      for (h in 1:(n-1))
      { 
        for (l in 1:repi)
        { 
          # Propose a new value for yh, and accept or reject in the usual
          # Metropolis fashion, using the interpolated energy function.

          pyh <- rnorm(1,yh,sy)
          a <- (h/n) * (Ejoint(px,yh) - Ejoint(px,pyh)) + 
               (1-h/n) * (Ejoint(cx,yh) - Ejoint(cx,pyh))
          if (runif(1) < exp(a))
          { yh <- pyh
          }
          else
          { ri <- ri+1
          }

        }

        s <- s + Ejoint(cx,yh)-Ejoint(px,yh)

      }

      # Decide whether or not to accept the outer update for x (along with
      # the final dragged y).

      py <- yh

      if (runif(1) < exp(s/n))
      { cx <- px
        cy <- py
      }
      else
      { r <- r+1
      }

    }

    x[i+1] <- cx
    y[i+1] <- cy

  }

  list (x=x, y=y, r=r/(k*rep), ri=ri/(k*rep*repi*(n-1)))
}


# The function for performing updates to x using transitions that drag along 
# y and z.  Arguments are as follows:
#
#   x0, y0, z0  the initial values for the chain
#   sx, sy, sz  the standard deviations for the normal proposal distributions
#               for the random walk metropolis updates
#   n           the number of interpolating distributions
#   k           the number of groups of updates to do
#   rep         the number of updates in each group (default 1)
#   repi        the number of inner Metropolis updates to do at each step
#
# The value returned is a list with components x, y, and z, containing the k+1
# values from the chain (including the initial value), components r containing
# the rejection rate for the updates, and component ri containing the rejection
# rate for the inner Metropolis updates (which change y and z simultaneously).

drag2 <- function (x0,y0,z0,sx,sy,sz,n,k,rep=1,repi=1)
{
  x <- c(x0,rep(0,k))
  y <- c(y0,rep(0,k))
  z <- c(z0,rep(0,k))

  cx <- x0
  cy <- y0
  cz <- z0

  r <- 0
  ri <- 0

  for (i in 1:k)
  { 
    for (j in 1:rep)
    { 
      # Propose a change to x from its current value, cx.

      px <- rnorm(1,cx,sx)

      # Begin the series of dragging updates with the dragged state, (yh,zh),
      # set to the current value, (cy,cz).  Start accumulating the sum needed
      # to compute the acceptance probability.

      yh <- cy
      zh <- cz
      s <- Ejoint2(cx,yh,zh)-Ejoint2(px,yh,zh)

      # Perform n-1 groups of inner Metropolis updates for yh and zh, keeping 
      # track of the sum needed to compute the acceptance probability.

      for (h in 1:(n-1))
      { 
        for (l in 1:repi)
        { 
          # Propose new values for yh and zh, and accept or reject in the usual
          # Metropolis fashion, using the interpolated energy function.

          pyh <- rnorm(1,yh,sy)
          pzh <- rnorm(1,zh,sz)
          a <- (h/n) * (Ejoint2(px,yh,zh) - Ejoint2(px,pyh,pzh)) + 
               (1-h/n) * (Ejoint2(cx,yh,zh) - Ejoint2(cx,pyh,pzh))
          if (runif(1) < exp(a))
          { yh <- pyh
            zh <- pzh
          }
          else
          { ri <- ri+1
          }

        }

        s <- s + Ejoint2(cx,yh,zh)-Ejoint2(px,yh,zh)

      }

      # Decide whether or not to accept the outer update for x (along with
      # the final dragged y and z).x

      py <- yh
      pz <- zh

      if (runif(1) < exp(s/n))
      { cx <- px
        cy <- py
        cz <- pz
      }
      else
      { r <- r+1
      }

    }

    x[i+1] <- cx
    y[i+1] <- cy
    z[i+1] <- cz

  }

  list (x=x, y=y, z=z, r=r/(k*rep), ri=ri/(k*rep*repi*(n-1)))
}


# Do the first set of tests, putting Postscript plots in test1.ps and test2.ps.

do.tests <- function()
{ 
  library(ts)

  iters <<- 10000

  postscript("test1.ps",paper="letter",height=6.5,width=6.4,horiz=F,pointsize=7)
  par (mfrow=c(2,3), mar=c(5,2,1,1)+0.1)

  set.seed(1)
  rj <<- met.joint (0, 0, 0.5, 0.5, iters)
  acf(rj$x,lag.max=30,main="",ylab="",xlab="Joint Metropolis",
      ci=0,ylim=c(-0.07,1))

  set.seed(1)
  rs <<- met.single (0, 0, 0.25, 0.25, iters)
  acf(rs$x,lag.max=30,main="",ylab="",xlab="Single-variable Metropolis",
      ci=0,ylim=c(-0.07,1))

  set.seed(1)
  rm <<- met.marg (0, 1.0, iters)
  acf(rm$x,lag.max=30,main="",ylab="",xlab="Marginal Metropolis",
      ci=0,ylim=c(-0.07,1))

  set.seed(1)
  rd.20 <<- drag (0, 0, 1.0, 0.2, 21, iters)
  acf(rd.20$x,lag.max=30,main="",ylab="",
      xlab="Dragging transitions, 20 steps", ci=0, ylim=c(-0.07,1))

  set.seed(1)
  rd.100 <<- drag (0, 0, 1.0, 0.2, 101, iters)
  acf(rd.100$x,lag.max=30,main="",ylab="",
      xlab="Dragging transitions, 100 steps", ci=0, ylim=c(-0.07,1))

  set.seed(1)
  rd.500 <<- drag (0, 0, 1.0, 0.2, 501, iters)
  acf(rd.500$x,lag.max=30,main="",ylab="",
      xlab="Dragging transitions, 500 steps", ci=0, ylim=c(-0.07,1))

  dev.off()

  postscript("test2.ps",paper="letter",height=4,width=6.4,horiz=F,pointsize=7)
  set.seed(1)
  plot(rm$x[seq(100,by=4,length=1000)],
       yfill(rm$x[seq(100,by=4,length=1000)]),
       xlim=c(-2,2),pch=20,xlab="x",ylab="y",
  )
  abline(h=0,v=0)
  dev.off()
}


# Do the second set of tests, putting a Postscript plot in test3.ps.  Should
# be done after the first set of tests, since it uses "rm" from that.

do.tests2 <- function()
{
  library(ts)

  iters <<- 10000

  postscript("test3.ps",paper="letter",height=6.5,width=6.4,horiz=F,pointsize=7)
  par (mfrow=c(2,3), mar=c(5,2,1,1)+0.1)

  set.seed(1)
  rj2 <<- met.joint2 (0, 0, 0, 0.3, 0.3, 0.3, iters)
  acf(rj2$x,lag.max=30,main="",ylab="",xlab="Joint Metropolis",
      ci=0,ylim=c(-0.07,1))

  set.seed(1)
  rs2 <<- met.single2 (0, 0, 0, 0.25, 0.25, 0.4, iters)
  acf(rs2$x,lag.max=30,main="",ylab="",xlab="Single-variable Metropolis",
      ci=0,ylim=c(-0.07,1))

  acf(rm$x,lag.max=30,main="",ylab="",xlab="Marginal Metropolis",
      ci=0,ylim=c(-0.07,1))

  set.seed(1)
  rd2.20 <<- drag2 (0, 0, 0, 1.0, 0.2, 0.2, 21, iters)
  acf(rd2.20$x,lag.max=30,main="",ylab="",
      xlab="Dragging transitions, 20 steps", ci=0, ylim=c(-0.07,1))

  set.seed(1)
  rd2.100 <<- drag2 (0, 0, 0, 1.0, 0.2, 0.2, 101, iters)
  acf(rd2.100$x,lag.max=30,main="",ylab="",
      xlab="Dragging transitions, 100 steps", ci=0, ylim=c(-0.07,1))

  set.seed(1)
  rd2.500 <<- drag2 (0, 0, 0, 1.0, 0.2, 0.2, 501, iters)
  acf(rd2.500$x,lag.max=30,main="",ylab="",
      xlab="Dragging transitions, 500 steps", ci=0, ylim=c(-0.07,1))

  dev.off()
}