## Replication script for article
## Inference tools for Markov Random Fields on lattices: The R package mrf2d
## Journal of Statistical Software
## Authors:
##  - Victor Freguglia
##  - Nancy Lopes Garcia
## Package mrf2d version 0.5
## Package ggplot2 version 3.3.2

## If the package is not installed please run the following commands
# install.packages("mrf2d") # CRAN version
## or
# devtools::install_github("Freguglia/mrf2d") # Development version

library("mrf2d")
set.seed(1)
## Please, set this seed and run all the commands in 
## sequence to obtain results identical to the article.

#### Section 3.1 ####

plot(mrfi(max_norm = 1))
plot(mrfi(max_norm = 2, norm_type = "m") + c(4, 0))
plot(mrfi(4) - mrfi(2))
plot(mrfi(6, norm_type = "m")[c(1, 2, 6, 9, 19, 41)])

#### Section 3.3 ####

th <- expand_array(-1, family = "onepar", mrfi(1), C = 1)

# Sample a field
z_sample <- rmrf2d(init_Z = c(200, 200), mrfi = mrfi(1), theta = th)

# Sample a field conditional to the border values being 0
## Define the logical matrix for the fixed region.
border <- matrix(FALSE, nrow = 100, ncol = 100) 
border[1, ] <- border[100, ] <- border[, 1] <- border[, 100] <- TRUE 
## Define the initial field
initial <- matrix(sample(0:1, 100*100, replace = TRUE), 
  nrow = 100, ncol = 100)
initial[border] <- 0
z_border <- rmrf2d(initial, mrfi = mrfi(1), theta = th, 
  fixed_region = border)

dplot(z_sample)
dplot(z_border)

#### Section 4.1 ####

data("field1", package = "mrf2d")
dplot(field1, legend = TRUE)

 # Define a large set of interacting positions
candidates <- mrfi(6, norm_type = "m")

 # Stochastic approximations for the candidate set
set.seed(1)
complete_sa <- fit_sa(field1, candidates, family = "oneeach", 
  gamma_seq = seq(from = 1, to = 0, length.out = 300), 
  cycles = 2, refresh_each = 301)

plot(complete_sa)

 # Threshold-based selection
thr_value <- 0.1
theta_vec <- smr_array(complete_sa$theta, "oneeach")
selected <- which(abs(theta_vec) > thr_value)

R1 <- candidates[selected]
R1

plot(candidates)
plot(R1) + ggplot2::xlim(-6.5, 6.5) + ggplot2::ylim(-6.5, 6.5)

pl <- fit_pl(field1, R1, family = "oneeach")
summary(pl)

sa <- fit_sa(field1, R1, family = "oneeach", 
  gamma_seq = seq(from = 1, to = 0, length.out = 300))
summary(sa)

z_sim <- rmrf2d(dim(field1), R1, pl$theta)

dplot(field1)
dplot(z_sim)

#### Section 4.2 ####

data("hfield1", package = "mrf2d")
cplot(hfield1)

hmrf_nofixed <- fit_ghm(hfield1, mrfi = R1, theta = pl$theta)
summary(hmrf_nofixed)

hmrf_poly <- fit_ghm(hfield1, mrfi = R1, theta = pl$theta, 
  fixed_fn = polynomial_2d(c(3, 3), dim(hfield1)))
summary(hmrf_poly)

library("ggplot2")
cplot(hfield1) + ggtitle("(a)")
dplot(hmrf_nofixed$Z_pred, legend = TRUE) + ggtitle("(b)")
cplot(hmrf_poly$fixed) + ggtitle("(c)")
dplot(hmrf_poly$Z_pred, legend = TRUE) + ggtitle("(d)")

#### Section 4.3 ####

data("bold5000", package = "mrf2d")
cplot(bold5000)

Rnn <- mrfi(1)
theta_nn <- expand_array(-1, family = "onepar", C = 3, mrfi = Rnn)

set.seed(1)
fit_brain <- fit_ghm(bold5000, Rnn, theta_nn, equal_vars = TRUE)
fit_brain_ind <- fit_ghm(bold5000, Rnn, theta_nn*0, equal_vars = TRUE)


summary(fit_brain)
summary(fit_brain_ind)


dplot(fit_brain$Z_pred, legend = TRUE)
dplot(fit_brain_ind$Z_pred, legend = TRUE)


#### Appendix A ####
library("ggplot2")
base_plot <- dplot(field1, legend = TRUE)

base_plot + ggtitle("This is a custom title")
base_plot + theme_void() + theme(legend.position = "none")
base_plot + scale_fill_manual(values = c("red", "blue"))
base_plot + theme(legend.position = "bottom")

mrfi_plot <- plot(mrfi(3) + c(5, 1))

# Custom colors
mrfi_plot + geom_tile(fill = "orange", color = "blue")
# Adding labels with geom_text
mrfi_plot + geom_text(aes(label = paste0("(", rx, ", ", ry, ")")))
# Custom title
mrfi_plot + ggtitle("Add a custom title") + 
  theme(plot.title = element_text(hjust = 0.5, size = 24))


#### Appendix B ####

# if the 'potts' package is not installed, run
# install.packages("potts")
library("potts")

dims <- c(32, 32)
nsims <- 100
C <- 2

setups <- expand.grid(package = c("mrf2d", "potts"), 
                      iterations = c(5, 15, 25, 35, 45, 55, 65, 75))

set.seed(1)
simulations <- lapply(1:nrow(setups), function(i){
  iters <- setups[i, 2]
  if(setups[i, 1] == "mrf2d"){
      sims <- sapply(seq_len(nsims), function(x){
        Z <- rmrf2d(dims, mrfi(1), expand_array(-1, "onepar", mrfi(1), C), 
                    cycles = iters)
        return(smr_stat(Z, mrfi(1), family = "onepar"))
      })
  } else {
      sims <- sapply(seq_len(nsims), function(x){
        Z <- matrix(sample(0:C, prod(dims), replace = TRUE), 
                    nrow = dims[1], ncol = dims[2])
        Z <- unpackPotts(potts(packPotts(Z + 1, ncolor = 3), 
                               param = c(0, 0, 0, 1), boundary = "free", 
                               nbatch = 1, nspac = iters)$final)
        return(smr_stat(Z-1, mrfi(1), family = "onepar"))
      })
  }
  return(sims)
})

library("dplyr")
df <- bind_rows(lapply(simulations, function(x) data.frame(stat = x)), .id = "id")
df$package = setups$package[as.numeric(df$id)]
df$iters = setups$iterations[as.numeric(df$id)]

ggplot(df, aes(y = stat, x = iters, group = iters)) + 
  geom_jitter(width = 2, height = 0, alpha = 0.5, aes(color = as.factor(iters))) +
  facet_wrap(~package) +
  theme_bw() +
  theme(legend.position = "none") +
  xlab("Number of cycles/iterations") +
  ylab("Sufficient statistics")

set.seed(2)
zlis <- lapply(1:100, function(x) rmrf2d(c(64, 64), mrfi(1), expand_array(-1, "onepar", mrfi(1), 2)))

est_mrf2d <- lapply(zlis, fit_pl, mrfi = mrfi(1), family = "onepar")
est_mrf2d <- unlist(lapply(est_mrf2d, function(x) smr_array(x$theta, "onepar")))

## This may take several minutes to run
est_potts <- lapply(zlis, function(x){
  t_stat <- potts::calc_t(x + 1, ncolor = 3)
  t_cache_mple <- generate_t_cache(x + 1, 3, t_stat, 64*64, 1, 
    singleton)
  objective <- function(phi){
    composite.ll(c(0, 0, phi), t_stat, t_cache_mple)
  }
  optim(1, objective, control = list(fnscale = -1), 
        method = "Brent", lower = -10, upper = 10)$par
})
est_potts <- -unlist(est_potts)

df_comp <- data.frame(mrf2d = est_mrf2d, potts = est_potts)

ggplot(df_comp, aes(x = est_mrf2d, y = est_potts)) + 
  geom_point() +
  geom_abline(slope = 1, intercept = 0) +
  geom_hline(yintercept = -1, linetype = "dashed") +
  geom_vline(xintercept = -1, linetype = "dashed") +
  xlab("mrf2d") +
  ylab("potts") +
  theme_bw()