Solved – How to conditionally run element of JAGS script based on user supplied variable

bayesianjagsr

Background: I often find that when working with JAGS that I end up with a lot of JAGS scripts as I explore different models. After a while I might settle on a set of models that I'm going to report, but then there is the issue of keeping settings consistent across these models that are meant to be consistent. For example, the priors for parameters that are common should be consistent, or the variable names for parameters should be consistent. While consistency is not that difficult to achieve, it's often a little annoying. For example, I might decide that I want to call something theta rather than beta, and now all the scripts need updating. So I'm trying to work out how I can have one script that can serve multiple purposes.

Example:
Here's a simple example. Assume I have two models, one with time as a predictor and one without. I supply to jags a logical variable timevariable to indicate whether I want to model time or not. Thus, I'd like to be able to write something like this.

for (i in 1:length(y)) {
    if (timevariable) {
        mu[i] <- beta1 + beta2 * time[i]
    } else {
        mu[i] <- beta1
    }
    y[i]  ~ dnorm(mu[i], tau)
}

However, if-then is not part of the lanugage. Note that I'm not trying to incorporate an if-then based on estimated values of parameters. Rather I just want to have a script that can take flexible input.

Question

What is a good way to conditionally run parts of a JAGS script based on a user supplied variable?

Initial thoughts

Put a macro name in the text and use a text replacement function in R like gsub to replace the macro name with the desired text. I guess this would work, but it seems like a bit of a hack. It would be nice to have everything inside the script.

Best Answer

Here is one example of implementing a basic macro substitution system for JAGS scripts.

Explanation of the system

Define a function that takes as arguments any optional elements of the script.
For any aspects of the script that vary across argument values, record a macro token. This should be some unique text. Starting and ending with some symbols may assist to make this unambiguous.
For each macro token include code for what the replacement string should be under all possible combinations of the function arguments.
Replace the macro tokens with appropriate placement text based on arguments.
The code below provides one example of how macros could be specified and how to apply the macros to the raw script (suggestions for doing this more elegantly are welcome).
The function returns a replaced model string that could if necessary be passed to a textConnection function for use with rjags.

I like this system for a few reasons:

the raw script is easy to read
the resulting script has proper indentation
you don't have to worry about errors related to commands appearing on the same line

Example

Specifically, this example, aims to allow the user to fit a particular type of multilevel nonlinear model. It is designed to allow for three functional forms: a two parameter power, three parameter power, and a three parameter exponential.

The macro section structures macros as a list of lists. The top level list contains one element for each macro token. For each macro token, there is the macro token text, and the conditional replacement text.

Finally, a for loop applies all the macro replacements to the raw script.

See below (scrolling is required):

jags_model <- function (f=c('power2', 'power3', 'exp3')) {
    f <- match.arg(f)

    # raw script
    script <- 
"model {
    # Model
    for (i in 1:length(y)) {
        $FUNCTION
        y[i]  ~ dnorm(mu[i], tau[subject[i]])
    }

    # Random coefficients
    for (i in 1:N) {    
        theta1[i] ~ dnorm(theta1.mu, theta1.tau)T(0, 1000)
        theta2[i] ~ dnorm(theta2.mu, theta2.tau)T(-10,  0)
        $THETA3DISTRIBUTION
        sigma[i] ~ dnorm(sigma.mu, sigma.tau)T(0, 100);
        tau[i] <- 1/(sigma[i]^2)
    }


    theta1.mu  ~ dunif(0, 100)
    theta2.mu   ~ dunif(-2, 0)
    $THETA3PRIOR.MU
    sigma.mu ~ dunif(0, 20)

    theta1.sigma ~ dunif(0, 100)
    theta2.sigma ~ dunif(0, 2)
    $THETA3PRIOR.SIGMA
    sigma.sigma ~ dunif(0, 10)

    # Transformations
    theta1.tau <- 1/(theta1.sigma^2)
    theta2.tau <- 1/(theta2.sigma^2)
    $THETA3.TAU 
    sigma.tau <- 1/(sigma.sigma^2)
}"


    # define macros
    macros <- list(list("$FUNCTION",  
              switch(f,
                     power2="mu[i] <- theta1[subject[i]] * pow(trial[i], theta2[subject[i]])", 
                     power3="mu[i] <- theta1[subject[i]] * pow(trial[i], theta2[subject[i]]) + theta3[subject[i]];",
                     exp3="mu[i] <- theta1[subject[i]] * exp(theta2[subject[i]] * (trial[i] - 1)) + theta3[subject[i]];") ), 
            list("$THETA3DISTRIBUTION",  
          switch(f, 
                 power3=, exp3= "theta3[i] ~ dnorm(theta3.mu, theta3.tau)T(0, 1000)", 
                 power2="") ), 
        list("$THETA3PRIOR.MU",  
              switch(f, 
                     power3=, exp3= "theta3.mu  ~ dunif(0, 100)", 
                     power2="") ), 
            list("$THETA3PRIOR.SIGMA",  
          switch(f, 
                 power3=, exp3= "theta3.sigma ~ dunif(0, 100)", 
                 power2="") ), 
        list("$THETA3.TAU",  
          switch(f, 
                 power3=, exp3= "theta3.tau <- 1/(theta3.sigma^2)", 
                 power2="") )
        )

    # apply macros
    for (m in seq(macros)) {
        script <- gsub(macros[[m]][1], macros[[m]][2], script, fixed=TRUE)
    }
    script
}

Demonstration

Thus, we can produce the processed JAGS model with

x <- jags_model(f='power3')

And if we want to view the resulting model, we can do

cat(x)

which results in

model {
    # Model
    for (i in 1:length(y)) {
        mu[i] <- theta1[subject[i]] * pow(trial[i], theta2[subject[i]]) + theta3[subject[i]];
        y[i]  ~ dnorm(mu[i], tau[subject[i]])
    }

    # Random coefficients
    for (i in 1:N) {    
        theta1[i] ~ dnorm(theta1.mu, theta1.tau)T(0, 1000)
        theta2[i] ~ dnorm(theta2.mu, theta2.tau)T(-10,  0)
        theta3[i] ~ dnorm(theta3.mu, theta3.tau)T(0, 1000)
        sigma[i] ~ dnorm(sigma.mu, sigma.tau)T(0, 100);
        tau[i] <- 1/(sigma[i]^2)
    }


    theta1.mu  ~ dunif(0, 100)
    theta2.mu   ~ dunif(-2, 0)
    theta3.mu  ~ dunif(0, 100)
    sigma.mu ~ dunif(0, 20)

    theta1.sigma ~ dunif(0, 100)
    theta2.sigma ~ dunif(0, 2)
    theta3.sigma ~ dunif(0, 100)
    sigma.sigma ~ dunif(0, 10)

    # Transformations
    theta1.tau <- 1/(theta1.sigma^2)
    theta2.tau <- 1/(theta2.sigma^2)
    theta3.tau <- 1/(theta3.sigma^2) 
    sigma.tau <- 1/(sigma.sigma^2)
}

Related Solutions

Solved – R lm simultaneous parameter tests

Do a full and reduced test by fitting a model with all the parameters, then another model with the reduced model (leave out the terms that you want to test with 0). Then do anova(fit1,fit2) to compute the F test.

If you want to test for values other than 0 then use the offset function in the reduced formula.

Solved – Calculating prediction variance from JAGS model of a bernoulli outcome in R

I think the best way to generate uncertainty values around each prediction is to send the desired predictions through the original JAGS model. This way, the model will sample from the distributions of all parameters, regardless of whether they are positive or negative, and generate a distribution around each prediction. This distribution will reflect the uncertainty associated with all parameters, at the value of each predictor of interest. Aspects of this distribution can then be used for plotting. To implement this in R, first generate a mean of x2, and a range of x1 values for which you want to make predictions for. Here I am going to predict over the entire range of x1 values in the data, divided into 101 intervals across the range of x1 values.

x1.range <- seq(min(x1),max(x1), by = (max(x1)-min(x1))/100)
x2.m     <- mean(x2)

I then modify the JAGS model to make predictions.

jags.model = "
model {
#model code
for (i in 1:N){
y[i] ~ dbern(p[i])
p[i] <- 1 / (1 + exp(-z[i]))
z[i] <- a + b*x1[i] + c*x2[i]
}

#predictions over x1, holding x2 constant at its mean value
for (i in 1:N.sim){
  z.x1[i] <- a + b*x1.range[i] + c*x2.m
}

a ~ dnorm(0, .0001)
b ~ dnorm(0, .0001)
c ~ dnorm(0, .0001)
}
"

Make sure to add the new data to your data list:

#setup data as list
data = list(y=y, x1=x1, x2=x2,
            x1.range = x1.range, x2.m = x2.m,
            N = length(y), N.sim = length(x1.range))

Compile the model, and the use the coda function to sample from the prediction posterior:

#compile the model
require(rjags)
j.model <- jags.model(file = textConnection(jags.model),
                      data=data,
                      n.chains=3)

#sample from the prediction posterior
x1.jags.out <- coda.samples(model=j.model, variable.names = c('z.x1'),n.iter=1000)

This will generate a mean and standard deviation for each predicted value of y, at each condition of x1 and x2 simulated, taking the full distribution and uncertainty of all parameter values into account. Finally, for plotting, grab the estimated mean and standard deviation (or whatever your favorite metric of variance is), and plot:

#generate plot, plotting mean and +/- 1 standard deviation. 
plot(x1.pred ~ x1.range, cex=0)
lines(smooth.spline(x1.pred~ x1.range), lwd=2, col='purple')
lines(smooth.spline((x1.pred+x1.pred.sd) ~ x1.range), lty=2, col='purple')
lines(smooth.spline((x1.pred-x1.pred.sd) ~ x1.range), lty=2, col='purple')

#shade error region
polygon(c(x1.range,rev(x1.range))
        ,c(x1.pred+x1.pred.sd,rev(x1.pred - x1.pred.sd))
        ,col=adjustcolor('purple',0.3)
        ,lty=0
)

If specific probabilities are desired, then these z values need to be inverse logit transformed.