Skip to content

Conversation

@kenkellner
Copy link
Collaborator

@kenkellner kenkellner commented Oct 23, 2025

Currently when a model uses several macros, it can lead to (valid) NIMBLE code which is not particularly simple/readable. For example, if you specify a simple occupancy model with macros like this:

occ <- nimbleCode({
  psi[1:nsites] <- LINPRED(~x[1:nsites], link=logit, coefPrefix=state_)
  p[1:nsites, 1:noccs] <- LINPRED(~x[1:nsites] + x2[1:nsites, 1:noccs], link=logit, coefPrefix=det_)

  z[1:nsites] ~ FORLOOP(dbern(psi[1:nsites]))
  y[1:nsites, 1:noccs] ~ FORLOOP(dbern(p[1:nsites, 1:noccs]*z[1:nsites]))
})

The resulting model code is

{
    for (i_1 in 1:nsites) {
        logit(psi[i_1]) <- state_Intercept + state_x * x[i_1]
    }
    state_Intercept ~ dnorm(0, sd = 1000)
    state_x ~ dnorm(0, sd = 1000)
    for (i_2 in 1:nsites) {
        for (i_3 in 1:noccs) {
            logit(p[i_2, i_3]) <- det_Intercept + det_x * x[i_2] +
 
                det_x2 * x2[i_2, i_3]
        }
    }
    det_Intercept ~ dnorm(0, sd = 1000)
    det_x ~ dnorm(0, sd = 1000)
    det_x2 ~ dnorm(0, sd = 1000)
    for (i_4 in 1:nsites) {
        z[i_4] ~ dbern(psi[i_4])
    }
    for (i_5 in 1:nsites) {
        for (i_6 in 1:noccs) {
            y[i_5, i_6] ~ dbern(p[i_5, i_6] * z[i_5])
        }
    }
}

There's nothing incorrect about that code, but it's not how someone would normally write it. They'd probably put the z model in the same loop as the calculation for psi, and the y model in the same loop as p, etc. Furthermore because we are cautious about not mixing up for loop indices when expanding macros, we end up with six indices (i_1 - i_6) when strictly speaking only two are needed.

This PR adds a function simplifyForLoops which takes NIMBLE code and attemps to collapse for loops that share indices, and also to replace indices with a smaller set of simpler ones. This is possible because the function is meant to be run on the final expanded code, when we know the entire structure. For example, when applied to the model above, the result is

{
    for (i in 1:nsites) {
        logit(psi[i]) <- state_Intercept + state_x * x[i]
        for (j in 1:noccs) {
            logit(p[i, j]) <- det_Intercept + det_x * x[i] + 
                det_x2 * x2[i, j]
            y[i, j] ~ dbern(p[i, j] * z[i])
        }
        z[i] ~ dbern(psi[i])
    }
    state_Intercept ~ dnorm(0, sd = 1000)
    state_x ~ dnorm(0, sd = 1000)
    det_Intercept ~ dnorm(0, sd = 1000)
    det_x ~ dnorm(0, sd = 1000)
    det_x2 ~ dnorm(0, sd = 1000)
}

Which is simpler and much closer to how a user would typically write the model.

Right now this is just a standalone function and is not run automatically at any point. Furthermore to use it you'd have to write the code with macros, create the model object, run the function on the output model code (post macro processing), and then re-create the model with the new code, which is a bit clunky.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants