update support for (mostly) simple math ops.

nimbleModel/R/getSymbolicParentNodes.R

@@ -1,3 +1,12 @@
+# Model code has nCompiler keywords changed (e.g., `chol` to `nChol`; see modelDefClass), so lists
+# of function names used in model processing also need the changed names.
+withChangedKeywords <- function(fnNames) {
+  unique(c(
+    fnNames,
+    unlist(nCompiler:::nKeyWords[intersect(names(nCompiler:::nKeyWords), fnNames)], use.names = FALSE)
+  ))
+}
+
 nimblePreevaluationFunctionNames <- c(
   "+",
   "-",
@@ -93,6 +102,7 @@ nimblePreevaluationFunctionNames <- c(
   "nimNumeric", "nimMatrix", "nimArray",
   "length"
 )
+nimblePreevaluationFunctionNames <- withChangedKeywords(nimblePreevaluationFunctionNames)
 
 nimbleOrRfunctionNames <- c(
   ":", "[", "(",
@@ -130,6 +140,7 @@ nimbleOrRfunctionNames <- c(
   "asRow", "asCol",
   "chol",
   "inverse",
+  "logdet",
   "forwardsolve", "backsolve", "solve",
   "nimEigen", "nimSvd",
   ">", "<", ">=", "<=",
@@ -147,12 +158,7 @@ nimbleOrRfunctionNames <- c(
   "diag",
   "length"
 )
-# Model code has nCompiler keywords changed (e.g., `chol` to `nChol`; see modelDefClass), so also
-# allow the changed names of the functions above.
-nimbleOrRfunctionNames <- unique(c(
-  nimbleOrRfunctionNames,
-  unlist(nCompiler:::nKeyWords[intersect(names(nCompiler:::nKeyWords), nimbleOrRfunctionNames)], use.names = FALSE)
-))
+nimbleOrRfunctionNames <- withChangedKeywords(nimbleOrRfunctionNames)
 
 
 # Determine the RHS pieces, expressed symbolically.

nimbleModel/R/mathCompatibility.R

@@ -6,45 +6,7 @@ linkInverses <- list(
   "probit" = quote(phi()),
   "cloglog" = quote(icloglog())
 )
-#' @export
-expit <- function(x) 1 / (1 + exp(-x))
-#' @export
-probit <- function(x) qnorm(x)
-#' @export
-ilogit <- function(x) 1 / (1 + exp(-x)) # same as expit
-#' @export
-iprobit <- function(x) pnorm(x)
-#' @export
-icloglog <- function(x) 1 - exp(-exp(x))
-#' @export
-cloglog <- function(x) log(-log(1 - x))
-#' @export
-nimEquals <- function(x1, x2) if (x1 == x2) 1 else 0 # "equals" conflicts with a usage in testthat
-#' @export
-logfact <- function(x) lfactorial(x) # lgamma(x+1) (from R) not equivalent to lgamma1p(x) (from our C). lgamma seems to be endowed with setting numerically 0 values to zero.
-#' @export
-loggam <- function(x) lgamma(x)
-#' @export
-logit <- function(x) log(x / (1 - x))
-#' @export
-phi <- function(x) pnorm(x) # same as iprobit
-#' @export
-pow <- function(x1, x2) x1^x2
-#' @export
-nimStep <- function(x) ifelse(x >= 0, 1L, 0L) # We rename step to nimStep before execution to avoid masking the R step function
-#' @export
-inverse <- function(x) solve(x)
-#' @export
-cube <- function(x) x * x * x # Doing is this way instead of x^3 makes it numerically identical to eigen's cube
-#' @export
-inprod <- function(v1, v2) sum(v1 * v2)
-#' @export
-logdet <- function(m) {
-  # this ensures that R returns NaN for negative determinants, to mimic C
-  out <- determinant(m, logarithm = TRUE)
-  if (out$sign == 1) {
-    return(out$modulus)
-  } else {
-    return(NaN)
-  }
-}
+
+# The BUGS math functions (e.g., `expit`, `cloglog`, `pow`, `inprod`) are provided by nCompiler.
+# `step`, `equals`, `logdet` and `inverse` in model code are changed to nCompiler's `nStep`, `nEquals`,
+# `nLogdet` and `nInverse` (see `nmKeywords` and nCompiler's `nKeyWords`).

nimbleModel/R/modelDef.R

@@ -555,7 +555,7 @@ modelDefClass <- R6Class(
 
     # Overwrite declInfo (both LHS and RHS) with constants replaced; only replaces scalar constants.
     replaceAllConstants = function() {
-      constantsEnv <- list2env(constants, parent = getDefaultNamespace())
+      constantsEnv <- list2env(constants, parent = makePreevaluationEnv())
       for (i in seq_along(declInfo)) {
         newCode <- replaceConstantsRecurse(declInfo[[i]]$code, constantsEnv)$code
         declInfo[[i]] <<- modelDeclClass$new(
@@ -671,7 +671,7 @@ modelDefClass <- R6Class(
 
     # Reparameterizations and altParams come from the distribution definitions (e.g., `cholesky = chol(prec)`),
     # after keywords were changed in `modelCode`, so change keywords again (e.g., `chol` to `nChol`).
-    # This is done after lifting and constant replacement, which look for the original names (e.g., `chol` in `neverReplaceable`).
+    # This is done after lifting and constant replacement; `liftedNonScalarCalls` and `neverReplaceable` include both the original and changed names.
     changeKeywordsInDecls = function() {
       for (iDecl in seq_along(declInfo)) {
         decl <- declInfo[[iDecl]]
@@ -1275,7 +1275,7 @@ replaceConstantsRecurse <- function(code, constantsEnv, do.eval = TRUE) {
         callChar <- as.character(code[[1]])
         if (exists(callChar, constantsEnv)) {
           if (!is.vectorized(code)) {
-            if (is.null(neverReplaceable[[callChar]])) {
+            if (!(callChar %in% neverReplaceable)) {
               if (isTRUE(callChar %in% nimblePreevaluationFunctionNames)) {
                 if (inherits(get(callChar, constantsEnv), "function")) {
                   testcode <- as.numeric(eval(code, constantsEnv))
@@ -1292,16 +1292,43 @@ replaceConstantsRecurse <- function(code, constantsEnv, do.eval = TRUE) {
   stop("unable to process `", safeDeparse(code), "`.")
 }
 
-neverReplaceable <- list(
-  # Only the names matter, any non-null value will do.
-  chol = TRUE,
-  inverse = TRUE,
-  CAR_calcNumIslands = TRUE,
-  CAR_calcC = TRUE,
-  CAR_calcM = TRUE,
-  CAR_calcEVs2 = TRUE,
-  CAR_calcEVs3 = TRUE
-)
+# Environment in which pre-evaluation finds functions in `nimblePreevaluationFunctionNames`.
+# It contains those that are not in base R (e.g., `pow`, `nStep`), from nimbleModel's namespace or its imports (nCompiler).
+# Its parent is `baseenv()`, so the user's global environment is not seen.
+makePreevaluationEnv <- function() {
+  ns <- topenv(environment(makePreevaluationEnv))
+  fns <- list()
+  for (fnName in setdiff(nimblePreevaluationFunctionNames, ls(baseenv(), all.names = TRUE))) {
+    fn <- get0(fnName, envir = ns, inherits = FALSE)
+    if (is.null(fn)) fn <- get0(fnName, envir = parent.env(ns), inherits = FALSE)
+    if (is.function(fn)) fns[[fnName]] <- fn
+  }
+  list2env(fns, parent = baseenv())
+}
+
+# Calls that are never pre-evaluated, even if all arguments are constants.
+neverReplaceable <- withChangedKeywords(c(
+  "chol",
+  "inverse",
+  "CAR_calcNumIslands",
+  "CAR_calcC",
+  "CAR_calcM",
+  "CAR_calcEVs2",
+  "CAR_calcEVs3"
+))
+
+# Non-scalar calls that are lifted when used as distribution arguments.
+# (Scalar argument expressions are lifted in general; see `isExprLiftable`.)
+liftedNonScalarCalls <- withChangedKeywords(c(
+  "chol",
+  "inverse",
+  "PDinverse_logdet",
+  "CAR_calcNumIslands",
+  "CAR_calcC",
+  "CAR_calcM",
+  "CAR_calcEVs2",
+  "CAR_calcEVs3"
+))
 
 liftedCallsDoNotAddIndexing <- c(
   "CAR_calcNumIslands"
@@ -1341,9 +1368,9 @@ isExprLiftable <- function(paramExpr, type = NULL) {
   }
   if (is.call(paramExpr)) {
     callText <- getCallText(paramExpr)
-    if (callText %in% names(neverReplaceable)) {
+    if (callText %in% liftedNonScalarCalls) {
       return(TRUE)
-    } # Special calls that are not lifted.
+    } # Special non-scalar calls that are lifted.
     if (length(paramExpr) == 1) {
       return(FALSE)
     } # Don't lift function calls with no arguments.

nimbleModel/tests/testthat/test-mathFunctions.R

@@ -0,0 +1,70 @@
+test_that("BUGS math functions work in model code", {
+  code <- nimbleCode({
+    a ~ dnorm(0, 1)
+    b1 <- step(a) + equals(a, 0.5) + pow(a, 2) + inprod(v[1:2], v[1:2])
+    b2 <- expit(a) + ilogit(a) + logit(0.3) + probit(0.3) + iprobit(a) + phi(a) +
+      cloglog(0.3) + icloglog(a) + loggam(2.5) + lfactorial(3) + cube(a)
+    d <- logdet(S[1:2, 1:2])
+    Sinv[1:2, 1:2] <- inverse(S[1:2, 1:2])
+    v[1:2] ~ dmnorm(mu[1:2], cov = S[1:2, 1:2])
+    logit(p) <- a
+    y ~ dbern(p)
+  })
+  S <- matrix(c(2, 0.5, 0.5, 1), 2)
+  m <- nimbleModel(code,
+    constants = list(mu = c(0, 0)), data = list(y = 1, S = S),
+    inits = list(a = 0.5, v = c(1, 2))
+  )
+  for (a in c(-1, 0.5, 2)) {
+    m$a <- a
+    m$calculate()
+    v <- c(1, 2)
+    expect_equal(m$b1, (a >= 0) + (a == 0.5) + a^2 + sum(v * v))
+    expect_equal(
+      m$b2,
+      2 * plogis(a) + qlogis(0.3) + qnorm(0.3) + 2 * pnorm(a) + log(-log(1 - 0.3)) +
+        (1 - exp(-exp(a))) + lgamma(2.5) + lfactorial(3) + a^3
+    )
+    expect_equal(m$p, plogis(a))
+  }
+  expect_equal(m$d, as.numeric(determinant(S)$modulus))
+  expect_equal(m$Sinv, solve(S))
+})
+
+test_that("expressions of constants using math functions are pre-evaluated", {
+  code <- nimbleCode({
+    y1 ~ dnorm(exp(c1), 1)
+    y2 ~ dnorm(pow(c1, 2), 1)
+    y3 ~ dnorm(step(c1) + equals(c1, 2), 1)
+    y4 ~ dnorm(logit(c1 / 4) + cube(c1), 1)
+  })
+  m <- nimbleModel(code, constants = list(c1 = 2), inits = list(y1 = 0, y2 = 0, y3 = 0, y4 = 0))
+  declCode <- lapply(m$modelDef$declInfo, function(d) d$code)
+  expect_length(declCode, 4) # Nothing lifted.
+  means <- sapply(declCode, function(x) x[[3]]$mean)
+  expect_equal(means, c(exp(2), 4, 2, 8))
+})
+
+test_that("user-written chol and inverse in multivariate arguments are lifted", {
+  code <- nimbleCode({
+    v[1:2] ~ dmnorm(mu[1:2], cholesky = chol(S[1:2, 1:2]), prec_param = 0)
+    w[1:2] ~ dmnorm(mu[1:2], cov = inverse(S[1:2, 1:2]))
+  })
+  S <- matrix(c(2, 0.5, 0.5, 1), 2)
+  m <- withNimbleModelOptions(
+    list(disallowMultivariateArgumentExpressions = FALSE),
+    nimbleModel(code, constants = list(mu = c(0, 0)), data = list(S = S), inits = list(v = c(1, 2), w = c(1, 2)))
+  )
+  declInfo <- m$modelDef$declInfo
+  liftedRHS <- lapply(declInfo[grepl("^lifted_", sapply(declInfo, function(d) d$targetNodeName))], function(d) d$valueExpr)
+  # `cov = inverse(S)` is reparameterized to `cholesky = chol(inverse(S))`, which is lifted as a whole.
+  expect_identical(liftedRHS, list(quote(nChol(S[1:2, 1:2])), quote(nChol(nInverse(S[1:2, 1:2])))))
+  expect_false("nChol" %in% all.names(declInfo[[2]]$valueExpr))
+  m$calculate()
+  dmnormLogDensity <- function(x, Sigma) {
+    U <- chol(Sigma)
+    sum(dnorm(backsolve(U, x, transpose = TRUE), log = TRUE)) - sum(log(diag(U)))
+  }
+  expect_equal(m$calculate("v[1:2]"), dmnormLogDensity(c(1, 2), S))
+  expect_equal(m$calculate("w[1:2]"), dmnormLogDensity(c(1, 2), solve(S)))
+})