# Code for fitting mediation models to P3b latency data from our caffeine study. Data were pre-cleaned, as described in the manuscript, using Brain Vision Analyzer 2. The files which were read into this script were average ERPs (to compute the PCA weights), and then cleaned and segmented stimulus-locked data, plus associated RT data from the .MRK files produced by Brain Vision Recorder.

library(lme4)
library(mediation)
library(psych)
library(ggplot2)
library(GPArotation)


# Read in and concatenate average ERPs

'INSERT FILE PATH - AVE ERPs' -> AvesFolder

setwd(AvesFolder)
dir() -> Files
colnames(read.table(Files[1], head=T)) -> ElectrodeNames
data.frame(matrix(ncol=length(ElectrodeNames))) -> Dat
colnames(Dat) <- sort(ElectrodeNames)
for(i in 1:length(Files)){
read.table(Files[i], head=T) -> File
File[,order(names(File))] -> File
rbind(Dat, File) -> Dat
print(Files[i])
}
Dat<- Dat[-1,]
setwd(thisFolder)

colnames(File)
colnames(Dat)
nrow(Dat)
str(Dat)

#Run PCA on concatenated averages and store weightings

fa.parallel(Dat) 

prcomp(Dat, scale = TRUE, tol = .18) -> PCAdat

infomaxQ(PCAdat$rotation) -> PCAinf

as.matrix(Dat)%*%PCAinf$loadings[,1] -> Fac1
as.matrix(Dat)%*%PCAinf$loadings[,2] -> Fac2
as.matrix(Dat)%*%PCAinf$loadings[,3] -> Fac3
as.matrix(Dat)%*%PCAinf$loadings[,4] -> Fac4
as.matrix(Dat)%*%PCAinf$loadings[,5] -> Fac5
as.matrix(Dat)%*%PCAinf$loadings[,6] -> Fac6

# Read in segmented data, convolve with first PCA component and pick peaks. This for loop runs through each file one by one

'INSERT FILE PATH - SEGMENTED DATA' -> SegFolder
'INSERT FILE PATH - .MRK FILES' -> RTFolder

#As testing was double blind, file names don't have condition in name so this list links files to the appropriate condition.
c('Placebo', 'Caffeine', 'Caffeine', 'Placebo','Placebo', 'Caffeine', 'Caffeine', 'Placebo', 'Caffeine', 'Placebo','Placebo', 'Caffeine','Placebo', 'Caffeine',
'Placebo', 'Caffeine','Placebo', 'Caffeine', 'Caffeine', 'Placebo', 'Caffeine', 'Placebo', 'Caffeine', 'Placebo', 'Caffeine', 'Placebo', 'Placebo', 'Caffeine', 'Caffeine', 'Placebo') -> Condition

paste('P', c(2,2,3,3,4,4,5,5,6,6,7,7,9,9,10,10,11,11,12,12,13,13,14,14,15,15,16,16, 18, 18), sep ='') -> PN

Peaks <- data.frame()

setwd(SegFolder)
dir() -> Files
colnames(read.table(Files[1], head=T)) -> ElectrodeNames
data.frame(matrix(ncol=length(ElectrodeNames))) -> Dat
colnames(Dat) <- sort(ElectrodeNames)
setwd(RTFolder)
dir() -> FilesRT
setwd(thisFolder)

for(i in 1:(length(Files))){

setwd(SegFolder)
read.table(Files[i], head=T) -> File
File[,order(names(File))] -> File
print(Files[i])

setwd(RTFolder)

RTs <- vector()
for(j in 1:nrow(File)/1200) {
RTs[j] <- as.numeric(tail(scan(file= FilesRT[i] ,sep="\t",skip=j,nlines=1,fill=T,quiet=T,what="numeric"), 1)[1])
}

as.matrix(File)%*%PCAinf$loadings[,1] -> Fac1

Segs <- matrix(nrow=length(Fac1)/1200, ncol=1200)
for(j in seq(1,length(Fac1)/1200)){
Segs[j,] <-  Fac1[seq(((j-1)*1200)+1,((j-1)*1200)+1200)]
}

Latencies <- matrix(nrow=1, ncol=length(Fac1)/1200)
Amplitudes <- matrix(nrow=1, ncol=length(Fac1)/1200)

RLLatencies <- matrix(nrow=1, ncol=length(Fac1)/1200)
RLAmplitudes <- matrix(nrow=1, ncol=length(Fac1)/1200)



for(k in seq(1,length(Fac1)/1200)){
rbind(seq(-600,1798, 2), Segs[k,]) -> Trial

Window <- Trial[,Trial[1,]>249&Trial[1,]<751]
Window[1, Window[2,]== max(Window[2,])][1]-> Latencies[k]
Window[2, Window[2,]== max(Window[2,])][1]- mean(Trial[2,Trial[1,]> -601 & Trial[1,]< -399])-> Amplitudes[k]

RLWindow <- Trial[,Trial[1,]>(RTs[k]-251) &Trial[1,]<(RTs[k]+251)] 
RLWindow[1, RLWindow[2,]== max(RLWindow[2,])][1]-RTs[k] -> RLLatencies[k]
RLWindow[2, RLWindow[2,]== max(RLWindow[2,])][1] - mean(Trial[2,Trial[1,]> -601 & Trial[1,]< -399]) -> RLAmplitudes[k]

}


data.frame(cbind(t(Latencies), t(Amplitudes), t(RLLatencies), t(RLAmplitudes))) -> PeaksFile
PeaksFile$RTs <- RTs
Files[i] -> PeaksFile$File
Condition[i] -> PeaksFile$Condition
PN[i] -> PeaksFile$PN
1:length(Latencies) -> PeaksFile$Trial

rbind(Peaks, PeaksFile) -> Peaks
}
setwd(thisFolder)

colnames(Peaks) <- c('P3bLs','P3bAs','RLP3bLs','RLP3bAs', 'RTs', 'File', 'Condition', 'PN', 'Trial')

# Exclude peaks at edges of windows or pre-stimulus-onset. Compute inferred peaks

Peaks[Peaks$P3bLs!= 250 & Peaks$P3bLs!=750 & Peaks$RLP3bLs!=-250&Peaks$RLP3bLs!=250 ,] -> PeaksT
PeaksT[PeaksT$RLP3bLs+ PeaksT$RTs>0,] -> PeaksT
PeaksT$RLP3bLs + PeaksT$RTs -> PeaksT$P3bLsInf
PeaksT$P3bLs - PeaksT$RTs -> PeaksT$RLP3bLsInf

#Centre and scale latencies. Exclude trials with RT of >1000ms. Create numeric version of condition variable 
for(i in unique(PeaksT$PN)){
PeaksT $P3bLsC[PeaksT $PN==i] <- scale(PeaksT $P3bLs[PeaksT $PN==i])
PeaksT $RLP3bLsC[PeaksT $PN==i] <- scale(PeaksT $RLP3bLs[PeaksT $PN==i])
PeaksT $P3bAsC[PeaksT $PN==i] <- scale(PeaksT $P3bAs[PeaksT $PN==i])
PeaksT $RLP3bAsC[PeaksT $PN==i] <- scale(PeaksT $RLP3bAs[PeaksT $PN==i])

PeaksT $P3bLsInfC[PeaksT $PN==i] <- scale(PeaksT $P3bLsInf[PeaksT $PN==i])
PeaksT $RLP3bLsInfC[PeaksT $PN==i] <- scale(PeaksT $RLP3bLsInf[PeaksT $PN==i])
}
PeaksT[PeaksT$RTs<1000,] -> PeaksT1
2-as.numeric(PeaksT1 $Condition) -> PeaksT1 $CondNum



#Make Fig. 2
tiff("CaffeineFig2.tif", , width = 8, height = 8, units = 'in', res = 600, compression = 'lzw', type='cairo')
BGCol <- 'white'
ColourScheme <- c('firebrick', 'slateblue')
alph <- .4

ggplot(aes(RTs, fill = Condition, colour = Condition), data = PeaksT) + geom_density(alpha = alph) + scale_colour_manual(values=ColourScheme) + labs(x='RTs', y = 'Probability density') + theme(legend.position=c(.85, .6), legend.background = element_rect(colour = 'grey')) + xlim(200,800) -> DensityRT
ggplot(aes(P3bLs, fill = Condition, colour = Condition), data = PeaksT) + geom_density(alpha = alph) + scale_colour_manual(values=ColourScheme) + labs(x='Stimulus-locked P3b latencies', y = 'Probability density') + theme(legend.position='none') + xlim(200,800) -> DensitySL
ggplot(aes(RLP3bLs, fill = Condition, colour = Condition), data = PeaksT) + geom_density(alpha = alph) + scale_colour_manual(values=ColourScheme) + labs(x='Response-locked P3b latencies', y = 'Probability density') + theme(legend.position='none') + xlim(-300,300)  -> DensityRL
multiplot(DensityRT, DensitySL, DensityRL, cols = 1)
dev.off()


#Fit models to check mediation assumptions
lmer(P3bLsC ~ CondNum + (0+CondNum |PN), REML =F, data = PeaksT1) -> a.fit
lmer(P3bLsC ~ (0+CondNum |PN), REML =F, data = PeaksT1) -> a.null
summary(a.fit)
AIC(a.fit, a.null)
anova(a.fit, a.null)
ggplot(aes(Condition, P3bLs), data = PeaksT1) + geom_boxplot(aes(fill = Condition), outlier.colour = 'grey65') + ylab('Stimulus-locked P3b latencies') + scale_fill_manual(values = ColourScheme) + theme(panel.background = element_rect(fill = BGCol), legend.key= element_rect(fill = 'white')) -> a.plot


lmer(RTs ~ P3bLsC + (1+ P3bLsC |PN), REML =F, data = PeaksT1) -> b.fit
lmer(RTs ~ (1+ P3bLsC |PN), REML =F, data = PeaksT1) -> b.null
summary(b.fit)
AIC(b.fit, b.null)
anova(b.fit, b.null)
r.squaredGLMM(b.fit)
ggplot(aes(RTs, P3bLs, colour = Condition), data = PeaksT1) + geom_point(alpha = alph) + geom_smooth(method='lm', se = 0) + ylab('Stimulus-locked P3b latencies') + scale_colour_manual(values = ColourScheme) + theme(panel.background = element_rect(fill = BGCol), legend.key= element_rect(fill = 'white')) -> b.plot


lmer(RLP3bLsC ~ CondNum + (0+CondNum |PN), REML =F, data = PeaksT1) -> a.fit.r
lmer(RLP3bLsC ~ (0+CondNum |PN), REML =F, data = PeaksT1) -> a.null.r
summary(a.fit.r)
AIC(a.fit.r, a.null.r)
anova(a.fit.r, a.null.r)
r.squaredGLMM(a.fit.r)
ggplot(aes(Condition, RLP3bLs), data = PeaksT1) + geom_boxplot(aes(fill = Condition), outlier.colour = 'grey65')  + ylab('Response-locked P3b latencies') + scale_fill_manual(values = ColourScheme) + theme(panel.background = element_rect(fill = BGCol), legend.key= element_rect(fill = 'white')) -> a.r.plot



lmer(RTs ~ RLP3bLsC + (1+ P3bLsC |PN), REML =F, data = PeaksT1) -> b.fit.r
lmer(RTs ~ (1+ RLP3bLsC |PN), REML =F, data = PeaksT1) -> b.null.r
summary(b.fit.r)
AIC(b.fit.r, b.null.r)
anova(b.fit.r, b.null.r)
r.squaredGLMM(b.fit.r)
ggplot(aes(RTs, RLP3bLs, colour = Condition), data = PeaksT1) + geom_point(alpha = alph) + geom_smooth(method='lm', se = 0)  + ylab('Response-locked P3b latencies') + scale_colour_manual(values = ColourScheme) + theme(panel.background = element_rect(fill = BGCol), legend.key= element_rect(fill = 'white')) -> b.r.plot

#Make Figure 3
tiff("CaffeineFig3.tif", , width = 8, height = 8, units = 'in', res = 600, compression = 'lzw', type='cairo')
multiplot(a.plot, b.plot, a.r.plot, b.r.plot, cols = 2)
dev.off()

# Run mediation on stimulus-locked data
lmer(RTs ~ P3bLsC + CondNum + (1+ (P3bLsC |PN)), REML =F, data = PeaksT1) -> out.fit
lmer(P3bLsC ~ CondNum + (1|PN), REML =F, data = PeaksT1) -> med.fit
summary(out.fit)
r.squaredGLMM(out.fit)
summary(med.fit)
r.squaredGLMM(med.fit)

Med.model <- mediate(med.fit, out.fit, mediator = 'P3bLsC', treat = 'CondNum', sims = 10000, group.out = 'PN', data = PeaksT1)
summary(Med.model)
plot(Med.model)

lmer(RTs ~ P3bLs + CondNum + (1+ (P3bLs |PN)), REML =F, data = PeaksT1) -> out.fit
lmer(P3bLs ~ CondNum + (1|PN), REML =F, data = PeaksT1) -> med.fit
summary(out.fit)
r.squaredGLMM(out.fit)
summary(med.fit)
r.squaredGLMM(med.fit)
Med.model <- mediate(med.fit, out.fit, mediator = 'P3bLs', treat = 'CondNum', sims = 10000, group.out = 'PN', data = PeaksT1)
summary(Med.model)

#Run mediation on response-locked data - not strictly required given assmption not met, but I checked it out of interest (no evidence of mediation). 
lmer(RTs ~ RLP3bLsC + CondNum + (1+RLP3bLsC|PN), REML =F, data = PeaksT1) -> out.fit.r
lmer(RLP3bLsC ~ CondNum + (1|PN), REML =F, data = PeaksT1) -> med.fit.r
summary(out.fit.r)
r.squaredGLMM(out.fit.r)
summary(med.fit.r)
r.squaredGLMM(med.fit.r)

Med.model.r <- mediate(med.fit.r, out.fit.r, mediator = 'RLP3bLsC', treat = 'CondNum', sims = 10000, group.out = "PN", data = PeaksT1)
summary(Med.model.r)


#Make Figure 4
ColScheme <- c('grey20', 'firebrick2')
openplotmat()
straightarrow(from = c(.1,.15), to = c(.9,.15), , lty = 1, lwd = 1, lcol = ColScheme[1], endhead=F)
straightarrow(from = c(.1,.15), to = c(.5,.85), , lty = 1, lwd = 1, lcol = ColScheme[1])
straightarrow(from = c(.5,.85), to = c(.9,.15), , lty = 1, lwd = 1, lcol = ColScheme[1])
curvedarrow(from = c(.1,.15), to = c(.9,.15), curve = -0.7, lty = 1, lwd = Med.model$n0*8, lcol = ColScheme[2])
curvedarrow(from = c(.1,.15), to = c(.9,.15), curve = -0.1, lty = 1, lwd =(1-Med.model$n0)*8, lcol = ColScheme[2])
textplain(c(.1,.12), .1, 'Condition', col = ColScheme[1], font = 2)
textplain(c(.5,.9), .1, 'Stimulus-locked P3b latency', col = ColScheme[1, font = 2])
textplain(c(.9,.12), .1, 'Reaction time', col = ColScheme[1], font = 2)
textplain(c(.5,.65), .1, paste('ACME =', round(Med.model$d0, 2)), col = ColScheme[2], font = 2)
textplain(c(.5,.18), .1, paste('ADE =', round(Med.model$z0, 2)), col = ColScheme[2], font = 2)
textplain(c(.73,.49), .1, paste('Marginal R squared =', round(r.squaredGLMM(b.fit)[1],4)), col = ColScheme[1], srt = -57)
textplain(c(.27,.49), .1, paste('Marginal R squared =', round(r.squaredGLMM(a.fit)[1],4)), col = ColScheme[1], srt = 57)
textplain(c(.5,.12), .1, paste('Marginal R squared =', round(r.squaredGLMM(ab.fit)[1],4)), col = ColScheme[1])
textplain(c(.5,.08), .1, paste('Total effect =', round(Med.model$tau.coef, 2)), col = ColScheme[2])

tiff("CaffeineFig4.tif", , width = 8, height = 8, units = 'in', res = 600, compression = 'lzw', type='cairo')
ColScheme <- c('grey20', 'firebrick2')
openplotmat()
polygon(x = c(.14, .48, .82), c(.1,.75, .1), col = adjustcolor('navy', alpha.f = .2), border = NA)
curvedarrow(from = c(.08,.15), to = c(.9,.15), curve = -0.75, lty = 1, lwd = Med.model$n0*9, lcol = ColScheme[2])
curvedarrow(from = c(.08,.15), to = c(.9,.15), curve = -0.1, lty = 1, lwd =(1-Med.model$n0)*9, lcol = ColScheme[2])
textplain(c(.08,.12), .1, 'Condition', col = ColScheme[1], font = 2)
textplain(c(.48,.86), .1, 'Stimulus-locked P3b latency', col = ColScheme[1], font = 2)
textplain(c(.91,.12), .1, 'Reaction time', col = ColScheme[1], font = 2)
textplain(c(.48,.8), .1, paste('ACME =', round(Med.model$d0, 2)), col = ColScheme[2], font = 2)
#textplain(c(.48,.53), .1, paste('(',round(Med.model$n0, 2)*100, '%', ')', sep = ''), col = ColScheme[2])
textplain(c(.48,.28), .1, paste('ADE =', round(Med.model$z0, 2)), col = ColScheme[2], font = 2)
textplain(c(.48,.16), .1, paste('Total effect =', round(Med.model$tau.coef, 2)), col = ColScheme[2], font = 2)
textplain(c(.68,.40), .1, paste('Marginal R squared =', round(r.squaredGLMM(b.fit)[1],4)), col = ColScheme[1], srt = -60, font = 3)
textplain(c(.28,.40), .1, paste('Marginal R squared =', round(r.squaredGLMM(a.fit)[1],4)), col = ColScheme[1], srt = 60, font = 3)
textplain(c(.48,.08), .1, paste('Marginal R squared =', round(r.squaredGLMM(ab.fit)[1],4)), col = ColScheme[1], font = 3)
dev.off()

#Control analyses
# Check if caffeine affects amplitude 
lmer(RLP3bAs ~ CondNum + (CondNum |PN), REML =F, data = PeaksT1) -> Amp.model
lmer(RLP3bAs ~ (CondNum |PN), REML =F, data = PeaksT1) -> Amp.null
summary(Amp.model)
summary(Amp.null)
anova(Amp.model, Amp.null)

# Check results hold for inferred latencies
lmer(P3bLsInfC ~ CondNum + (0+CondNum |PN), REML =F, data = PeaksT1) -> a.fit
lmer(P3bLsInfC ~ (0+CondNum |PN), REML =F, data = PeaksT1) -> a.null
summary(a.fit)
r.squaredGLMM(a.fit)
lmer(RTs ~ P3bLsInfC + (1+ P3bLsC |PN), REML =F, data = PeaksT1) -> b.fit
lmer(RTs ~ (1+ P3bLsInfC |PN), REML =F, data = PeaksT1) -> b.null
summary(b.fit)
AIC(b.fit, b.null)
anova(b.fit, b.null)
r.squaredGLMM(b.fit)

lmer(RTs ~ P3bLsInfC + CondNum + (1+ (P3bLsInfC |PN)), REML =F, data = PeaksT1) -> out.fit
lmer(P3bLsInfC ~ CondNum + (1|PN), REML =F, data = PeaksT1) -> med.fit
summary(out.fit)
r.squaredGLMM(out.fit)
summary(med.fit)
r.squaredGLMM(med.fit)

Med.model.Inf <- mediate(med.fit, out.fit, mediator = 'P3bLsInfC', treat = 'CondNum', sims = 10000, group.out = 'PN', data = PeaksT1)
summary(Med.model.Inf)


lmer(RLP3bLsInfC ~ CondNum + (0+CondNum |PN), REML =F, data = PeaksT1) -> a.fit.r
lmer(RLP3bLsInfC ~ (0+CondNum |PN), REML =F, data = PeaksT1) -> a.null.r
summary(a.fit.r)
r.squaredGLMM(a.fit.r)
lmer(RTs ~ RLP3bLsInfC + (1+ RLP3bLsInfC |PN), REML =F, data = PeaksT1) -> b.fit.r
lmer(RTs ~ (1+ RLP3bLsC |PN), REML =F, data = PeaksT1) -> b.null.r
summary(b.fit.r)
AIC(b.fit.r, b.null.r)
anova(b.fit.r, b.null.r)
r.squaredGLMM(b.fit.r)


lmer(RTs ~ RLP3bLsInfC + CondNum + (1+ (RLP3bLsInfC |PN)), REML =F, data = PeaksT1) -> out.fit
lmer(RLP3bLsInfC ~ CondNum + (1|PN), REML =F, data = PeaksT1) -> med.fit
summary(out.fit)
r.squaredGLMM(out.fit)
summary(med.fit)
r.squaredGLMM(med.fit)

Med.model.Inf <- mediate(med.fit, out.fit, mediator = 'RLP3bLsInfC', treat = 'CondNum', sims = 10000, group.out = 'PN', data = PeaksT1)
summary(Med.model.Inf)

#Note - code for the Pz analysis not included here - I just used the Pz electrode rather than 'Fac1' for peak picking and ran the rest of the script as normal.