add ability to sort results of get{Dependencies,Parents}

nimbleModel/R/modelBaseClass.R

@@ -360,17 +360,17 @@ modelBase_nClass <- nClass(
     },
     getDependencies = function(nodes, self = TRUE, downstream = FALSE, immediateOnly = FALSE,
                                nodesAsChars = getNimbleModelOption('nodesAsChars'),
-                               returnScalarComponents = FALSE
+                               returnScalarComponents = FALSE, .sort = FALSE
                                ) {
       nimbleModel::getDependencies(modelDef, nodes, self, downstream, immediateOnly,
-                                   nodesAsChars, returnScalarComponents)
+                                   nodesAsChars, returnScalarComponents, .sort)
     },
-    getParents = function(nodes, self = TRUE, upstream = FALSE, immediateOnly = FALSE,
+    getParents = function(nodes, self = FALSE, upstream = FALSE, immediateOnly = FALSE,
                           nodesAsChars = getNimbleModelOption('nodesAsChars'),
-                          returnScalarComponents = FALSE
+                          returnScalarComponents = FALSE, .sort = FALSE
                           ) {
       nimbleModel::getParents(modelDef, nodes, self, upstream, immediateOnly,
-                              nodesAsChars, returnScalarComponents)
+                              nodesAsChars, returnScalarComponents, .sort)
     },
     # TODO: not working because `nimbleModel::getNodes` needs the model not just modelDef.
     # Once we integrate modelClass with modelBase_nClass, we should be able to pass `self`.

nimbleModel/R/modelDef.R

@@ -1039,28 +1039,30 @@ getDependencies <- function(modelDef, nodes,
                             self = TRUE,
                             downstream = FALSE, immediateOnly = FALSE,
                             nodesAsChars = getNimbleModelOption('nodesAsChars'),
-                            returnScalarComponents = FALSE
+                            returnScalarComponents = FALSE, .sort = FALSE
                             ) {
   traverseGraph(modelDef$downstreamRules, modelDef$declRules,
     nodes = nodes,
     down = TRUE, self = self,
     follow = downstream, immediateOnly = immediateOnly,
-    nodesAsChars = nodesAsChars, returnScalarComponents = returnScalarComponents
-  )
+    nodesAsChars = nodesAsChars, returnScalarComponents = returnScalarComponents,
+    .sort = .sort, modelDef = modelDef
+    )
 }
 
 getParents <- function(modelDef, nodes,
                        self = FALSE,
                        upstream = FALSE, immediateOnly = FALSE,
                        nodesAsChars = getNimbleModelOption('nodesAsChars'),
-                       returnScalarComponents = FALSE
+                       returnScalarComponents = FALSE, .sort = FALSE
                        ) {
   traverseGraph(modelDef$upstreamRules, modelDef$declRules,
     nodes = nodes,
     down = FALSE, self = self,
     follow = upstream, immediateOnly = immediateOnly,
-    nodesAsChars = nodesAsChars, returnScalarComponents = returnScalarComponents
-  )
+    nodesAsChars = nodesAsChars, returnScalarComponents = returnScalarComponents,
+    .sort = .sort, modelDef = modelDef
+    )
 }
 
 

nimbleModel/R/nodeRules.R

@@ -478,8 +478,8 @@ nodeRangeClass <- R6Class(
         varName = varName
       )
     },
-    toVarRange = function() {
-      varRangeClass$new(indexRanges, rangeToIndexSlot = rangeToIndexSlot, varName = varName, fromStochRule = decl$stoch)
+    toVarRange = function(fromStochRule = decl$stoch) {
+      varRangeClass$new(indexRanges, rangeToIndexSlot = rangeToIndexSlot, varName = varName, fromStochRule = fromStochRule)
     },
 
     toChar = function() {

nimbleModel/R/processModelGraph.R

@@ -242,8 +242,15 @@ traverseGraph <- function(streamRules, declRules,
                           nodes, down, self = TRUE,
                           follow = FALSE, immediateOnly = FALSE,
                           nodesAsChars = getNimbleModelOption('nodesAsChars'),
-                          returnScalarComponents = FALSE
+                          returnScalarComponents = FALSE, .sort = FALSE,
+                          modelDef = NULL
                           ) {
+
+  # A single varRange can have elements that don't share a sortID when converted to calcRange representation,
+  # so we can't sort varRanges.
+  if (.sort && !nodesAsChars)
+    warning("`.sort=TRUE` is provided only for back compatibility and requires the use of character representations of nodes")
+
   if (inherits(nodes, "varRangeClass")) nodes <- list(nodes) # We use `lapply` on 'nodes' later.
 
   results <- traverseGraphRecurse(streamRules, nodes, down, follow, immediateOnly)
@@ -323,12 +330,48 @@ traverseGraph <- function(streamRules, declRules,
     return(NULL)
   }
   results <- removeDuplicateVarRanges(results)
-  if (nodesAsChars) {
-    return(unlist(sapply(results, \(x) x$toVarChars(expandScalars = returnScalarComponents))))
+
+  # Remove RHSonly by passing through declRules.
+  results <- flatten(lapply(results, \(vr)
+                            lapply(modelDef$declRules[[getVarName(vr)]]$rules, \(rule) {
+                              nodeRange <- rule$apply(vr)
+                              if(is.null(nodeRange)) return(NULL) else return(nodeRange$toVarRange(fromStochRule = vr$fromStochRule))
+                              })))
+  if (!length(results)) {
+    return(NULL)
+  }  
+
+  if (.sort) {
+    # Ordering is only relevant at calcRange stage and a single nodeRange can contain
+    # elements with various sortIDs, so we convert to nodeChars first and then get their
+    # sortID by creating a temporary calcRange for each.
+    nodeChars <- unlist(lapply(results, \(vr)
+                               lapply(modelDef$calcRules[[getVarName(vr)]]$rules,
+                                      \(rule) {
+                                        tmp <- rule$apply(vr)
+                                        if(!is.null(tmp)) return(tmp$toNodeChars()) else return(NULL)
+                                      }
+                                      )))
+    calcRanges <- flatten(lapply(nodeChars, function(node) {
+      lapply(modelDef$calcRules[[getVarName(node)]]$rules, function(rule) {
+        rule$makeCalcRange(rule$apply(node))
+      })
+    }))
+    if (length(nodeChars) != length(calcRanges))
+      stop("unexpected mismatch between node character representation and calcRanges in `getNodes` sorting")
+    ord <- order(sapply(calcRanges, \(x) x$sortID))
+    results <- nodeChars[ord]
+    names(results) <- NULL
+    if (returnScalarComponents)
+      results <- unlist(lapply(results, \(x) varRangeClass$new(x)$toVarChars(expandScalars = TRUE)))
   } else {
-    if (returnScalarComponents)   # TODO: put into new messaging system
-      warning("one must request result as characters via `nodesAsChars` in order to use `returnScalarComponents`")
-  } 
+    if (nodesAsChars) {
+      return(unlist(lapply(results, \(x) x$toVarChars(expandScalars = returnScalarComponents))))
+    } else {
+      if (returnScalarComponents)   # TODO: put into new messaging system
+        warning("one must request result as characters via `nodesAsChars` in order to use `returnScalarComponents`")
+    }
+  }
   return(results)                                                  
 }
 

nimbleModel/tests/testthat/test-modelGraph.R

@@ -1257,7 +1257,7 @@ test_that("basic check of graph interface", {
 
     result <- getParents(modelDef, 'y', upstream = TRUE)
     expect_identical(sapply(result, function(node) node$varName),
-                     c('theta','mu0','sigma'))
+                     c('theta','mu0'))
 
 })
 
@@ -1290,6 +1290,8 @@ test_that("getDependencies deals with repeated parents", {
         y ~ dnorm(mu, sd = tau)
         z ~ dnorm(mu, 1)
         w <- z + 3
+        tau ~ dunif(0,1)
+        mu ~ dunif(0,1)
     })
     modelDef <- modelDefClass$new(code)
 
@@ -1338,6 +1340,7 @@ test_that("getParents traversal with mix of stoch/determ edges", {
 test_that("same dependent on RHS", {
     code <- quote({
         y ~ dnorm(mu, sd = mu)
+        mu ~ dunif(0,1)
     })
     modelDef <- modelDefClass$new(code)
 
@@ -1357,7 +1360,6 @@ test_that("basic hierarchical models", {
         sigma ~ dunif(0, 1)
         for(i in 1:3)
             z[k[i]] ~ dnorm(y[k[i]], 1)
-
     })
     k <- c(2,4,7)
     model <- nimbleModel(code, constants = list(k = k))
@@ -1432,18 +1434,18 @@ test_that("basic hierarchical models", {
 
     result <- getParents(model$modelDef, 'y', upstream = TRUE)
     expect_identical(sapply(result, function(node) node$varName),
-                     c('mu','tau','mu0','sigma','bnd'))
+                     c('mu','tau','sigma'))
 
     result <- getParents(model$modelDef, 'y[1:5]', upstream = TRUE)
     expect_identical(sapply(result, function(node) node$varName),
-                     c('mu','tau','mu0','sigma','bnd'))
+                     c('mu','tau','sigma'))
     expect_equal(result[[1]]$indexRanges, 
                list(newIndexRange(quote(1:5))))
 
     result <- getParents(model$modelDef, varRangeClass$new(list(newIndexRange(quote(1:5))), varName = 'y'),
                          upstream = TRUE)
     expect_identical(sapply(result, function(node) node$varName),
-                     c('mu','tau','mu0','sigma','bnd'))
+                     c('mu','tau','sigma'))
     expect_equal(result[[1]]$indexRanges, 
                list(newIndexRange(quote(1:5))))
 
@@ -1488,14 +1490,11 @@ test_that("basic hierarchical models", {
                      c('w','mu','y'))
 
     result <- getParents(model$modelDef, 'pr[2,2]')
-    expect_identical(sapply(result, function(node) node$varName),
-                     'S')
-    expect_equal(result[[1]]$indexRanges, 
-               list(newIndexRange(quote(1:10)), newIndexRange(quote(1:10))))
+    expect_identical(result, NULL)
 
     result <- getParents(model$modelDef, 'mu[3]')
     expect_identical(sapply(result, function(node) node$varName),
-                     c('mu0','pr','mu00','z'))
+                     c('mu0','pr','mu00'))
     expect_equal(result[[2]]$indexRanges, 
                list(newIndexRange(quote(1:10)), newIndexRange(quote(1:10))))
 
@@ -1536,15 +1535,15 @@ test_that("basic hierarchical models", {
 
     result <- getParents(model$modelDef, 'y[1,1:3]')
     expect_identical(sapply(result, function(node) node$varName),
-                     c('mn','pr','X','beta'))
+                     c('mn','pr','beta'))
 
     result <- getParents(model$modelDef, 'y[1,2]')
     expect_identical(sapply(result, function(node) node$varName),
-                     c('mn','pr','X','beta'))
+                     c('mn','pr','beta'))
 
     result <- getParents(model$modelDef, c('y[1,2]','y[2,3]'))
     expect_identical(sapply(result, function(node) node$varName),
-                     c('mn','pr','X','beta'))
+                     c('mn','pr','beta'))
 
 })
 
@@ -1638,12 +1637,19 @@ test_that("state-space model", {
     expect_equal(result[[1]],
                  varRangeClass$new(list(newIndexRange(2)), varName = 'y'))
 
+    code <- quote({
+        for(i in 2:4)
+            y[i]~dnorm(y[i-1],1)
+        y[1] ~ dnorm(0,1)
+    })
+    model <- nimbleModel(code)
+    modelDef <- model$modelDef
+
     result <- getParents(modelDef, 'y[3]', upstream = TRUE)
     expect_length(result, 2)
     expect_equal(result[[2]],
                  varRangeClass$new(list(newIndexRange(1)), varName = 'y'))
 
-
     code <- quote({
         for(i in 1:5)
             y[i] ~ dnorm(z[i], sd = tau)
@@ -1776,10 +1782,10 @@ test_that("complicated input varRange", {
                      vr$indexRanges[c(2,1)])
 
     result <- getDependencies(modelDef, vr)
-    expect_length(result, 2)
+    expect_length(result, 1)
     expect_identical(sapply(result, function(node) node$varName),
-                     c('theta','y'))
-    expect_equal(result[[2]],
+                     c('y'))
+    expect_equal(result[[1]],
                  varRangeClass$new(list(newIndexRange(matrix(c(3,5,5,1), nrow = 2)),
                                         newIndexRange(quote(2:3))),
                             rangeToIndexSlot = list(c(1,3), 2),
@@ -2558,12 +2564,12 @@ test_that("missing indexing", {
     modelCode <- quote({
         y <- sum(mu[])
     })
+    m <- nimbleModel(modelCode, dimensions = list(mu = 5))
+    modelDef <- m$modelDef
 
-    modelDef <- modelDefClass$new(modelCode, dimensions = list(mu = 5))
-
-    expect_equal(getParents(modelDef, 'y')[[1]],
+    expect_equal(getNodes(m, includeRHSonly = TRUE)[[2]]$toVarRange(),
                  varRangeClass$new(list(newIndexRange(quote(1:5))),
-                                        varName = 'mu', fromStochRule = FALSE))
+                                        varName = 'mu'))
 
     expect_equal(getDependencies(modelDef, 'mu[2]')[[1]],
                  varRangeClass$new(list(), varName = 'y', fromStochRule = FALSE))