cpmmc:

# Random Effect Model Experiments ----


# Wrapper Functions
# ------------------------------------


# Calculate acceptance probability
acceptance_prob <- function(theta_chain) {
  length(unique(theta_chain)) / length(theta_chain)
}


# MC diagnostics
#--------------------------------
h <- function(chain, max_lag) 1+2*sum((acf(chain, plot = F)$acf[seq_len(max_lag)])^2)


# For MH
# --------------------------------
log_target_density <- function(data, theta){
  (sum(dnorm(data, mean = theta, sd = sqrt(2), log = T))) + dnorm(theta, log = T, sd = 1)
}

log_target_density_theta <- function(theta){
  log_target_density(data, theta)
}

proposal_sampler <- function(state){
  rnorm(1, mean = state, sd = 0.02)
}

log_proposal_density <- function(state, proposal){
  dnorm(proposal, mean = state, log=T, sd = 0.02)
}

# For PMMC / CPMMC
# --------------------------------
simulation_study <- function(T_, N_, theta_0, rho, nsum, burnin, u_0, data){

  # Instantiate model ----
  rem <- normal_random_effect_model(data,
                                    theta_0 = theta_0,
                                    u_0 = u_0,
                                    rho = rho,
                                    log_theta_prior_density = function(x) dnorm(x, log = T, sd = 1),
                                    log_theta_proposal_density = function(old_theta, new_theta) dnorm(new_theta-old_theta, log = T, sd=0.02),
                                    theta_proposal_sampler = function(x) rnorm(1, mean = x, sd = 0.02)
  )


  # Create results store
  Z_results <- list()
  W_results <- list()

  # Run model
  start_time <- Sys.time()
  for(i in seq_len(nsim)){
    if (i %% 100 == 0){
      current_time <- Sys.time()
      cat(paste0(">> Iteration: ",i, "\n"))
      cat(paste0(">> Time elapsed: ", current_time-start_time, " \n \n" ))

    }

    rem <- single_transition(rem)
    W_results[[i]] <- rem$log_new_marginal_estimator -
      sum(dnorm(rem$data, mean = rem$proposed_chain[[i]], sd = sqrt(2), log = T))

    Z_results[[i]] <- rem$log_old_marginal_estimator -
      sum(dnorm(rem$data, mean = rem$chain[[i]], sd = sqrt(2), log = T))
  }

  return(
    list(
      w = sapply(W_results, function(x) x[[1]]),
      z = sapply(Z_results, function(x) x[[1]]),
      object = rem
    )
  )
}



# Experiment 2
# --------------------------------
## Parameters ----

T_ <- 1024
N_ <- 19
rho <- 0.9894
nsim <- 10^4
burnin <- 100
theta_0 <- 0

# Generate data and initialise ----
u_0 <- array(rnorm(T_*N_), dim = c(N_,1,T_))
data <- rnorm(T_,0.5, sd = sqrt(2))

# Run comparisons
exp_2_cpm <- simulation_study(T_, N_, theta_0, rho, nsim, burnin, u_0, data)
exp_2_pm <- simulation_study(T_, N_, theta_0, 0, nsim, burnin, u_0, data)

mh <- metropolis_hastings(theta_0, log_target_density_theta, log_proposal_density, proposal_sampler)
mh <- run_chain(mh, nsim)
exp_2_mh <- mh

serialize_robject("./Simulation_Study/cpm_exp2", exp_2_cpm)
serialize_robject("./Simulation_Study/pm_exp2", exp_2_pm)
serialize_robject("./Simulation_Study/mh_exp2", exp_2_mh)


# Experiment 3
# --------------------------------
## Parameters ----

T_ <- 2048
N_ <- 28
rho <- 0.9925
nsim <- 10^4
burnin <- 100
theta_0 <- 0

# Generate data and initialise ----
u_0 <- array(rnorm(T_*N_), dim = c(N_,1,T_))
data <- rnorm(T_,0.5, sd = sqrt(2))

# Run comparisons
exp_3_cpm <- simulation_study(T_, N_, theta_0, rho, nsim, burnin, u_0, data)
exp_3_pm <- simulation_study(T_, N_, theta_0, 0, nsim, burnin, u_0, data)

mh <- metropolis_hastings(theta_0, log_target_density_theta, log_proposal_density, proposal_sampler)
mh <- run_chain(mh, nsim)
exp_3_mh <- mh

serialize_robject("./Simulation_Study/cpm_exp3", exp_3_cpm)
serialize_robject("./Simulation_Study/pm_exp3", exp_3_pm)
serialize_robject("./Simulation_Study/mh_exp3", exp_3_mh)



# Experiment 4
# --------------------------------
## Parameters ----

T_ <- 4096
N_ <- 39
rho <- 0.9947
nsim <- 10^4
burnin <- 100
theta_0 <- 0

# Generate data and initialise ----
u_0 <- array(rnorm(T_*N_), dim = c(N_,1,T_))
data <- rnorm(T_,0.5, sd = sqrt(2))

# Run comparisons
exp_4_cpm <- simulation_study(T_, N_, theta_0, rho, nsim, burnin, u_0, data)
exp_4_pm <- simulation_study(T_, N_, theta_0, 0, nsim, burnin, u_0, data)

mh <- metropolis_hastings(theta_0, log_target_density_theta, log_proposal_density, proposal_sampler)
mh <- run_chain(mh, nsim)
exp_4_mh <- mh

serialize_robject("./Simulation_Study/cpm_exp4", exp_4_cpm)
serialize_robject("./Simulation_Study/pm_exp4", exp_4_pm)
serialize_robject("./Simulation_Study/mh_exp4", exp_4_mh)


# Experiment 5
# --------------------------------
## Parameters ----

T_ <- 8192
N_ <- 56
rho <- 0.9962
nsim <- 10^4
burnin <- 100
theta_0 <-0

# Generate data and initialise ----
u_0 <- array(rnorm(T_*N_), dim = c(N_,1,T_))
data <- rnorm(T_,0.5, sd = sqrt(2))

# Run comparisons
exp_5_cpm <- simulation_study(T_, N_, theta_0, rho, nsim, burnin, u_0, data)
# Time elapsed: 23.7469944794973 mins
exp_5_pm <- simulation_study(T_, N_, theta_0, 0, nsim, burnin, u_0, data)
# Time elapsed: 23.771907377243 mins
mh <- metropolis_hastings(theta_0, log_target_density_theta, log_proposal_density, proposal_sampler)
mh <- run_chain(mh, nsim)
# Time elapsed: 6.619749 secs

exp_5_mh <- mh

serialize_robject("./Simulation_Study/cpm_exp5", exp_5_cpm)
serialize_robject("./Simulation_Study/pm_exp5", exp_5_pm)
serialize_robject("./Simulation_Study/mh_exp5", exp_5_mh)



# Experiment 6 (5.2)
# --------------------------------
## Parameters ----

T_ <- 8192
N_ <- 56
rho <- 0.9962
nsim <- 10^4
burnin <- 100
theta_0 <- 0.5

# Generate data and initialise ----
u_0 <- array(rnorm(T_*N_), dim = c(N_,1,T_))
data <- rnorm(T_,0.5, sd = sqrt(2))

# Run comparisons
exp_6_cpm <- simulation_study(T_, N_, theta_0, rho, nsim, burnin, u_0, data)
# Time elapsed: 23.7850300033887 mins
exp_6_pm <- simulation_study(T_, N_, theta_0, 0, nsim, burnin, u_0, data)
# Time elapsed: 24.0089103539785 mins
mh <- metropolis_hastings(theta_0, log_target_density_theta, log_proposal_density, proposal_sampler)
mh <- run_chain(mh, nsim)
# Time elapsed: 6.619749 secs
exp_6_mh <- mh

serialize_robject("./Simulation_Study/cpm_exp6", exp_6_cpm)
serialize_robject("./Simulation_Study/pm_exp6", exp_6_pm)
serialize_robject("./Simulation_Study/mh_exp6", exp_6_mh)




# Experiment 1
# --------------------------------
## Parameters ----

T_ <- 8192
N_ <- 80
rho <- 0.9963
nsim <- 10^4
burnin <- 100
theta_0 <- 0
# Generate data and initialise ----
u_0 <- array(rnorm(T_*N_), dim = c(N_,1,T_))
data <- rnorm(T_,0.5, sd = sqrt(2))

long_run_CPM <- simulation_study(T_, N_, theta_0, rho, nsum, burnin, u_0, data)
serialize_robject("./Simulation_Study/cpm_exp1", long_run_CPM)
# Time elapsed: 34.5729867140452 mins
long_run_PM <- simulation_study(T_, N_, theta_0, 0, nsum, burnin, u_0, data)
serialize_robject("./Simulation_Study/pm_exp1", long_run_PM)
# Time elapsed: 29.9212190111478 mins
mh <- metropolis_hastings(theta_0, log_target_density_theta, log_proposal_density, proposal_sampler)
long_run_MH <- run_chain(mh, nsim)
serialize_robject("./Simulation_Study/mh_exp1", long_run_MH)
# Time elapsed: 6.403306 secs

# Results
# -----------------------------------------------
file_templates <- c("./Simulation_Study/mh_exp", "./Simulation_Study/pm_exp","./Simulation_Study/cpm_exp")
file_locations <- unlist(lapply(1:6, function(x) paste0(file_templates,x)))

# Aggregate model runs
results <- list()
for (fp in file_locations){
  print(fp)
  results[[fp]] <- list()
  object <- unserialize_robject(fp)
  if (grepl(pattern = 'pm', x = fp)){
    thetas <- sapply(object$object$chain, function(x) x[[1]])
    results[[fp]][['w']] <- object$w
    results[[fp]][['z']] <- object$z
  } else if (grepl(pattern = 'mh', x = fp)){
    thetas <- sapply(object$chain, function(x) x[[1]])
  }
  results[[fp]][['IF']] <- h(thetas, max_lag = 40)
  results[[fp]][['thetas']] <- thetas
  results[[fp]][['A']] <- acceptance_prob(thetas)

}
#Save summary
serialize_robject('./Simulation_Study/Results_Summary.blob', results)
results <- unserialize_robject('./Simulation_Study/Results_Summary.blob')

# Long run Plots
#----------------------------------------------
fp <- './Simulation_Study/cpm_exp1'
thetas <- results[[fp]]$theta
w <- results[[fp]]$w
z <- results[[fp]]$z



# Detailed plot of long run CPM
# -------------------------------------------------
obj <- unserialize_robject('./Simulation_Study/cpm_exp1')

par(mfrow=c(2,3))
burnin <- 1
plot_ind <- burnin:(length(z)-burnin)
plot(plot_ind + burnin, w[plot_ind], type='l', col ='blue', ylab='', xlab='')
points(plot_ind + burnin,z[plot_ind], type='l', col='red')

r = w-z
plot(plot_ind + burnin, r[plot_ind], col='red', type='l', ylab='', xlab='')

hist(z[1000:10^4], breaks = 100, freq = F, main ='Z', xlab ='')
hist(r, breaks = 100, freq=F, main = 'R', xlab ='')
acf(thetas, main = '')


prop <- sapply(obj$object$proposed_chain, function(x) x[[1]])
acc <- sapply(obj$object$chain, function(x) x[[1]])
plot(plot_ind + burnin, prop[plot_ind], col='red', type='l',ylab='', xlab='')
points(plot_ind + burnin, acc[plot_ind], col='blue', type='l',ylab='', xlab='')



# Mixing speed, ACF
# -------------------------------------------------
model_runs <- c(
                 "./Simulation_Study/mh_exp1",
                 "./Simulation_Study/cpm_exp1",
                 "./Simulation_Study/pm_exp1")

titles <- c('MH', 'CPM', 'PM')
par(mfrow = c(1,3))
for (i in seq_along(model_runs)){
  fp <- model_runs[i]
  title <- titles[i]
  theta <- results[[fp]]$theta
  acf(theta, main = title)
}