inst/gganimate-coupledcpf.R

In pierrejacob/CoupledCPF: Coupling of Conditional Particle Filters

#### This scripts creates an animation
#### showing two trajectories generated by CCPF
#### from iteration 1 until they meet
# remove all objects from R environment
rm(list = ls())
# load packages
library(CoupledCPF)
library(dplyr)
library(reshape2)
# 
set.seed(17)
module_tree <<- Module("module_tree", PACKAGE = "CoupledCPF")
TreeClass <<- module_tree$Tree
#
# We are going to work with the nonlinear growth model of Gordon, Salmond and Smith
# first we import it
gss <- get_gss()
# gss contains the model distributions, i.e. initialization, transition and measurement
names(gss)
# generate some data from the model
dimension <- 1
theta <- c()
datalength <- 50
x <- matrix(0, nrow = datalength + 1, ncol = 1)
observations <- matrix(0, nrow = datalength, ncol = 1)
x[1,] <- gss$rinit(1, theta, gss$rinit_rand(1, theta))
# equivalent to
#x[1,] <- rnorm(1, 0, sd = sqrt(2))
for (i in 1:datalength){
  x[i+1,] <- gss$rtransition(x[i,], theta, i, gss$rtransition_rand(1, theta))
  # equivalent to 
  # x[i+1,] <- 0.5 * x[i,] + 25 * x[i,] / (1 + x[i,]^2) + 8 * cos(1.2*(i - 1)) + sqrt(10) * rnorm(1)
  observations[i,] <- (x[i+1,]^2) / 20 + rnorm(1)
}
# plot the latent process
plot(0:datalength, x, type = "l")
# plot the observations
plot(1:datalength, observations, type = "l")
#

# now let us run coupled chains on this example
# number of particles
N <- 512
# with ancestor sampling
with_as <- TRUE
# coupled resampling scheme
coupled_resampling <- CR_indexmatching
# initial distribution of the chains: draw a trajectory from a particle filter
rinit <- function(){
  return(CPF(N, gss, theta, observations, ref_trajectory = NULL, with_as = FALSE))
}
# single kernel of CPF
single_kernel_RB <- function(chain_state, h = function(x) x){
  return(CPF_RB(N, gss, theta, observations, ref_trajectory = chain_state, with_as = with_as, h = h))
}
# coupled kernel of CPF
coupled_kernel_RB <- function(chain_state1, chain_state2, h = function(x) x){
  return(CPF_coupled_RB(N, gss, theta, observations, ref_trajectory1 = chain_state1, ref_trajectory2 = chain_state2,
                        coupled_resampling = coupled_resampling, with_as = with_as, h = h))
}
# run coupled chains until they meet
coupled_chains_ <- function(single_kernel, coupled_kernel, rinit, m = 1, max_iterations = Inf, preallocate = 10){
  chain_state1 <- rinit()
  chain_state2 <- rinit()
  dimstate <- length(chain_state1)
  samples1 <- matrix(nrow = m+preallocate+1, ncol = dimstate)
  samples2 <- matrix(nrow = m+preallocate, ncol = dimstate)
  nrowsamples1 <- m+preallocate+1
  samples1[1,] <- chain_state1
  samples2[1,] <- chain_state2
  current_nsamples1 <- 1
  chain_state1 <- single_kernel(chain_state1)$new_trajectory
  current_nsamples1 <- current_nsamples1 + 1
  samples1[current_nsamples1,] <- chain_state1
  iter <- 1
  meet <- FALSE
  finished <- FALSE
  meetingtime <- Inf
  while (!finished && iter < max_iterations){
    iter <- iter + 1
    if (meet){
      chain_state1 <- single_kernel(chain_state1)$new_trajectory
      chain_state2 <- chain_state1
    } else {
      res_coupled_kernel <- coupled_kernel(chain_state1, chain_state2)
      chain_state1 <- res_coupled_kernel$new_trajectory1
      chain_state2 <- res_coupled_kernel$new_trajectory2
      # to avoid numerical instabilities, meeting if distance is very small
      if (sqrt(sum((chain_state1 - chain_state2)^2)) < 1e-4 && !meet){
        # recording meeting time tau
        meet <- TRUE
        meetingtime <- iter
      }
    }
    if ((current_nsamples1+1) > nrowsamples1){
      # print('increase nrow')
      new_rows <- nrowsamples1-1
      nrowsamples1 <- nrowsamples1 + new_rows
      samples1 <- rbind(samples1, matrix(NA, nrow = new_rows, ncol = dimstate))
      samples2 <- rbind(samples2, matrix(NA, nrow = new_rows, ncol = dimstate))
    }
    samples1[current_nsamples1+1,] <- chain_state1
    samples2[current_nsamples1,] <- chain_state2
    current_nsamples1 <- current_nsamples1 + 1
    # stop after max(m, tau) steps
    if (iter >= max(meetingtime, m)){
      finished <- TRUE
    }
  }
  samples1 <- samples1[1:current_nsamples1,,drop=F]
  samples2 <- samples2[1:(current_nsamples1-1),,drop=F]
  return(list(samples1 = samples1, samples2 = samples2,
              meetingtime = meetingtime, iteration = iter, finished = finished))
}


# run coupled chains
res_ <- coupled_chains_(single_kernel_RB, coupled_kernel_RB, rinit, m = 30)
res_$meetingtime
#
df1 <- melt(res_$samples1)
names(df1) <- c("iteration", "time", "value")
df1$chain <- "1"
df1$iteration <- df1$iteration-1 
df2 <- melt(res_$samples2)
names(df2) <- c("iteration", "time", "value")
df2$chain <- "2"
df <- rbind(df1,df2) %>% filter(iteration > 0)
df$time <- df$time - 1

library(ggplot2)
ggplot(df %>% filter(iteration == 30), aes(x = time, y = value, group = chain, colour = chain)) + geom_point() + geom_line() +
  theme_minimal() + theme(legend.position = "none") +
  ggthemes::scale_color_colorblind() +
  labs(title = "Trajectories generated by coupled conditional particle filters",
       subtitle = "Iteration: {closest_state} / 20",
       x = "time",
       y = "latent state space") +
  geom_line(data = data.frame(time = 0:length(observations[,1]), value = x[,1]), aes(group = NULL, colour = NULL),
            linetype = 2)


library(gganimate)
ggplot(df %>% filter(iteration <= 25), aes(x = time, y = value, group = chain, colour = chain)) + geom_point() + geom_line() +
  theme_minimal() + theme(legend.position = "none") + ggthemes::scale_color_colorblind() +
  labs(title = "Trajectories generated by coupled conditional particle filters",
       subtitle = "iteration: {closest_state} / 20",
       x = "time",
       y = "latent state space") +
  transition_states(iteration, 3, 1) + 
  ease_aes('linear')

# anim_save()