# code for "Fixed and Random effects: making an informed choice"
# Andrew Bell, Malcolm Fairbrother, Kelvyn Jones.
# Affiliation: University of Sheffield, andrew.j.d.bell@sheffield.ac.uk

library(lme4)
library(plm) # for clustered standard errors
library(lmtest)
library(ggplot2)
library(reshape2)
library(gridExtra)

c.mer <- function(mod) { # function to extract FEs, SEs, RE variances, and REs from a fitted mer object
	out <- c(fixef(mod), vcov(mod)@factors$correlation@sd, unlist(ranef(mod)$grp[,1]))
	out <- c(out, c(unlist(lapply(VarCorr(mod), diag))))
	if(length(c(unlist(lapply(VarCorr(mod), diag))))<2) out <- c(out, 1) # pad this to save time (symmetry) later
	c(out, attr(VarCorr(mod), "sc")^2, 1)
	} # intercept, slope, SE on intercept, SE on slope, 60/30 REs, 1/2 RE variance(s), residual variance (if exists), pad

c.FE <- function(mod) { # function to extract FEs, SEs, robust SEs from/for a fitted FE model
	out <- c(coef(summary(mod))[1:2,1:2]) # 2 FEs and 2 SEs
	out <- c(out, 1, coef(summary(mod))[3:nrow(coef(summary(mod))),1]) # pad, then 29/59 dummies
	out <- c(out, 1, 1, summary(mod)$sigma) # pads, then residual variance (if exists)
	if(!length(summary(mod)$sigma)) out <- c(out, 1) # pad this to save time (symmetry) later
	out <- c(out, as.matrix(coeftest(mod, vcov=vcovHC(mod, cluster="group")))[2,2]) # uses plm package for cluster-robust SE
	out
	} # intercept, slope, SE on intercept, SE on slope, pad, 29/59 FEs, pad, pad, residual variance (if exists), robust SE

dgp <- function(n_I, n_J, var_j0, var_j1, chi, outlier, unbalanced) {
	dat <- expand.grid(grp=1:n_J, ind=1:(n_I/n_J)) # create grp and ind indices
	dat$U0_j <- rnorm(n_J, mean=1, sd=sqrt(var_j0))[dat$grp]
	dat$U1_j <- rnorm(n_J, mean=1, sd=sqrt(var_j1))[dat$grp]
	if(chi) dat$U0_j <- 1+(rchisq(n_J, df=var_j0/2)-(var_j0/2))[dat$grp] 	# mean of 1, variance is var_j0
	if(outlier) dat$U0_j[abs(dat$U0_j)==max(abs(dat$U0_j))] <- 3*dat$U0_j[abs(dat$U0_j)==max(abs(dat$U0_j))]	# make an outlier
	sam <- 1:3
	while(length(unique(dat$U0_j[sam]))!=n_J) sam <- sample(1:nrow(dat), n_I, replace=T) 	
	if(unbalanced) dat <- dat[sam,]						# sample if unbalanced, while keeping n_J
	dat$X <- rnorm(n_I, sd=0.25)
	dat <- within(dat, y1 <- X + U0_j + rnorm(n_I))				# random intercepts only, Normal
	dat <- within(dat, y1bi <- rbinom(nrow(dat), 1, plogis(X + U0_j)))		# random intercepts only, logit
	dat <- within(dat, y2 <- U1_j*X + U0_j + rnorm(n_I))				# random intercepts and slopes, Normal
	dat <- within(dat, y2bi <- rbinom(nrow(dat), 1, plogis(U1_j*X + U0_j)))	# random intercepts and slopes, logit
	dat
	}

sim <- function(n_I, n_J, var_j0, var_j1, chi, outlier, unbalanced) { # function to generate data, fit models, return results
	dat <- dgp(n_I=n_I, n_J=n_J, var_j0=var_j0, var_j1=var_j1, chi=chi, outlier=outlier, unbalanced=unbalanced)
	RI1  <- lmer(y1 ~ X + (1 | grp), dat)
	RS1  <- lmer(y1 ~ X + (X | grp), dat)
	RI2  <- lmer(y2 ~ X + (1 | grp), dat)
	RS2  <- lmer(y2 ~ X + (X | grp), dat)
	LI1  <- glmer(y1bi ~ X + (1 | grp), family=binomial, dat)
	LS1  <- glmer(y1bi ~ X + (X | grp), family=binomial, dat)
	LI2  <- glmer(y2bi ~ X + (1 | grp), family=binomial, dat)
	LS2  <- glmer(y2bi ~ X + (X | grp), family=binomial, dat)
	FE1  <- lm(y1 ~ X + as.factor(grp), dat)
	FE2  <- lm(y2 ~ X + as.factor(grp), dat)
	LE1  <- glm(y1bi ~ X + as.factor(grp), family=binomial, dat)
	LE2  <- glm(y2bi ~ X + as.factor(grp), family=binomial, dat)
	cbind(sapply(list(RI1, RS1, RI2, RS2, LI1, LS1, LI2, LS2), c.mer), sapply(list(FE1, FE2, LE1, LE2), c.FE), c(1,1,1,1,unique(dat$U0_j),1,1,1,1)) # matrix 68 by 13
	}

# DGP conditions:
n_I <- 600
n_J <- c(60, 30)
var_j0 <- 4
var_j1 <- 4
chi <- c(F, T)
outlier <- c(F, T)
unbalanced <- c(F, T)

# combine conditions:

com <- expand.grid(n_I=n_I, n_J=n_J, var_j0=var_j0, var_j1=var_j1, chi=chi, outlier=outlier, unbalanced=unbalanced)
com <- com[apply(com[,5:7], 1, sum)<2,]
rownames(com) <- 1:nrow(com)

sims <- 1000 # set number of simulations per combination of conditions

options(warn=2) # to elevate warnings to errors

# run the simulations
d0 <- date()
set.seed(123)
res <- apply(com, 1, function(x) lapply(1:sims, function(y) {
	out <- 0
	class(out) <- "try-error"
	while(class(out)=="try-error") out <- try(do.call("sim", as.list(x)), silent=T)
	out
	} ) ) # a list of 8 lists of (sims) matrices of dimensions 68 by 13
d1 <- date()

# plot Optimism (Figure 1)

# first, using "non-robust" SEs
resout <- lapply(res, function(x) array(unlist(x), dim=c(dim(x[[1]]), length(x))))
opt <- data.frame(sapply(resout, function(X) apply(X[2,,], 1, function(x) sqrt(sum((x-mean(x))^2)))/apply(X[4,,], 1, function(x) sqrt(sum(x^2)))))[-13,]
names(opt) <- c("nor60", "nor30", "chi60", "chi30", "out60", "out30", "unb60", "unb30")
rownames(opt) <- c("RI1", "RS1", "RI2", "RS2", "LI1", "LS1", "LI2", "LS2", "FE1", "FE2", "LE1", "LE2")
opt <- opt[c("FE1", "LE1", "RI1", "LI1", "RS1", "LS1", "FE2", "LE2", "RI2", "LI2", "RS2", "LS2"),]

toplot <- melt(opt)
toplot$mod <- rep(rownames(opt), 8)
toplot$RS <- F
toplot$RS[grep("2", toplot$mod)] <- T
toplot$RS <- factor(toplot$RS)
levels(toplot$RS) <- c("No RSs in DGP", "RSs in DGP")
toplot$mod <- substr(toplot$mod, 1, 2)
toplot$bin <- F
toplot$bin[grep("L", toplot$mod)] <- T
toplot$J60 <- F
toplot$J60[grep("60", toplot$variable)] <- T
toplot$scen <- rep(c("Normal","Chi-Sq","Outlier","Unbal."), each=24)
toplot$scen <- factor(toplot$scen, levels=c("Normal", "Chi-Sq", "Outlier", "Unbal."), ordered=T)
toplot$EIS <- factor(substr(toplot$mod, 2, 2))
levels(toplot$EIS) <- c("FE Model", "RI Model", "RS Model", "FE (Robust) Model")

# now to incorporate robust SEs
resout <- lapply(res, function(x) array(unlist(x), dim=c(dim(x[[1]]), length(x))))
resout[[7]] <- resout[[7]][,,apply(apply(resout[[7]], 3, is.na), 2, function(x) !any(x))] # some rows had NAs for the robust SEs
Ropt <- data.frame(sapply(resout, function(X) apply(X[2,9:12,], 1, function(x) sqrt(sum((x-mean(x))^2)))/apply(X[nrow(X),9:12,], 1, function(x) sqrt(sum(x^2)))))
names(Ropt) <- c("nor60", "nor30", "chi60", "chi30", "out60", "out30", "unb60", "unb30")
rownames(Ropt) <- c("FE1", "FE2", "LE1", "LE2")
Ropt <- melt(Ropt)

Ropt$mod <- rep(c("FE1", "FE2", "LE1", "LE2"), 8)
Ropt$RS <- F
Ropt$RS[grep("2", Ropt$mod)] <- T
Ropt$RS <- factor(Ropt$RS)
levels(Ropt$RS) <- c("No RSs in DGP", "RSs in DGP")
Ropt$mod <- substr(Ropt$mod, 1, 2)
Ropt$bin <- F
Ropt$bin[grep("L", Ropt$mod)] <- T
Ropt$J60 <- F
Ropt$J60[grep("60", Ropt$variable)] <- T
Ropt$scen <- rep(c("Normal","Chi-Sq","Outlier","Unbal."), each=8)
Ropt$scen <- factor(Ropt$scen, levels=c("Normal", "Chi-Sq", "Outlier", "Unbal."), ordered=T)
Ropt$EIS <- "FE (Robust) Model"

toplot <- rbind(toplot, Ropt)

toplot$EIS <- factor(toplot$EIS, levels=c("FE Model", "FE (Robust) Model", "RI Model", "RS Model"))

ggplot(toplot, aes(x = scen, y = value, shape=bin)) + geom_hline(yintercept=1) + geom_point(aes(size=3, colour= J60), position = position_dodge(width = 0.25)) + facet_wrap(RS~EIS, ncol=4) + theme(legend.position="none") + labs(x="", y="Optimism") + scale_colour_manual(values = c('red','blue'))

ggsave("Figure1.pdf", width=9, height=6)

# plot Figure 2

means <- lapply(resout, function(x) apply(x, c(1,2), function(xx) mean(xx)))
Bmeans <- data.frame(sapply(means, function(x) x[2,-13]))
names(Bmeans) <- c("nor60", "nor30", "chi60", "chi30", "out60", "out30", "unb60", "unb30")
rownames(Bmeans) <- c("RI1", "RS1", "RI2", "RS2", "LI1", "LS1", "LI2", "LS2", "FE1", "FE2", "LE1", "LE2")
Bmeans <- Bmeans[c("FE1", "LE1", "RI1", "LI1", "RS1", "LS1", "FE2", "LE2", "RI2", "LI2", "RS2", "LS2"),]
toplot <- melt(Bmeans)
toplot$mod <- rep(rownames(Bmeans), 8)
toplot$RS <- F
toplot$RS[grep("2", toplot$mod)] <- T
toplot$RS <- factor(toplot$RS)
levels(toplot$RS) <- c("No Random Slopes in DGP", "Random Slopes in DGP")
toplot$mod <- substr(toplot$mod, 1, 2)
toplot$bin <- F
toplot$bin[grep("L", toplot$mod)] <- T
toplot$J60 <- F
toplot$J60[grep("60", toplot$variable)] <- T
toplot$scen <- rep(c("Normal","Chi-Sq","Outlier","Unbalanced"), each=24)
toplot$scen <- factor(toplot$scen, levels=c("Normal", "Chi-Sq", "Outlier", "Unbalanced"), ordered=T)
toplot$EIS <- factor(substr(toplot$mod, 2, 2))
levels(toplot$EIS) <- c("FE Model", "RI Model", "RS Model")
Btoplot <- toplot[toplot$RS=="No Random Slopes in DGP" & toplot$EIS=="RI Model",]
Btoplot$panel <- "Btoplot"

opt <- data.frame(sapply(resout, function(X) apply(X[2,,], 1, function(x) sqrt(sum((x-mean(x))^2)))/apply(X[4,,], 1, function(x) sqrt(sum(x^2)))))[-13,]
names(opt) <- c("nor60", "nor30", "chi60", "chi30", "out60", "out30", "unb60", "unb30")
rownames(opt) <- c("RI1", "RS1", "RI2", "RS2", "LI1", "LS1", "LI2", "LS2", "FE1", "FE2", "LE1", "LE2")
opt <- opt[c("FE1", "LE1", "RI1", "LI1", "RS1", "LS1", "FE2", "LE2", "RI2", "LI2", "RS2", "LS2"),]
toplot <- melt(opt)
toplot$mod <- rep(rownames(opt), 8)
toplot$RS <- F
toplot$RS[grep("2", toplot$mod)] <- T
toplot$RS <- factor(toplot$RS)
levels(toplot$RS) <- c("No Random Slopes in DGP", "Random Slopes in DGP")
toplot$mod <- substr(toplot$mod, 1, 2)
toplot$bin <- F
toplot$bin[grep("L", toplot$mod)] <- T
toplot$J60 <- F
toplot$J60[grep("60", toplot$variable)] <- T
toplot$scen <- rep(c("Normal","Chi-Sq","Outlier","Unbalanced"), each=24)
toplot$scen <- factor(toplot$scen, levels=c("Normal", "Chi-Sq", "Outlier", "Unbalanced"), ordered=T)
toplot$EIS <- factor(substr(toplot$mod, 2, 2))
levels(toplot$EIS) <- c("FE Model", "RI Model", "RS Model", "FE (Robust) Model")
toplot$EIS <- factor(toplot$EIS, levels=c("FE Model", "FE (Robust) Model", "RI Model", "RS Model"))
Otoplot <- toplot[toplot$RS=="No Random Slopes in DGP" & toplot$EIS=="RI Model",]
Otoplot$panel <- "Otoplot"

means <- lapply(resout, function(x) apply(x, c(1,2), function(xx) mean(xx)))
Smeans <- data.frame(sapply(means, function(x) x[nrow(x)-3,1:8]))
names(Smeans) <- c("nor60", "nor30", "chi60", "chi30", "out60", "out30", "unb60", "unb30")
rownames(Smeans) <- c("RI1", "RS1", "RI2", "RS2", "LI1", "LS1", "LI2", "LS2")
Smeans <- Smeans[c("RI1", "LI1", "RS1", "LS1", "RI2", "LI2", "RS2", "LS2"),]
toplot <- melt(Smeans)
toplot$mod <- rep(rownames(Smeans), 8)
toplot$RS <- F
toplot$RS[grep("2", toplot$mod)] <- T
toplot$RS <- factor(toplot$RS)
levels(toplot$RS) <- c("No Random Slopes in DGP", "Random Slopes in DGP")
toplot$mod <- substr(toplot$mod, 1, 2)
toplot$bin <- F
toplot$bin[grep("L", toplot$mod)] <- T
toplot$J60 <- F
toplot$J60[grep("60", toplot$variable)] <- T
toplot$scen <- rep(c("Normal","Chi-Sq","Outlier","Unbalanced"), each=16)
toplot$scen <- factor(toplot$scen, levels=c("Normal", "Chi-Sq", "Outlier", "Unbalanced"), ordered=T)
toplot$EIS <- factor(substr(toplot$mod, 2, 2))
levels(toplot$EIS) <- c("RI Model", "RS Model")
Stoplot <- toplot[toplot$RS=="No Random Slopes in DGP" & toplot$EIS=="RI Model",]
Stoplot$panel <- "Stoplot"

REs <- lapply(res, function(x) lapply(x, function(xx) {
	FE <- apply(xx[,9:12], 2, function(xxx) c(xxx[1], xxx+xxx[1])[c(1,7:(length(xxx)-3))])
	RE <- apply(xx[,1:8], 2, function(xxx) 1+xxx[5:(length(xxx)-4)])
	true <- matrix(xx[,13][5:(nrow(xx)-4)])
	cbind(RE, FE, true)
	} )) # a list of eight lists of 1000 matrices (each 60 by 13)
REs <- data.frame(sapply(REs, function(x) rowMeans(sapply(x, function(xx) sqrt(rowMeans(apply(xx, 1, function(xxx) sqrt((xxx[-13]-xxx[13])^2)))))))) # 12 by 8
names(REs) <- c("nor60", "nor30", "chi60", "chi30", "out60", "out30", "unb60", "unb30")
rownames(REs) <- c("RI1", "RS1", "RI2", "RS2", "LI1", "LS1", "LI2", "LS2", "FE1", "FE2", "LE1", "LE2")
toplot <- melt(REs)
toplot$mod <- rep(rownames(REs), 8)
toplot$RS <- F
toplot$RS[grep("2", toplot$mod)] <- T
toplot$RS <- factor(toplot$RS)
levels(toplot$RS) <- c("No Random Slopes in DGP", "Random Slopes in DGP")
toplot$mod <- substr(toplot$mod, 1, 2)
toplot$bin <- F
toplot$bin[grep("L", toplot$mod)] <- T
toplot$J60 <- F
toplot$J60[grep("60", toplot$variable)] <- T
toplot$scen <- rep(c("Normal","Chi-Sq","Outlier","Unbalanced"), each=24)
toplot$scen <- factor(toplot$scen, levels=c("Normal", "Chi-Sq", "Outlier", "Unbalanced"), ordered=T)
toplot$EIS <- factor(substr(toplot$mod, 2, 2))
levels(toplot$EIS) <- c("FE Model", "RI Model", "RS Model")
Utoplot <- toplot[toplot$RS=="No Random Slopes in DGP" & toplot$EIS=="RI Model",]
Utoplot$panel <- "Utoplot"

Uplot <- ggplot(Utoplot, aes(x = scen, y = value, shape=bin)) + geom_hline(yintercept=0) + geom_point(aes(size=3, colour= J60), position = position_dodge(width = 0.25)) + facet_wrap(RS~EIS) + theme(axis.text.y = element_text(size = 16), axis.title.y = element_text(size = 16), axis.text.x = element_text(size = 16), strip.text.x = element_text(size = 16), legend.position="none") + labs(x="", y="RMSE of the Random Intercepts") + scale_colour_manual(values = c('red','blue')) + coord_cartesian(ylim=c(0,max(Utoplot$value)))

Splot <- ggplot(Stoplot, aes(x = scen, y = value, shape=bin)) + geom_hline(yintercept=4) + geom_point(aes(size=3, colour= J60), position = position_dodge(width = 0.25)) + facet_wrap(RS~EIS) + theme(axis.text.y = element_text(size = 16), axis.title.y = element_text(size = 16), axis.text.x = element_text(size = 16), strip.text.x = element_text(size = 16), legend.position="none") + labs(x="", y="Mean Estimate of Level-2 Variance") + scale_colour_manual(values = c('red','blue'))

Oplot <- ggplot(Otoplot, aes(x = scen, y = value, shape=bin)) + geom_hline(yintercept=1) + geom_point(aes(size=3, colour= J60), position = position_dodge(width = 0.25)) + facet_wrap(RS~EIS, ncol=4) + theme(axis.text.y = element_text(size = 16), axis.title.y = element_text(size = 16), axis.text.x = element_text(size = 16), strip.text.x = element_text(size = 16), legend.position="none") + labs(x="", y="Optimism of the SEs") + scale_colour_manual(values = c('red','blue')) + coord_cartesian(ylim = c(0.9,1.1))

Bplot <- ggplot(Btoplot, aes(x = scen, y = value, shape=bin)) + geom_hline(yintercept=1) + geom_point(aes(size=3, colour= J60), position = position_dodge(width = 0.25)) + facet_wrap(RS~EIS) + theme(axis.text.y = element_text(size = 16), axis.title.y = element_text(size = 16), axis.text.x = element_text(size = 16), strip.text.x = element_text(size = 16), legend.position="none") + labs(x="", y="Mean Estimate of Beta") + scale_colour_manual(values = c('red','blue')) + coord_cartesian(ylim = c(0.9,1.1))

pdf("Figure2.pdf", width = 12, height = 12)
grid.arrange(Bplot, Oplot, Splot, Uplot, ncol=2)
dev.off()