From 48740992fef401589e2548a351cdacb6f88a4f6b Mon Sep 17 00:00:00 2001 From: perrydv Date: Fri, 6 Nov 2026 05:31:23 -0800 Subject: [PATCH] update support for (mostly) simple math ops. Many move to nCompiler, at least for now. --- nimbleModel/R/getSymbolicParentNodes.R | 18 +++-- nimbleModel/R/mathCompatibility.R | 46 ++---------- nimbleModel/R/modelDef.R | 57 +++++++++++---- .../tests/testthat/test-mathFunctions.R | 70 +++++++++++++++++++ 4 files changed, 128 insertions(+), 63 deletions(-) create mode 100644 nimbleModel/tests/testthat/test-mathFunctions.R diff --git a/nimbleModel/R/getSymbolicParentNodes.R b/nimbleModel/R/getSymbolicParentNodes.R index 7b5c4ff..54dcf3d 100644 --- a/nimbleModel/R/getSymbolicParentNodes.R +++ b/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. diff --git a/nimbleModel/R/mathCompatibility.R b/nimbleModel/R/mathCompatibility.R index 09308db..fa247c6 100644 --- a/nimbleModel/R/mathCompatibility.R +++ b/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`). diff --git a/nimbleModel/R/modelDef.R b/nimbleModel/R/modelDef.R index d9670ff..ed20204 100644 --- a/nimbleModel/R/modelDef.R +++ b/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. diff --git a/nimbleModel/tests/testthat/test-mathFunctions.R b/nimbleModel/tests/testthat/test-mathFunctions.R new file mode 100644 index 0000000..f5b8a6d --- /dev/null +++ b/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))) +})