###################################################################################################################
# Replication materials of "pexm: a JAGS module for applications involving the piecewise exponential distribution".
# Authors: Vinicius D. Mayrink, Joao D. N. Duarte and Fabio N. Demarqui.
# Departamento de Estatistica, ICEx, Universidade Federal de Minas Gerais.
###################################################################################################################
# Commands to reproduce results in the Section 4 of the paper
###################################################################################################################

# Important remark: Make sure that the R working directory is the same one where this R script is located.
# This is important for the correct identification of the files containg the Bayesian model written in the JAGS syntax.

###################################################################
# Section 4: R code related to the real data application 
##################################################################

set.seed(1981)

# Load some required R libraries.
if(!require("survival")){ install.packages("survival"); library("survival")}
if(!require("rjags")){ install.packages("rjags"); library("rjags")}
library("pexm")
loadpexm()

wd <- getwd() # Identify the working directory.

m <- 10 # Number of intervals (PE distribution).

# Load the real data set.
data(kidney, package = "survival")
head(kidney)
n <- max(kidney$id)
id <- matrix(kidney$id, n, 2, byrow=TRUE)
age <- matrix(kidney$age, n, 2, byrow = TRUE) 
sex <- matrix(kidney$sex-1, n, 2, byrow = TRUE) 
sex <- sex[,1]
delta <- matrix(kidney$status, n, 2, byrow = TRUE) 
censored <- 1 - delta
y <- matrix(kidney$time, n, 2, byrow = TRUE)
t <- y
t_cen <- array(0, c(n, 2))
t[delta == 0] <- NA
t_cen[delta == 0] <- y[delta == 0]

# Organizing the information to be passed to JAGS.
data_pexm <- list(t = t, t_cen = t_cen, censored = censored, n = n, m = m, sex = sex, age = age)
data_pois <- list(t = y, delta = delta, n = n, m = m, sex = sex, age = age)
parameters <- c("beta_sex","beta_age","kappa","loglik","lambda")
#
inits1 <- list( .RNG.name = "base::Super-Duper", .RNG.seed = 5 ) 
inits2 <- list( .RNG.name = "base::Wichmann-Hill", .RNG.seed = 2 )
inits_G = list(inits1, inits2)
#
inits1 <- list( .RNG.name = "base::Super-Duper", .RNG.seed = 5 ) 
inits2 <- list( .RNG.name = "base::Wichmann-Hill", .RNG.seed = 2 )
inits_N = list(inits1, inits2)
#
# MCMC setup.
burnin <- 10000
lag <- 1
npost <- 10000

# Models

# M1: Gamma distribution to associate adjacent rates
M1_Gamma_pexm <- "data{ 
  for (j in 1:m) { 
    a[j] <- 562 * (j - 1) / m # Time grid.
  } 
} 
model {  
  for (i in 1:n) {
    for (k in 1:2) {
      # Infection time bounded below by the censoring time.
      censored[i, k] ~ dinterval(t[i, k], t_cen[i, k])
      t[i, k] ~ dpex(haz[i, k,], a[])
      for(j in 1:m){ haz[i, k, j] <- lambda[j] * exp(beta_sex * sex[i] + beta_age * age[i,k] ) * z[i] }
      loglik[i, k] <- log( dpex(t[i, k], haz[i, k,], a[]) )
    }
  }
  for(i in 1:n) { 
    z[i] ~ dgamma(eta, eta) # Frailties.
  }  
  kappa <- 1 / eta	
  eta ~ dgamma(0.001, 0.001)
  beta_sex ~ dnorm(0, 0.001)
  beta_age ~ dnorm(0, 0.001)
  rate[1] <- 0.01 / 1
  lambda[1] ~ dgamma(0.01, rate[1])  
  for(j in 2:m){
    rate[j] <- 0.01 / lambda[j-1]
    lambda[j] ~ dgamma(0.01, rate[j])
  }  
}
"

# M2: Normal distribution to associate adjacent rates.
M2_Normal_pexm <- "
data{ for(j in 1:m){ a[j] <- 562 * (j - 1) / m } } # Time grid.
model{  
  for (i in 1:n){ 
    for(k in 1:2){
      # Infection time bounded below by the censoring time.
      censored[i, k] ~ dinterval(t[i, k], t_cen[i, k])
      t[i, k] ~ dpex(haz[i, k,], a[])
      for(j in 1:m){ haz[i, k, j] <- lambda[j] * exp( beta_sex * sex[i] + beta_age * age[i, k] ) * z[i] }
      loglik[i,k] <- log( dpex(t[i, k], haz[i, k,], a[]) )
    }
  }
  for(i in 1:n){ z[i] ~ dgamma(eta, eta) } # frailties.
  kappa <- 1 / eta	
  eta ~ dgamma(0.001, 0.001)
  beta_sex ~ dnorm(0, 0.001)
  beta_age ~ dnorm(0, 0.001)
  xi[1] ~ dnorm(0, 0.0001)
  lambda[1] <- exp(xi[1])  
  for(j in 2:m){
    xi[j] ~ dnorm(xi[j-1], 0.0001)
    lambda[j] <- exp(xi[j])
  }  
}
"
# M1: Gamma distribution to associate adjacent rates
M1_Gamma_pois <- "
data{ 
  for(j in 1:(m+1)){ a[j] <- 562 * (j - 1) / m } # Time grid (including the max. infection time).
  for(i in 1:n){ for(k in 1:2){ zeros[i, k] <- 0 } } # Poisson-zeros strategy.
}
model{ 
  for(i in 1:n){ 
    for(k in 1:2){
      for(j in 1:m){
        # Indicator (event-time in interval j).
        d[i, k, j] <- delta[i, k] * step(t[i, k] - a[j]) * step(a[j+1] - t[i,k])
        # Length of overlap between (0, t[i, k]) and interval j.
        le[i, k, j] <- (min(t[i, k], a[j+1]) - a[j]) * step(t[i, k] - a[j])
        # Proportional hazard with frailty z[i] and PE hazard rate lambda[j].
        theta[i, k, j] <- lambda[j] * exp( beta_sex * sex[i] + beta_age * age[i, k] ) * z[i]
        # Likelihood.
        mu[i, k, j] <- le[i, k, j] * theta[i, k, j]
        lik[i, k, j] <- dpois(d[i, k, j], mu[i, k, j])
      }
      loglik[i, k] <- sum(log(lik[i, k,]))
      zeros[i, k] ~ dpois(-loglik[i, k])
    }
  } 
  for(i in 1:n){ z[i] ~ dgamma(eta, eta) } # Frailties 
  kappa <- 1 / eta	
  eta ~ dgamma(0.001, 0.001)
  beta_sex ~ dnorm(0, 0.001)
  beta_age ~ dnorm(0, 0.001)
  rate[1] <- 0.01 / 1
  lambda[1] ~ dgamma(0.01, rate[1])  
  for(j in 2:m){
    rate[j] <- 0.01 / lambda[j-1]
    lambda[j] ~ dgamma(0.01, rate[j])
  } 
}
"
# M2: Normal distribution to associate adjacent rates.
M2_Normal_pois <- "data{ 
  for(j in 1:(m+1)){ 
    a[j] <- 562 * (j - 1) / m  # Time grid (including the max. infection time).
  }
  for(i in 1:n){ 
    for(k in 1:2){ 
      zeros[i, k] <- 0  # Poisson-zeros strategy.
    } 
  }
}
model{ 
  for(i in 1:n){
    for(k in 1:2){
      for(j in 1:m){
        # Indicator (event-time in interval j)
        d[i, k, j] <- delta[i, k] * step(t[i, k] - a[j]) * step(a[j+1] - t[i, k])
        # Length of overlap between (0,t[i,k]) and interval j.
        le[i, k, j] <- (min(t[i, k], a[j+1]) - a[j]) * step(t[i, k] - a[j])
        # Proportional hazard with frailty z[i] and PE hazard rate lambda[j].
        theta[i,k,j] <- lambda[j]*exp( beta_sex * sex[i] + beta_age * age[i,k] ) * z[i]
        # Likelihood. 
        mu[i, k, j] <- le[i, k, j] * theta[i, k, j]
        lik[i, k, j] <- dpois(d[i, k, j], mu[i, k, j])
      }
      loglik[i, k] <- sum(log(lik[i, k,]))
      zeros[i, k] ~ dpois(-loglik[i, k])
    }
  } 
  for(i in 1:n){ z[i] ~ dgamma(eta, eta) } # Frailties. 
  kappa <- 1 / eta	
  eta ~ dgamma(0.001, 0.001)
  beta_sex ~ dnorm(0, 0.001)
  beta_age ~ dnorm(0, 0.001)
  xi[1] ~ dnorm(0, 0.0001)
  lambda[1] <- exp(xi[1])  
  for(j in 2:m){
    xi[j] ~ dnorm(xi[j-1], 0.0001)
    lambda[j] <- exp(xi[j])
  } 
}
"

# Compile the models.
M1jags_pexm <- rjags::jags.model(textConnection(M1_Gamma_pexm),  data = data_pexm, inits = inits_G, n.chains = 2, n.adapt = burnin)
M2jags_pexm <- rjags::jags.model(textConnection(M2_Normal_pexm), data = data_pexm, inits = inits_N, n.chains = 2, n.adapt = burnin)
M1jags_pois <- rjags::jags.model(textConnection(M1_Gamma_pois),  data = data_pois, inits = inits_G, n.chains = 2, n.adapt = burnin)
M2jags_pois <- rjags::jags.model(textConnection(M2_Normal_pois), data = data_pois, inits = inits_N, n.chains = 2, n.adapt = burnin)

# Run JAGS.
set.seed(1981)
time0 <- proc.time()
output_M1_pexm <- rjags::coda.samples(M1jags_pexm, variable.names = parameters, n.iter = npost, thin = lag)
time_M1_pexm <- proc.time() - time0
#
set.seed(1981)
time0 <- proc.time()
output_M2_pexm <- rjags::coda.samples(M2jags_pexm, variable.names = parameters, n.iter = npost, thin = lag)
time_M2_pexm <- proc.time() - time0
#
set.seed(1981)
time0 <- proc.time()
output_M1_pois <- rjags::coda.samples(M1jags_pois, variable.names = parameters, n.iter = npost, thin = lag)
time_M1_pois <- proc.time() - time0
#
set.seed(1981)
time0 <- proc.time()
output_M2_pois <- rjags::coda.samples(M2jags_pois, variable.names = parameters, n.iter = npost, thin = lag)
time_M2_pois <- proc.time() - time0

# Computational times required to run each MCMC
# This result depends on the computer being used for the simulations.
# The outcome is discussed in the end of Section 4.
round(c(time_M1_pexm[1], time_M2_pexm[1], time_M1_pois[1], time_M2_pois[1]), 4)

###########################################################
# Code to reproduce Table 1 (Section 4).
# Table of posterior estimates (mean, median, s.d. and HPD) 
# for each combination (parameter, implementation, model). 
# Notation: bs = beta_sex, ba = beta_age, k = kappa. 
###########################################################
library("coda")
#
pname <- "beta_sex"
aux1 <- as.matrix(output_M1_pexm); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res1 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M1_pois); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res2 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M2_pexm); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res3 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M2_pois); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res4 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
T_beta_sex <- rbind(res1, res2, res3, res4); 
rownames(T_beta_sex) <- c("bs Module  M.I","bs Poisson M.I","bs Module  M.II","bs Poisson M.II"); 
colnames(T_beta_sex) <- c("mean","median","s.d.","HPD inf.","HPD sup.");
#
pname <- "beta_age"
aux1 <- as.matrix(output_M1_pexm); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res1 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M1_pois); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res2 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M2_pexm); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res3 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M2_pois); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res4 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
T_beta_age <- rbind(res1, res2, res3, res4); 
rownames(T_beta_age) <- c("ba Module  M.I","ba Poisson M.I","ba Module  M.II","ba Poisson M.II"); 
colnames(T_beta_age) <- c("mean","median","s.d.","HPD inf.","HPD sup.");
#
pname <- "kappa"
aux1 <- as.matrix(output_M1_pexm); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res1 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M1_pois); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res2 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M2_pexm); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res3 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
aux1 <- as.matrix(output_M2_pois); aux2 = as.mcmc(aux1); HPD = HPDinterval(aux2)
res4 <- c(apply(aux1, 2, mean)[pname], apply(aux1, 2, median)[pname],
         apply(aux1, 2, sd)[pname], HPD[pname,])
T_kappa <- rbind(res1, res2, res3, res4); 
rownames(T_kappa) <- c("k Module  M.I","k Poisson M.I","k Module  M.II","k Poisson M.II"); 
colnames(T_kappa) <- c("mean","median","s.d.","HPD inf.","HPD sup.");
#
T_final = rbind(T_beta_sex, T_beta_age, T_kappa)
round(T_final, 4)

###############################################
# Code to reproduce Figure 3 (Section 4).
# Trace plots displaying the first chain:
# Module in red, Poisson-zeros in black, 
# row 1 = Model I, row 2 = Model II.
###############################################
co1 <- "black"; co2 <- "tomato"; co3 = "royalblue";
labels = c("beta_sex","beta_age","kappa")
# Effective sample sizes for the chains.
eff_M1_pexm <- round(coda::effectiveSize(output_M1_pexm[1])[labels], 1)
eff_M1_pois <- round(coda::effectiveSize(output_M1_pois[1])[labels], 1)
eff_M2_pexm <- round(coda::effectiveSize(output_M2_pexm[1])[labels], 1)
eff_M2_pois <- round(coda::effectiveSize(output_M2_pois[1])[labels], 1)
#
aux_pexm <- as.matrix(output_M1_pexm[1])
aux_pois <- as.matrix(output_M1_pois[1])
l_sex <- range(aux_pexm[,"beta_sex"]) 
l_age <- range(aux_pexm[,"beta_age"]) 
l_k <- range(aux_pexm[,"kappa"])
par(mfrow = c(2,3))
plot(aux_pois[, "beta_sex"], type = "l", xlab = "iterations", 
     ylab = expression(paste(beta[sex],"  (Model I)", sep = "")),
     main = paste0("ESS: ",eff_M1_pexm[1]," (pexm), ",eff_M1_pois[1]," (Pois.)"),
     ylim = l_sex, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 1.5, col = co1)
lines(aux_pexm[,"beta_sex"], lwd = 1.5, col = co2)
abline(h = -2.430, lwd = 5, col = co3, lty = 2)
abline(h = -1.493, lwd = 5, col = co3, lty = 1)
abline(h = -0.600, lwd = 5, col = co3, lty = 2)
plot(aux_pois[, "beta_age"], type = "l", xlab = "iterations", 
     ylab = expression(paste(beta[age],"  (Model I)", sep = "")),
     main = paste0("ESS: ",eff_M1_pexm[2]," (pexm), ",eff_M1_pois[2]," (Pois.)"),
     ylim = l_age, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 1.5, col = co1)
lines(aux_pexm[, "beta_age"], lwd = 1.5, col = co2)
abline(h = -0.018, lwd = 5, col = co3, lty = 2)
abline(h = 0.006, lwd = 5, col = co3, lty = 1)
abline(h = 0.032, lwd = 5, col = co3, lty = 2)
plot(aux_pois[, "kappa"], type = "l", xlab = "iterations", 
     ylab = expression(paste(kappa,"  (Model I)", sep = "")),
     main = paste0("ESS: ",eff_M1_pexm[3]," (pexm), ",eff_M1_pois[3]," (Pois.)"),
     ylim = l_k, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 1.5, col = co1)
lines(aux_pexm[, "kappa"], lwd = 1.5, col = co2)
abline(h = 0.061, lwd = 5, col = co3, lty = 2)
abline(h = 0.499, lwd = 5, col = co3, lty = 1)
abline(h = 1.160, lwd = 5, col = co3, lty = 2)
#
aux_pexm <- as.matrix(output_M2_pexm[1])
aux_pois <- as.matrix(output_M2_pois[1])
plot(aux_pois[, "beta_sex"], type = "l", xlab = "iterations",
     ylab = expression(paste(beta[sex],"  (Model II)", sep = "")),
     main = paste0("ESS: ",eff_M2_pexm[1]," (pexm), ",eff_M2_pois[1]," (Pois.)"),
     ylim = l_sex, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 1.5, col = co1)
lines(aux_pexm[, "beta_sex"], lwd = 1.5, col = co2)
abline(h = -2.467, lwd = 5, col = co3, lty = 2)
abline(h = -1.500, lwd = 5, col = co3, lty = 1)
abline(h = -0.624, lwd = 5, col = co3, lty = 2)
plot(aux_pois[, "beta_age"], type = "l", xlab = "iterations", 
     ylab = expression(paste(beta[age],"  (Model II)", sep = "")),
     main = paste0("ESS: ",eff_M2_pexm[2]," (pexm), ",eff_M2_pois[2]," (Pois.)"),
     ylim = l_age, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 1.5, col = co1)
lines(aux_pexm[, "beta_age"], lwd = 1.5 , col = co2)
abline(h = -0.018, lwd = 5, col = co3, lty = 2)
abline(h =  0.007, lwd = 5, col = co3, lty = 1)
abline(h =  0.036, lwd = 5, col = co3, lty = 2)
plot(aux_pois[, "kappa"], type = "l", xlab = "iterations",
     ylab = expression(paste(kappa,"  (Model II)", sep = "")),
     main = paste0("ESS: ",eff_M2_pexm[3]," (pexm), ",eff_M2_pois[3]," (Pois.)"),     
     ylim = l_k, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 1.5, col = co1)
lines(aux_pexm[, "kappa"], lwd = 1.5, col = co2)
abline(h = 0.089, lwd = 5, col = co3, lty = 2)
abline(h = 0.523, lwd = 5, col = co3, lty = 1)
abline(h = 1.195, lwd = 5, col = co3, lty = 2)

##############################################
# Code to reproduce Figure 4 (Section 4).
# Densities (histogram shape) representing 
# the posterior distribution:
# Module in red, Poisson-zeros in black, 
# row 1 = Model I, row 2 = Model II.
##############################################
co1 = "black"; co2 = "tomato"
aux_pexm <- as.matrix(output_M1_pexm)
aux_pois <- as.matrix(output_M1_pois)
l_sex <- range(aux_pexm[, "beta_sex"])
l_age <- range(aux_pexm[, "beta_age"])
l_k <- range(aux_pexm[, "kappa"]);
par(mfrow = c(2, 3))
den_pexm <- density(aux_pexm[, "beta_sex"])
den_pois <- density(aux_pois[, "beta_sex"])
max_den <- max(c(den_pexm$y, den_pois$y))
plot(den_pois$x, den_pois$y, type = "l", ylab = "density", xlab = "", main = expression(paste(beta[sex], "  (Model I)", sep = "")),  
     ylim = c(0, max_den), xlim = l_sex, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 7, col = co1)
lines(den_pexm$x, den_pexm$y, lwd = 3, col = co2)
den_pexm <- density(aux_pexm[, "beta_age"])
den_pois <- density(aux_pois[, "beta_age"])
max_den <- max(c(den_pexm$y, den_pois$y))
plot(den_pois$x, den_pois$y, type = "l", ylab = "", xlab = "", main = expression(paste(beta[age], "  (Model I)", sep = "")),  
     ylim = c(0, max_den), xlim = l_age, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 7, col = co1)
lines(den_pexm$x,  den_pexm$y, lwd = 3, col = co2)
den_pexm <- density(aux_pexm[, "kappa"])
den_pois <- density(aux_pois[, "kappa"])
max_den <- max(c(den_pexm$y, den_pois$y))
plot(den_pois$x, den_pois$y, type = "l", ylab = "", xlab = "", main = expression(paste(kappa, "  (Model I)", sep = "")),  
     ylim = c(0, max_den), xlim = l_k, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 7, col = co1)
lines(den_pexm$x, den_pexm$y, lwd = 3, col = co2)
#
aux_pexm <- as.matrix(output_M2_pexm)
aux_pois <- as.matrix(output_M2_pois)
den_pexm <- density(aux_pexm[, "beta_sex"])
den_pois <- density(aux_pois[, "beta_sex"])
max_den <- max(c(den_pexm$y, den_pois$y))
plot(den_pois$x, den_pois$y, type = "l", ylab = "density", xlab = "", main = expression(paste(beta[sex], "  (Model II)", sep = "")),  
     ylim = c(0, max_den), xlim = l_sex, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 7, col = co1)
lines(den_pexm$x, den_pexm$y, lwd = 3, col = co2)
den_pexm <- density(aux_pexm[, "beta_age"])
den_pois <- density(aux_pois[, "beta_age"])
max_den <- max(c(den_pexm$y, den_pois$y))
plot(den_pois$x, den_pois$y, type = "l", ylab = "", xlab = "", main = expression(paste(beta[age], "  (Model II)", sep = "")),  
     ylim = c(0, max_den), xlim = l_age, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 7, col = co1)
lines(den_pexm$x, den_pexm$y, lwd = 3, col = co2)
den_pexm <- density(aux_pexm[, "kappa"])
den_pois <- density(aux_pois[, "kappa"])
max_den <- max(c(den_pexm$y, den_pois$y))
plot(den_pois$x, den_pois$y, type = "l", ylab = "", xlab = "", main = expression(paste(kappa, "  (Model II)", sep = "")),  
     ylim = c(0, max_den), xlim = l_k, cex.lab = 2, cex.axis = 2, cex.main = 2, lwd = 7, col = co1)
lines(den_pexm$x, den_pexm$y, lwd = 3, col = co2)


###############################################################
# Code to reproduce Figure 5 (Section 4).
# This code can also be used to create Figure B.2 (Appendix B). 
# Posterior estimates of all lambda[j]'s in the PE distribution.
# Graph displaying the estimated baseline hazard function: 
# time axis vs. posterior mean of lambda.
# Module in red, Poisson strategy in black.
###############################################################
stepf <- function(rate, grid)
{
  nr <- length(rate)
  y <- 0; x <- 0
  j <- 1 
  for(i in 1:nr){
    y <- c(y, rate[i], rate[i])
    x <- c(x, grid[j], grid[j+1])
    j <- j+1
  }
  y <- c(y, 0); x <- c(x, max(grid))
  return(list(x = x, y = y))
}
#
par(mfrow = c(2,2))
lnames <- paste0("lambda[", 1:m, "]")
ma <- 0.55
co1 = "black"; co2 = "tomato"
#
aux_M1_pexm <- as.matrix(output_M1_pexm)
aux_M1_pois <- as.matrix(output_M1_pois)
aux_M2_pexm <- as.matrix(output_M2_pexm)
aux_M2_pois <- as.matrix(output_M2_pois)
#
lam_pexm <- apply(aux_M1_pexm[, lnames], 2, mean)
lam_pois <- apply(aux_M1_pois[, lnames], 2, mean)
a <- NULL; for(j in 1:m){ a[j] <- 562 * (j-1) / m }; a <- c(a, 562);
rates_pexm <- stepf(lam_pexm, a)
rates_pois <- stepf(lam_pois, a)
plot(rates_pois$x, rates_pois$y, xlab = "time", ylab = "baseline hazard function", col = co1, 
     main = expression("Model I (mean)"), type = "l", ylim = c(0, ma),lwd = 7, cex.lab = 2, cex.axis = 2, cex.main = 2)
lines(rates_pexm$x, rates_pexm$y, col = co2, lwd = 3)
#
lam_pexm <- apply(aux_M1_pexm[, lnames], 2, median)
lam_pois <- apply(aux_M1_pois[, lnames], 2, median)
a <- NULL; for(j in 1:m){ a[j] <- 562 * (j-1) / m }; a <- c(a, 562);
rates_pexm <- stepf(lam_pexm, a)
rates_pois <- stepf(lam_pois, a)
plot(rates_pois$x, rates_pois$y, xlab = "time", ylab = "baseline hazard function", col = co1, 
     main = expression("Model I (median)"), type = "l", ylim = c(0, ma), lwd = 7, cex.lab = 2, cex.axis = 2, cex.main = 2)
lines(rates_pexm$x, rates_pexm$y, col = co2, lwd = 3)
#
lam_pexm <- apply(aux_M2_pexm[, lnames], 2, mean)
lam_pois <- apply(aux_M2_pois[, lnames], 2, mean)
a <- NULL; for(j in 1:m){ a[j] <- 562 * (j-1) / m }; a <- c(a, 562);
rates_pexm <- stepf(lam_pexm, a)
rates_pois <- stepf(lam_pois, a)
plot(rates_pois$x, rates_pois$y, xlab = "time", ylab = "baseline hazard function", col = co1, type = "l", 
     main = expression("Model II (mean)"), ylim = c(0, ma), lwd = 7, cex.lab = 2, cex.axis = 2, cex.main = 2)
lines(rates_pexm$x, rates_pexm$y, col = co2, lwd = 3)
#
lam_pexm <- apply(aux_M2_pexm[, lnames], 2, median)
lam_pois <- apply(aux_M2_pois[, lnames], 2, median)
a <- NULL; for(j in 1:m){ a[j] <- 562 * (j-1) / m }; a <- c(a, 562);
rates_pexm <- stepf(lam_pexm, a)
rates_pois <- stepf(lam_pois, a)
plot(rates_pois$x, rates_pois$y, xlab = "time", ylab = "baseline hazard function", col = co1, 
     main = expression("Model II (median)"), type = "l", ylim = c(0, ma), lwd = 7, cex.lab = 2, cex.axis = 2, cex.main = 2)
lines(rates_pexm$x, rates_pexm$y, col = co2, lwd = 3)


##############################################################
# Code to reproduce Table 2 (Section 4).
# Standard deviations for each lambda in the PE distribution.
##############################################################
lnames <- paste0("lambda[",1:m,"]")
aux_pexm <- as.matrix(output_M1_pexm)
aux_pois = as.matrix(output_M1_pois)
lam_M1_pexm <- apply(aux_pexm[, lnames], 2, sd)
lam_M1_pois <- apply(aux_pois[, lnames], 2, sd)
aux_pexm <- as.matrix(output_M2_pexm)
aux_pois = as.matrix(output_M2_pois)
lam_M2_pexm <- apply(aux_pexm[, lnames], 2, sd)
lam_M2_pois <- apply(aux_pois[, lnames], 2, sd)
tab <- round(rbind(lam_M1_pexm, lam_M1_pois, lam_M2_pexm, lam_M2_pois), 4)
rownames(tab) <- c("I Module", "I Poisson", "II Module", "II Poisson")
round(tab, 4)