# SHORT-MET-ABBREV.R - Abbreviated version of short-met.r.


# DO ONE METROPOLIS UPDATE.
#
# Arguments:
#
#    initial.x    The initial state (a vector)
#    lpr          Function returning the log probability of a state, plus an
#                 arbitrary constant
#    pf           Function returning the random offset (a vector) for a proposal
#    w            The stepsize for proposals, multiplies the offset
#    initial.lpr  The value of lpr(initial.x).
#
# The value returned is a list containing the following elements:
#
#    next.x       The new state
#    next.lpr     The value of lpr(next.x)
#    rejected     TRUE if the proposal was rejected

metropolis.update <- function (initial.x, lpr, pf, w, initial.lpr)
{
  # Propose a candidate state, and evalute its log probability.

  proposed.x <- initial.x + w*pf()
  proposed.lpr <- lpr(proposed.x)

  # Decide whether to accept or reject the proposed state as the new state.

  if (runif(1)<exp(proposed.lpr-initial.lpr)) # accept
  { next.x <- proposed.x
    next.lpr <- proposed.lpr
    rejected <- FALSE
  }
  else # reject
  { next.x <- initial.x
    next.lpr <- initial.lpr
    rejected <- TRUE
  }

  # Return the new state, its log probability, and whether a rejection occurred.

  list (next.x=next.x, next.lpr=next.lpr, rejected=rejected)
}


# THE SHORT-CUT METROPOLIS METHOD.  Simulates a short-cut Metropolis update
# consisting of M*L Metropolis updates.
#
# Arguments:
#
#    initial.x    The initial state (a vector of length n)
#    lpr          Function returning the log probability of a state, plus an
#                 arbitrary constant
#    pf           Function returning the random offset (a vector) for a proposal
#    w            The stepsize for proposals, multiplies the offset
#    L            Number of updates in each group
#    M            Number of groups to simulate
#    min.rej      Minimum number of rejections for a good group of L updates 
#    max.rej      Maximum number of rejections for a good group of L updates
#
# The value returned is a list containing the following elements:
#
#    states1      A M*L+1 by n matrix whose rows contain the initial state 
#                 and the states after each Metropolis update, including
#                 updates in groups after which a reversal occurred
#    states2      A M+1 by n matrix whose rows contain the initial state and
#                 the states after each group of L Metropolis updates; the
#                 last row contains the final state from the whole sequence

short.cut.metropolis <- function (initial.x, lpr, pf, w, L, M, 
                                  min.rej=0, max.rej=L-1)
{
  # The function below puts together the list of results that we return.

  results <- function () list (states1=states1, states2=states2)

  # The function below performs the L Metropolis updates in a group, storing the
  # results in states1.  The last.x and last.lpr variables should contain the
  # previous state and its log probability; they are updated by this function.
  # The variable k indexes where in states1 the new states should be stored.  
  # It is incremented in this function.  The n.rejected variable is set to the 
  # number of the L updates that were rejections.  

  do.group <- function ()
  { n.rejected <<- 0
    for (l in 1:L)
    { update <- metropolis.update (last.x, lpr, pf, w, last.lpr)
      states1[k+1,] <<- last.x <<- update$next.x
      last.lpr <<- update$next.lpr
      k <<- k + 1
      n.rejected <<- n.rejected + update$rejected
    }
  }

  # Allocate space for the results.

  K <- M*L
  states1 <- matrix(NA,K+1,length(initial.x))
  states2 <- matrix(NA,M+1,length(initial.x))

  # Store the initial state in the first row of states1 and of states2, and 
  # evaluate its log probability.

  states1[1,] <- initial.x
  states2[1,] <- initial.x
  initial.lpr <- lpr(initial.x)
 
  # Do groups of Metropolis updates starting from the initial state, until we've
  # done many as were asked for, or until the number of rejections in a group 
  # is outside the limits.  States after each Metropolis update are saved
  # in states1.  The final state for each group is saved in states2.  Note 
  # that for a group with the number of rejections outside the limits, the 
  # state saved in states2 is not state the last state in states1, but rather
  # state from the start of the group.  The indexes in states1 of the start 
  # (upper0) and end (upper1) of these groups in are recorded for later use 
  # in copying states.

  last.x <- initial.x
  last.lpr <- initial.lpr
  k <- 1

  upper0 <- k
  while (k<=K)
  { states2[2+(k-1)/L,] <- last.x
    do.group()
    if (n.rejected<min.rej || n.rejected>max.rej)
    { upper1 <- k
      break
    }
    else
    { states2[1+(k-1)/L,] <- states1[k,]
    }
  }

  if (k>K) return (results())  # Return if we've done all M groups.

  # Copy already-computed states as "new" states as we move backwards through
  # the groups previously simulated.  Don't copy the last group causing the
  # reversal, for which the number of rejections was outside the limits.

  j <- upper1 - L - 1
  while (k<=K && j>=upper0)
  { states1[(k+1):(k+L),] <- states1[j:(j-L+1),]
    states2[1+(k+L-1)/L,] <- states2[1+(j-L)/L,]
    k <- k + L 
    j <- j - L
  }

  if (k>K) return (results())  # Return if we've done all M groups.

  # Restore the initial state, then do more groups of Metropolis updates, 
  # until we've done as many as were asked for, or the number of rejections in
  # a group is outside the limits.  Record the indexes of the start (lower0) 
  # and end (lower1) of these groups in states1. 

  last.x <- initial.x
  last.lpr <- initial.lpr

  lower0 <- k
  while (k<=K)
  { states2[2+(k-1)/L,] <- last.x
    do.group()
    if (n.rejected<min.rej || n.rejected>max.rej)
    { lower1 <- k
      break
    }
    else
    { states2[1+(k-1)/L,] <- states1[k,]
    }
  }

  if (k>K) return (results())  # Return if we've done all M groups.

  # Copy already-computed states as "new" states, going back and forth over
  # the "lower" groups and the "upper" groups.

  repeat
  {
    # Copy the lower states backwards, excluding the group causing the reversal.

    j <- lower1 - L - 1
    while (k<=K && j>=lower0)
    { states1[(k+1):(k+L),] <- states1[j:(j-L+1),]
      states2[1+(k+L-1)/L,] <- states2[1+(j-L)/L,]
      k <- k + L
      j <- j - L
    }

    if (k>K) return (results())  # Return if we've done all M groups.

    # Copy the upper states forwards, including the group causing the reversal.

    j <- upper0+1
    while (k<=K && j<=upper1)
    { states1[(k+1):(k+L),] <- states1[j:(j+L-1),]
      states2[1+(k+L-1)/L,] <- states2[1+(j+L-2)/L,]
      k <- k + L
      j <- j + L
    }

    if (k>K) return (results())  # Return if we've done all M groups.

    # Copy the upper states backwards, excluding the group causing the reversal.

    j <- upper1 - L - 1
    while (k<=K && j>=upper0)
    { states1[(k+1):(k+L),] <- states1[j:(j-L+1),]
      states2[1+(k+L-1)/L,] <- states2[1+(j-L)/L,]
      k <- k + L
      j <- j - L
    }

    if (k>K) return (results())  # Return if we've done all M groups.
   
    # Copy the lower states forwards, including the group causing the reversal.

    j <- lower0 + 1
    while (k<=K && j<=lower1)
    { states1[(k+1):(k+L),] <- states1[j:(j+L-1),]
      states2[1+(k+L-1)/L,] <- states2[1+(j+L-2)/L,]
      k <- k + L
      j <- j + L
    }

    if (k>K) return (results())  # Return if we've done all M groups.
  }
}