stochastic indexing notes

[paciorek]

To capture some thoughts from an in-person discussion.

If we have

y[i+1] ~ dnorm(mu[k[i]-3],1)

we'd like to capture the fact that the mu -> y calcRule has stochastic indexing and that the indexing expr is k[i]-3.

I think a tricky thing is that we could want: getDependencies(mu[1]) or getDependencies(mu[3:5])). In both cases the result from nimbleModel processing may be y[2:100] but we need the info about the mu indexing captured to be able to provide it to "calc_stoch" so that it can check if k[i]-3 equals 1 or 3,4,5.

Here's a bit of benchmarking that shows that if we could implement something that conditions within the calculation we can save a good proportion of the time difference between our current naive approach and the underlying fundamental calculations. Results will obviously differ depending on the proportion of calculations that actually need to be done.

## Check within calculate loop.
bench0 <- nFunction(
    fun = function(y='numericVector', k = 'integerVector', mu = 'numericVector') {
        value <- 0
        n <- length(y)
        for(i in 1:n)
            if(k[i] == 1)
                value <- value + dnorm(y[i], mu[1], 1, log = TRUE)
    }
)

## Naive/current.
bench1 <- nFunction(
    fun = function(y='numericVector', k='integerVector',mu = 'numericVector') {
        value <- 0
        n <- length(y)
        for(i in 1:n)
                value <- value + dnorm(y[i], mu[k[i]], 1, log = TRUE)
    }
)

## Oracle
bench2 <- nFunction(
    fun = function(y='numericVector', inds = 'integerVector', mu = 'numericVector') {
        value <- 0
        n <- length(inds)
        for(i in 1:n)
                value <- value + dnorm(y[inds[i]], mu[1], 1, log = TRUE)
    }
)

n <- 1e7
y <- rnorm(n)
k <- rep(1:10, n/10)
mu <- rnorm(10)
inds <- which(k==1)   # 10% "sparsity"

cbench0 <- nCompile(bench0)
cbench1 <- nCompile(bench1)
cbench2 <- nCompile(bench2)

rbenchmark::benchmark(cbench0(y, k, mu), cbench1(y,k,mu), cbench2(y, inds, mu), replicates = 10)

> benchmark(cbench0(y, k, mu), cbench1(y,k,mu), cbench2(y, inds, mu))
                  test replications elapsed relative user.self sys.self
1    cbench0(y, k, mu)          100   7.720    1.621     4.374    3.345
2    cbench1(y, k, mu)          100  13.916    2.922    10.372    3.536
3 cbench2(y, inds, mu)          100   4.762    1.000     2.460    2.300