improve findSamplersOnNodes performance

[hoxo-m]

Summary

This PR improves the performance of findSamplersOnNodes() in MCMC_configuration.R.

It preserves the same return values/behavior, while reducing node-matching overhead by flattening sampler target nodes once and mapping matches back to sampler indices.

Scope

• No API changes
• No behavior changes
• Internal performance improvement only

Motivation

While reconfiguring a large MCMC setup, I found that removeSamplers() was slow, and profiling showed that most of the time was spent in findSamplersOnNodes().

In one case, calling conf$findSamplersOnNodes("r") took about 4 minutes before this change.

Implementation

The change reduces repeated matching work by:

• flattening sampler target nodes once,
• building a parallel sampler-index vector, and
• mapping matched node positions back to sampler indices.

This keeps the returned sampler indices unchanged while reducing overhead for large configurations.

Results (local benchmark)

• In my use case, conf$findSamplersOnNodes("r"): 4 min → 1.4 s
• In the reproducible example below, findSamplersOnNodes("A"): 8.47s → 0.06s

Outputs were unchanged (identical(...) == TRUE).

Reproducible example

library(nimble)

M <- 50

code <- nimbleCode({
  for (i in 1:M) {
    for (j in 1:M) {
      y[i, j] ~ dnorm(A[i, j], tau_y)
      A[i, j] ~ dnorm(0, tau_A)
    }
  }
  tau_y ~ dgamma(1, 1)
  tau_A ~ dgamma(1, 1)
})

set.seed(1)
A_true <- matrix(rnorm(M * M, 0, 1), M, M)
y <- A_true + matrix(rnorm(M * M, 0, 0.3), M, M)

constants <- list(M = M)
data <- list(y = y)
inits <- list(A = matrix(0, M, M), tau_y = 10, tau_A = 1)

model <- nimbleModel(code, constants = constants, data = data, inits = inits)
conf  <- configureMCMC(model, monitors = c("A", "tau_y", "tau_A"))

findSamplersOnNodes2 <- function(conf, nodes) {
  samplerConfs <- conf$samplerConfs
  model <- conf$model
  
  if(length(samplerConfs) == 0) return(integer())
  nodes <- model$expandNodeNames(nodes, returnScalarComponents = TRUE, sort = TRUE)
  samplerConfNodesList <- lapply(samplerConfs, function(sc) sc$targetAsScalar)
  
  # Match requested nodes in the flattened node list and map matches back to sampler indices.
  samplerIndices <- 1:length(samplerConfs)
  samplerConfNodesLengths <- unlist(lapply(samplerConfNodesList, length))
  flatSamplerConfNodes <- unlist(samplerConfNodesList)
  flatSamplerIndices <- rep.int(samplerIndices, times = samplerConfNodesLengths)
  flatNodePositions <- unlist(lapply(nodes, function(n) which(n == flatSamplerConfNodes)))
  matchedSamplerIndices <- flatSamplerIndices[flatNodePositions]
  unique(matchedSamplerIndices)
}

system.time({
  out1 <- conf$findSamplersOnNodes("A")
})
#> elapsed
#>    8.47
system.time({
  out2 <- findSamplersOnNodes2(conf, "A")
})
#> elapsed
#>    0.06
identical(out1, out2)
#> TRUE

# Verify that outputs are identical in a case where some sampler targets have length > 1
conf$removeSamplers("A")
for (i in 1:M) {
  conf$addSampler(
    target = paste0("A[", i, ", ]"),
    type   = "RW_block",
    silent = TRUE
  )
}

system.time({
  out3 <- conf$findSamplersOnNodes("A")
})
#> elapsed
#>    0.20
system.time({
  out4 <- findSamplersOnNodes2(conf, "A")
})
#> elapsed
#>    0.03
identical(out3, out4)
#> TRUE

[danielturek]

@hoxo-m Thank you for identifying this inefficiency in findSamplersOnNodes, and taking the time to make this PR.