## ---------------------------------------------------
## Replication script for the manuscript:
##
## Robust Mediation Analysis: The R Package robmed
## Andreas Alfons, Nufer Y. Ates, Patrick J.F. Groenen
## ---------------------------------------------------


## Set command prompt as in manuscript
# options(prompt = "R> ", continue = "+  ", width = 75, useFancyQuotes = FALSE)


## -----------------------
## Section 1: Introduction
## -----------------------

## Page 3: Package installation.  This is commented out to prevent that
## anything is installed without the user's knowledge.

# install.packages("robmed")


## ----------------------
## Section 2: Methodology
## ----------------------

## Page 11: Produce Figure 5.  This figure shows the loss functions and weight
## functions of OLS regression and the robust MM-estimator.  The code is not
## shown in the manuscript as it is not relevant for the use of package robmed.

## load packages
library("dplyr")
library("ggplot2")
library("robmed")
library("scales")

## control parameters
tuning.psi <- reg_control()$tuning.psi
col <- hue_pal()(4)[c(1, 3)]
lty <- c(2, 1)
line_size <- 2/3

## grid of x-values
x <- seq(-4, 4, length.out = 100)

## labels to be used for methods and panels
methods <- c("OLS", "MM")
functions <- c("Loss", "Weight")

## loss functions
OLS_loss <- data.frame(x = x,
                       y = x^2,
                       Method = factor(methods[1], levels = methods),
                       Function = factor(functions[1], levels = functions))
robust_loss <- data.frame(x = x,
                          y = Mpsi(x, tuning.psi, "bisquare", deriv = -1),
                          Method = factor(methods[2], levels = methods),
                          Function = factor(functions[1], levels = functions))

## weight functions
OLS_weight <- data.frame(x = x,
                         y = 1,
                         Method = factor(methods[1], levels = methods),
                         Function = factor(functions[2], levels = functions))
robust_weight <- data.frame(x = x,
                            y = Mwgt(x, tuning.psi, "bisquare"),
                            Method = factor(methods[2], levels = methods),
                            Function = factor(functions[2], levels = functions))

## data frame for plot
df <- rbind(OLS_loss, robust_loss, OLS_weight, robust_weight) %>% filter(y <= 4)

## create plot
ggplot() +
  geom_line(aes(x = x, y = y, color = Method, linetype = Method),
            data = df, size = line_size) +
  labs(x = NULL, y = NULL) +
  scale_linetype_manual(values = lty) +
  facet_wrap(~ Function, scales = "free_y") +
  theme_bw() + theme(axis.text = element_text(size = 12),
                     axis.title = element_text(size = 13),
                     legend.key = element_rect(color = NA),
                     legend.key.width = unit(1.5, "line"),
                     legend.text = element_text(size = 12),
                     legend.title = element_text(size = 13),
                     strip.text = element_text(size = 12))


## -------------------------
## Section 3: Implementation
## -------------------------

## Page 12: Load package and data for use in code examples
library("robmed")
data("BSG2014", package = "robmed")


## -------------------------
## Section 3.1: Example data
## -------------------------

## Page 14: Summary of the variables used in the case studies
keep <- c("ValueDiversity", "TaskConflict", "TeamCommitment", "TeamScore",
          "SharedLeadership", "AgeDiversity", "GenderDiversity",
          "ProceduralJustice", "InteractionalJustice", "TeamPerformance")
summary(BSG2014[, keep])


## ------------------------------
## Section 3.2: Formula interface
## ------------------------------

## Page 15: Example of the formula interface for a simple mediation model
TeamCommitment ~ m(TaskConflict) + ValueDiversity

## Page 15: Example of the formula interface for a serial multiple mediator
## model.  The serial mediators are in consecutive order from left to right.
TeamScore ~ serial_m(TaskConflict, TeamCommitment) + ValueDiversity

## Page 15: Example of the formula interface for a parallel multiple mediator
## model with control variables
TeamPerformance ~ parallel_m(ProceduralJustice, InteractionalJustice) +
  SharedLeadership + covariates(AgeDiversity, GenderDiversity)


## ----------------------------------------------
## Section 4: Illustrations: Using package robmed
## ----------------------------------------------

## Page 16: Seed to be used for the random number generator
seed <- 20211117


## -----------------------------
## Section 4.1: Simple mediation
## -----------------------------

## Page 17: Store formula object for later use
f_simple <- TeamCommitment ~ m(TaskConflict) + ValueDiversity

## Page 17: Perform robust bootstrap test and OLS bootstrap test
set.seed(seed)
robust_boot_simple <- test_mediation(f_simple, data = BSG2014,
                                     robust = TRUE)
set.seed(seed)
ols_boot_simple <- test_mediation(f_simple, data = BSG2014,
                                  robust = FALSE)

## Page 17: Show summary of results from the robust bootstrap test, which also
## shows the diagnostic plot in Figure 6
summary(robust_boot_simple)

## Page 20: The diagnostic plot in Figure 6 can also be produced with the
## commands below
weight_plot(robust_boot_simple) +
  scale_color_manual("", values = c("black", "#00BFC4")) +
  theme(legend.position = "top")

## Page 20: Show summary of results from the OLS bootstrap test.  In this
## example, estimates on the original data and the usual normal-theory t tests
## are reported for the effects other than the indirect effect, similar to the
## PROCESS macro.
summary(ols_boot_simple, type = "data")

## Page 22: The usual methods are available to extract effect estimates and
## confidence intervals
coef(robust_boot_simple)
confint(robust_boot_simple)

## Page 22: They can also be used to extract estimates and confidence intervals
## based on the original data rather than based on the bootstrap distribution
coef(ols_boot_simple, type = "data")
confint(ols_boot_simple, type = "data")

## Page 22/23: They further allow to select only estimates and confidence
## intervals for specific effects
coef(robust_boot_simple, parm = "Indirect")
confint(robust_boot_simple, parm = "Indirect")

## Page 23: p values of the indirect effect can be computed as the smallest
## significance level alpha for which the (1-alpha)*100% confidence interval
## does not contain 0
p_value(robust_boot_simple, parm = "Indirect")
p_value(ols_boot_simple, parm = "Indirect")

## Page 23: Combine results of the OLS bootstrap test and the robust bootstrap
## test into a list to compare them with several plots
boot_list <- list("OLS bootstrap" = ols_boot_simple,
                  "ROBMED" = robust_boot_simple)

## Page 23: Produce Figure 7 showing density estimates of the indirect effect
density_plot(boot_list)

## Page 23: Produce Figure 8 visualizing point estimates and confidence intervals
ci_plot(boot_list, parm = c("a", "b", "Direct", "Indirect"))


## Page 25: Produce Figure 9 containing a diagnostic plot with tolerance ellipses
ellipse_plot(boot_list, horizontal = "ValueDiversity",
             vertical = "TaskConflict")

## Page 26: Example of further customization of plots shown in Figure 10
setup <- setup_ellipse_plot(boot_list, horizontal = "ValueDiversity",
                            vertical = "TaskConflict")
ggplot() +
  geom_path(aes(x = x, y = y, color = Method), data = setup$ellipse) +
  geom_point(aes(x = x, y = y, fill = Weight), data = setup$data,
             shape = 21) +
  scale_fill_gradient(limits = 0:1, low = "white", high = "black") +
  labs(x = setup$horizontal, y = setup$vertical)


## --------------------------------------
## Section 4.2: Serial multiple mediators
## --------------------------------------

## Page 27: Store formula object for later use
f_serial <- TeamScore ~ serial_m(TaskConflict, TeamCommitment) +
  ValueDiversity

## Page 27: Perform robust bootstrap test and OLS bootstrap test
set.seed(seed)
robust_boot_serial <- test_mediation(f_serial, data = BSG2014,
                                     robust = TRUE)
set.seed(seed)
ols_boot_serial <- test_mediation(f_serial, data = BSG2014,
                                  robust = FALSE)

## Page 27/28: Print objects to show only results for indirect effects
robust_boot_serial
ols_boot_serial

## Page 28: Produce Figure 11 containing the diagnostic plot of the robust
## bootstrap test
weight_plot(robust_boot_serial) +
  scale_color_manual("", values = c("black", "#00BFC4")) +
  theme(legend.position = "top")


## --------------------------------------------------------------
## Section 4.3: Parallel multiple mediators and control variables
## --------------------------------------------------------------

## Page 29: Store formula object for later use
f_parallel <- TeamPerformance ~
  parallel_m(ProceduralJustice, InteractionalJustice) +
  SharedLeadership + covariates(AgeDiversity, GenderDiversity)

## Page 29: Perform robust bootstrap test and OLS bootstrap test
set.seed(seed)
robust_boot_parallel <- test_mediation(f_parallel, data = BSG2014,
                                       robust = TRUE)
set.seed(seed)
ols_boot_parallel <- test_mediation(f_parallel, data = BSG2014,
                                    robust = FALSE)

## Page 30: Print objects to show only results for indirect effects
robust_boot_parallel
ols_boot_parallel

## Page 31: Produce Figure 12 containing the diagnostic plot with a tolerance
## ellipse, in this example with partial residuals on the vertical axis
ellipse_plot(robust_boot_parallel, horizontal = "SharedLeadership",
             vertical = "TeamPerformance", partial = TRUE)

## Page 31: Perform the robust bootstrap test and include a test on whether the
## indirect effects differ in magnitude
set.seed(seed)
test_mediation(f_parallel, data = BSG2014, contrast = "absolute")

## Page 32: It is not actually necessary to run the entire bootstrap procedure
## again.  Instead, the existing object can be reanalyzed.
retest(robust_boot_parallel, contrast = "absolute")


## ---------------------------------------------------------------
## Appendix B: Additional summary output of robust bootstrap tests
## ---------------------------------------------------------------

## Page 40: Show the full summary of the robust bootstrap test for the serial
## multiple mediator model from Section 4.2, but suppress the diagnostic plot
## (since the plot has already been shown in Figure 11 in Section 4.2).
summary(robust_boot_serial, plot = FALSE)

## Page 41: Show the full summary of the robust bootstrap test for the serial
## multiple mediator model from Section 4.3, which also creates the diagnostic
## plot in Figure 13
summary(robust_boot_parallel)