### Single-Trial Linear Ballistic Accumulator 
###
### Leendert van Maanen, October, 15, 2011
### Comments and questions can be addressed to l.vanmaanen@uva.nl
### 
### Citation: Van Maanen, L., Brown, S.D., Eichele, T., Wagenmakers, E.J., 
###		Ho, T.C., Serences, J. T., & Forstmann, B. U. (2011). Neural 
###		correlates of trial-to-trial fluctuations in response-caution. 
###		Journal of Neuroscience, 31, 17488-17495.

## steps:
# 1. compute posterior distributions of the group level parameters (eg, using winBugs)
# 2. sample from those
# 3. draw single trial drift
# 4. compute single-trial start point
# 5. accept/reject based on a_i < 0: reject; a_i > A: reject; otherwise accept

accept_da <- function(A, chi, v, s, t0, rt, n) {
## this function draws a drift rate (d) from N(v,s) and computes the
## necessary start point value (a). Checks if that value is allowed
	d <- rnorm(n, v, s)
	a <- chi+A - d*(rt-t0)
	matrix(c(d[a>=0&a<=A],a[a>=0&a<=A]), ncol=2, nrow=length(a[a>=0&a<=A]))
}

stlba.mh <- function(rt,chi,A,v,s,t0,n) {
## here, A, chi etc, are the posterior distributions that we get for these
## params from (eg) winBUGS. (t is the RT)
## we will sample until we get an allowed combination
	if (length(A) >1) {
		res <- NULL
		while (is.null(res)||nrow(res)<n) {
		res <- rbind(res, accept_da(sample(A,n,replace=T), sample(chi,n,replace=T), 
			sample(v,n,replace=T), sample(s,n,replace=T), sample(t0,n,replace=T), rt, n))
		}
	} else {
		res <- NULL
		while (is.null(res)||nrow(res)<n) {
		res <- rbind(res, accept_da(sample(rep(A,n),n,replace=T), sample(rep(chi,n),n,replace=T), 
			sample(rep(v,n),n,replace=T), sample(rep(s,n),n,replace=T), sample(rep(t0,n),n,replace=T), rt, n))
		}
	}
	res[1:n,]
}

stlba.ml=function(rt,chi,A,v,s,t0) {
	# Assumes that observed RT was at node=1.
  	if (length(chi)==1) chi <- rep(chi, length(rt))
  	if (length(A)==1) A <- rep(A, length(rt))
    if (length(v)==1) v <- rep(v, length(rt))
	if (length(s)==1) s <- rep(s, length(rt))
	if (length(t0)==1) t0 <- rep(t0, length(rt))
	toolow=rt< (chi/v + t0)
  	toohigh=rt> (((chi+A)/v) + t0)
  	justright=(!toohigh)&(!toolow)
  
  	d <- v
  	d[toolow] <- (chi[toolow])/(rt[toolow]-t0[toolow])
  	d[toohigh] <- (chi[toohigh]+A[toohigh])/(rt[toohigh]-t0[toohigh])
  
  	a <- (chi+A)-(rt-t0)*v
  	a[toolow] <- A[toolow]
  	a[toohigh] <- 0
  
  	list(a=a,d=d)
}


######### Simulation 1 ######
### Use the maximum likelihood values of the group-level LBA parameters, eg.:
A=1; s=.2; v=1; t0=0; chi =.5

## generate true value for drift rate and start points
n=100
truestarts <- runif(n,0,A)
truedrifts <- rnorm(n, v, s)

## compute RTs
rt=(chi+A-truestarts)/truedrifts + t0

### compute Maximum Likelihood values of STLBA parameters
MLs=stlba.ml(rt=rt,chi=chi, A=A,v=v,s=s,t0=t0)

### Compute posterior distributions of STLBA parameters:
Nsamples <- 1e03
par(mfcol=c(2,5), mar=c(4,4,1,1))
#plot five distributions as example:
for (i in floor(seq(1,n,len=5))) {
	post <- stlba.mh(rt[i],chi,A, v, s, t0, Nsamples)
	plot(density(post[,2]), main=paste("Posterior a[",i,"]", sep=''), xlim=c(0,2)); abline(v=truestarts[i], lty=2); abline(v=MLs$a[i], lty=3)
	legend("topleft", c("true value", "ML"), lty=2:3, bty='n')
	plot(density(post[,1]), main=paste("Posterior d[",i,"]", sep=''), xlim=c(0,2)); abline(v=truedrifts[i], lty=2); abline(v=MLs$d[i], lty=3)
}


######### Simulation 2 ######
### Using eg winBUGS, obtain probability density for the group-level LBA parameters, eg:
Nsamples <- 1e02
A=rnorm(Nsamples,1,.1); s=rnorm(Nsamples,.2,.02); v=rnorm(Nsamples,1,.1) 
t0=rnorm(Nsamples,0.3,.01); chi=rnorm(Nsamples,.5,.05)

## generate true value for drift rate and start points
n=100
truestarts <- runif(n,0,mean(A)) # we assume that the mean of each group-level distributions was the true value
truedrifts <- rnorm(n, mean(v), mean(s))

## compute RTs
rt=(mean(chi)+mean(A)-truestarts)/truedrifts + mean(t0)

### Compute posterior distributions of STLBA parameters:
Nsamples <- 1e03
par(mfcol=c(2,5), mar=c(4,4,1,1))
#plot five distributions as example:
for (i in floor(seq(1,n,len=5))) {
	post <- stlba.mh(rt[i],chi,A, v, s, t0, Nsamples)
	plot(density(post[,2]), main=paste("Posterior a[",i,"]", sep=''), xlim=c(0,2)); abline(v=truestarts[i], lty=2)
	legend("topleft", c("true value"), lty=2, bty='n')
	plot(density(post[,1]), main=paste("Posterior d[",i,"]", sep=''), xlim=c(0,2)); abline(v=truedrifts[i], lty=2)
}