diff --git a/nimbleModel/DESCRIPTION b/nimbleModel/DESCRIPTION index 12acd75..b3e53fe 100644 --- a/nimbleModel/DESCRIPTION +++ b/nimbleModel/DESCRIPTION @@ -42,6 +42,7 @@ Collate: modelFunctions.R modelValues.R model_utils.R + nimCopy.R nodeRules.R opDefs.R options.R diff --git a/nimbleModel/NAMESPACE b/nimbleModel/NAMESPACE index 0e99885..e78e6aa 100644 --- a/nimbleModel/NAMESPACE +++ b/nimbleModel/NAMESPACE @@ -83,6 +83,7 @@ export(messageIfVerbose) export(calc_dmnormAltParams) export(calc_dwishAltParams) +export(nimCopy) export(nimbleCode) export(nimbleModel) export(makeInstrList) diff --git a/nimbleModel/R/copier.R b/nimbleModel/R/copier.R index ac9f39c..e45a7ad 100644 --- a/nimbleModel/R/copier.R +++ b/nimbleModel/R/copier.R @@ -217,13 +217,15 @@ multiCopierMV_nClass <- nCompiler::nClass( #' @export makeMultiCopier <- function(source, nodes, - sourceType = c("model", "modelValues"), ...) { - # This will be called from a nimble2 keyword processor + sourceType = c("model", "modelValues"), + logProb = FALSE, logProbOnly = FALSE, model = NULL, ...) { + # This is called from setup code made by nimble2 keyword processing (for nimCopy and values). + # logProb and logProbOnly are as for nimCopy (see nimCopy.R). For a modelValues source, + # `model` is used to find the logProb elements of nodes. # The ... is to absorb further arguments that at the time of this writing are not fleshed out. - if (is.character(nodes)) { - nodes <- nodes |> lapply(\(x) nimbleModel:::varRangeClass$new(x)) - } sourceType <- match.arg(sourceType) + if (sourceType == "model") model <- source + nodes <- nimCopy_prepareNodes(model, nodes, logProb, logProbOnly) multiCopier <- switch(sourceType, model = multiCopier_nClass$new(), diff --git a/nimbleModel/R/nimCopy.R b/nimbleModel/R/nimCopy.R new file mode 100644 index 0000000..67f9f3b --- /dev/null +++ b/nimbleModel/R/nimCopy.R @@ -0,0 +1,199 @@ +# nimCopy: copy values between models and modelValues objects. +# +# As in nimble: +# - "nodes" are really arbitrary indexed variable ranges and are copied as given, not expanded to nodes. +# If nodesTo is missing, it defaults to nodes. +# - logProb = TRUE also copies the logProb elements of the stochastic nodes in nodes +# (including nodes only partly in nodes). logProbOnly = TRUE copies only those. +# - row is the row of `from` if it is a modelValues object. rowTo is the row of `to` +# if it is a modelValues object, defaulting to row. +# +# A difference from nimble: +# Finding the logProb elements of nodes needs a model, so with logProb or logProbOnly, +# and nodes given, one of from and to must be a model. +# This might necessitate enhancing modelValues to be know about logProb nodes. +# +# In the future, there are several efficiency issues that could be revisited: +#. - Instead of doing a full get for the "to" nodes just to get lengths, calculate these from indices or objects. +#. - Check on when assignment ends up creating whole object copies or not. +#. This may be particularly the case when uncompiled code works with compiled objects. +#. - Consider get-and-set for each, rather than get-all and set-all, but preliminary investigation +#. suggested that overhead of large allocation is minor compared to benefits of the latter. +# +#' Copy values between models and modelValues objects +#' +#' @param from a model or modelValues object to copy from. +#' @param to a model or modelValues object to copy to. +#' @param nodes variables or indexed ranges of variables to copy \code{from}, as a character vector or list +#' of varRanges. If missing or NULL, all variables are copied. +#' @param nodesTo variables or indexed ranges of variables to copy \code{to} If missing or NULL, +#' the same as \code{nodes}. +#' @param row the row of \code{from}, if it is a modelValues object. +#' @param rowTo the row of \code{to}, if it is a modelValues object. Defaults to \code{row}. +#' @param logProb whether to also copy the logProb elements of the stochastic nodes in +#' \code{nodes} (or all logProb variables if \code{nodes} is missing). One can also manually include +#' logProb elements in \code{nodes} or \code{nodesTo}. See details. +#' @param logProbOnly whether to copy only the logProb elements (or variables). +#' +#' @return NULL, invisibly. +#' @export +nimCopy <- function(from, to, nodes = NULL, nodesTo = NULL, row = NA, rowTo = NA, + logProb = FALSE, logProbOnly = FALSE) { + is_fromMV <- inherits(from, "modelValues") + is_toMV <- inherits(to, "modelValues") + if (is_fromMV && is.na(row)) stop("nimCopy: row is needed to copy from a modelValues object.") + if (is_toMV && is.na(rowTo)) { + if (is.na(row)) stop("nimCopy: row or rowTo is needed to copy to a modelValues object.") + rowTo <- row + } + model <- if (!is_fromMV) from else if (!is_toMV) to else NULL + if (is.null(nodes)) { + fromNodes <- nimCopy_prepareNodes(NULL, nimCopy_allVars(from, logProb, logProbOnly)) + } else { + fromNodes <- nimCopy_prepareNodes(model, nodes, logProb, logProbOnly) + } + toNodes <- if (is.null(nodesTo)) fromNodes else nimCopy_prepareNodes(model, nodesTo, logProb, logProbOnly) + + vals <- lapply(fromNodes, \(vr) as.numeric(nimCopy_get(from, is_fromMV, vr, row))) + vals <- unlist(vals, use.names = FALSE) + toLengths <- vapply(toNodes, \(vr) length(nimCopy_get(to, is_toMV, vr, rowTo)), numeric(1)) + if (sum(toLengths) != length(vals)) { + stop("nimCopy: from has ", length(vals), " values to copy, but to has ", sum(toLengths), ".") + } + offset <- 0 + for (i in seq_along(toNodes)) { + nimCopy_set(to, is_toMV, toNodes[[i]], rowTo, vals[offset + seq_len(toLengths[i])]) + offset <- offset + toLengths[i] + } + invisible(NULL) +} + +# The values of varRange vr in a model, or in row `row` of a modelValues object. +nimCopy_get <- function(source, isMV, vr, row) { + var <- vr$varName + value <- if (isMV) source[var, row] else source[[var]] + eval(vr$toExpr(), envir = structure(list(value), names = var), enclos = baseenv()) +} + +# Set the values of varRange vr from v, in flattened (column-major) order. +nimCopy_set <- function(target, isMV, vr, row, v) { + var <- vr$varName + expr <- vr$toExpr() + if (is.name(expr)) { + # A whole variable: keep its dimensions. + value <- if (isMV) target[var, row] else target[[var]] + value[] <- v + } else { + env <- new.env(parent = baseenv()) + env[[var]] <- if (isMV) target[var, row] else target[[var]] + env$LocalV_ <- v + eval(substitute(E <- LocalV_, list(E = expr)), env) + value <- env[[var]] + } + if (isMV) target[var, row] <- value else target[[var]] <- value + invisible(NULL) +} + +# All variables of a model or modelValues object, for nimCopy with nodes missing. +nimCopy_allVars <- function(source, logProb = FALSE, logProbOnly = FALSE) { + vars <- if (inherits(source, "modelValues")) { + modelValues_varNames(source) + } else { + source$getVarNames(includeLogProb = TRUE) + } + isLP <- startsWith(vars, "logProb_") + if (logProbOnly) vars[isLP] else if (logProb) vars else vars[!isLP] +} + +# Make nodes into a list of varRanges, and add the logProb elements of their stochastic nodes +# (logProb = TRUE) or use only those (logProbOnly = TRUE). Elements of nodes that are already +# logProb variables are kept with logProb = TRUE and dropped with logProbOnly = TRUE, as in nimble. +nimCopy_prepareNodes <- function(model, nodes, logProb = FALSE, logProbOnly = FALSE) { + if (inherits(nodes, "varRangeClass")) nodes <- list(nodes) + if (is.character(nodes)) nodes <- lapply(nodes, \(x) varRangeClass$new(x)) + if (!(logProb || logProbOnly)) { + return(nodes) + } + if (is.null(model)) { + stop( + "nimCopy: with logProb = TRUE or logProbOnly = TRUE, from or to must be a model, ", + "to find the logProb elements of the nodes." + ) + } + isLP <- vapply(nodes, \(vr) startsWith(vr$varName, "logProb_"), logical(1)) + lpNodes <- list() + if (any(!isLP)) { + lpNodes <- lapply(logProbNodes(model, nodes[!isLP]), \(x) varRangeClass$new(x)) + } + if (logProbOnly) lpNodes else c(nodes, lpNodes) +} + +# The logProb elements of the stochastic nodes in, or partly in, nodes, as character strings +# like "logProb_x[2:3, 1]". As in genLogProbExpr, node x[i, 1:2] has logProb element +# logProb_x[i, 1]: node-internal ranges collapse to their start. So a block of nodes +# usually gives a block of logProb elements. Other cases give one element per node. +logProbNodes <- function(model, nodes) { + # getNodes takes character node names more generally than varRanges (e.g. a whole variable). + if (!is.character(nodes)) nodes <- vapply(nodes, \(vr) safeDeparse(vr$toExpr()), character(1)) + nodeRanges <- model$getNodes(nodes, stochOnly = TRUE) + ans <- lapply(nodeRanges, \(nr) { + block <- logProb_block(nr) + if (is.null(block)) logProb_per_node(model, nr) else block + }) + unlist(ans, use.names = FALSE) +} + +# The logProb block for a nodeRange whose ranges are each a sequence (or scalar) in one +# index slot, or NULL otherwise. Internal ranges have constant starts in this case, +# since a range that depends on the loop indices is not a sequence range. +logProb_block <- function(nr) { + lpName <- makeLogProbName(nr$varName) + ranges <- nr$indexRanges + if (!length(ranges)) { + return(lpName) + } + isSimple <- vapply(ranges, \(r) inherits(r, c("indexRangeSequenceClass", "indexRangeScalarClass")), logical(1)) + if (!all(isSimple) || !all(lengths(nr$rangeToIndexSlot) == 1)) { + return(NULL) + } + inds <- character(length(ranges)) + for (r in seq_along(ranges)) { + ir <- ranges[[r]] + if (inherits(ir, "indexRangeScalarClass")) { + ind <- ir$value + } else if (nr$boolExternalIndexRanges[r] && ir$end > ir$start) { + ind <- paste0(ir$start, ":", ir$end) + } else { + ind <- ir$start + } + inds[nr$rangeToIndexSlot[[r]]] <- ind + } + paste0(lpName, "[", paste(inds, collapse = ", "), "]") +} + +# The logProb element of each node in a nodeRange, from the declarations' logProb expressions. +logProb_per_node <- function(model, nr) { + instrList <- makeInstrList(model, list(nr)) + if (is.null(instrList)) { + return(NULL) + } + declInfo <- model$modelDef$declInfo + ans <- lapply(instrList$instr_nList, \(instr) { + di <- declInfo[[instr$declID]] + lpExpr <- di$genLogProbExpr() + if (is.null(lpExpr)) { + return(NULL) + } # not stochastic + if (is.name(lpExpr)) { + return(as.character(lpExpr)) + } + indexVarNames <- di$declRule$loopIndexingRule$indexVarNames + indexExprs <- as.list(lpExpr)[-(1:2)] + vapply(seq_len(instr$nNodes), \(k) { + idxEnv <- list2env(as.list(structure(set_instr_idx(instr, k), names = indexVarNames)), parent = baseenv()) + inds <- vapply(indexExprs, \(e) eval(e, idxEnv), numeric(1)) + paste0(as.character(lpExpr[[2]]), "[", paste(inds, collapse = ", "), "]") + }, character(1), USE.NAMES = FALSE) + }) + unlist(ans, use.names = FALSE) +} diff --git a/nimbleModel/tests/testthat/test-nimCopy.R b/nimbleModel/tests/testthat/test-nimCopy.R new file mode 100644 index 0000000..7300330 --- /dev/null +++ b/nimbleModel/tests/testthat/test-nimCopy.R @@ -0,0 +1,97 @@ +library(nimbleModel) +library(testthat) +library(nCompiler) + +# Uncompiled nimCopy and its logProb expansion. Compiled nimCopy (in nimble2 nimbleFunctions) +# uses the same expansion through makeMultiCopier and is tested in nimble2. + +make_nimCopy_model <- function(y = matrix(1:12, 3), mu = 0.2, x = matrix(1:6, 3)) { + code <- nimbleCode({ + mu ~ dnorm(0, sd = 2) + for (i in 1:3) { + x[i, 1:2] ~ dmnorm(m0[1:2], pr[1:2, 1:2]) + for (j in 1:4) y[i, j] ~ dnorm(x[i, 1], 1) + } + w <- mu + 1 + for (k in 1:2) q[1:2, k] ~ dmnorm(m0[1:2], pr[1:2, 1:2]) + }) + m <- nimbleModel(code, + data = list(y = y), + inits = list(mu = mu, m0 = c(0, 0), pr = diag(2), x = x, q = diag(2)) + ) + m$calculate() + m +} + +test_that("logProbNodes gives logProb blocks, including for partly covered nodes", { + m <- make_nimCopy_model() + expect_identical(nimbleModel:::logProbNodes(m, "mu"), "logProb_mu") + expect_identical(nimbleModel:::logProbNodes(m, "x"), "logProb_x[1:3, 1]") + expect_identical(nimbleModel:::logProbNodes(m, "x[2:3, 1]"), "logProb_x[2:3, 1]") # partly covered + expect_identical(nimbleModel:::logProbNodes(m, "q[2, 1:2]"), "logProb_q[1, 1:2]") # internal range first + expect_identical(nimbleModel:::logProbNodes(m, "y[1:2, 2:3]"), "logProb_y[1:2, 2:3]") + expect_identical(nimbleModel:::logProbNodes(m, "y"), "logProb_y[1:3, 1:4]") + expect_identical(nimbleModel:::logProbNodes(m, c("w", "mu")), "logProb_mu") # deterministic w has none + expect_length(nimbleModel:::logProbNodes(m, "w"), 0) +}) + +test_that("logProbNodes gives one element per node when there is no block", { + code <- nimbleCode({ + for (i in 1:3) z[i, i:(i + 1)] ~ dmnorm(m0[1:2], pr[1:2, 1:2]) # internal range depends on i + for (i in 1:3) v[2 * i] ~ dnorm(0, 1) + }) + m <- suppressMessages(nimbleModel(code, inits = list(m0 = c(0, 0), pr = diag(2)))) + expect_identical( + nimbleModel:::logProbNodes(m, "z"), + c("logProb_z[1, 1]", "logProb_z[2, 2]", "logProb_z[3, 3]") + ) + expect_identical(nimbleModel:::logProbNodes(m, "v"), c("logProb_v[2]", "logProb_v[4]", "logProb_v[6]")) +}) + +test_that("uncompiled nimCopy copies between models and modelValues, with logProb and logProbOnly", { + m <- make_nimCopy_model() + mv <- modelValues(m)$new() + resize(mv, 2) + + # model to modelValues, with logProb + nimCopy(m, mv, nodes = c("mu", "x[2:3, 1:2]"), row = 2, logProb = TRUE) + expect_equal(mv["mu", 2], 0.2) + expect_equal(mv["x", 2], rbind(0, m$x[2:3, ])) + expect_equal(mv["logProb_mu", 2], m$logProb_mu) + expect_equal(mv["logProb_x", 2], cbind(c(0, m$logProb_x[2:3, 1]), 0)) + expect_equal(mv["x", 1], matrix(0, 3, 2)) # other row untouched + + # nodes missing: all variables, or only logProb variables + mv2 <- modelValues(m)$new() + resize(mv2, 1) + nimCopy(m, mv2, row = 1, logProbOnly = TRUE) + expect_equal(mv2["logProb_y", 1], m$logProb_y) + expect_equal(mv2["x", 1], matrix(0, 3, 2)) + nimCopy(m, mv2, row = 1) + expect_equal(mv2["x", 1], m$x) + expect_equal(mv2["pr", 1], diag(2)) + + # modelValues to model, with nodesTo + m3 <- make_nimCopy_model(y = matrix(0, 3, 4), mu = 0, x = matrix(0, 3, 2)) + nimCopy(mv, m3, nodes = "x[2:3, 1:2]", nodesTo = "x[1:2, 1:2]", row = 2) + expect_equal(m3$x, rbind(m$x[2:3, ], 0)) + + # model to model, with logProb; a whole variable keeps its dimensions + nimCopy(m, m3, nodes = "y", logProb = TRUE) + expect_equal(m3$y, m$y) + expect_equal(m3$logProb_y, m$logProb_y) + + # modelValues to modelValues needs a model for logProb expansion of nodes, but not with nodes missing + expect_error(nimCopy(mv, mv2, nodes = "x", row = 2, rowTo = 1, logProb = TRUE), "must be a model") + nimCopy(mv, mv2, row = 2, rowTo = 1, logProb = TRUE) + expect_equal(mv2["x", 1], mv["x", 2]) + expect_equal(mv2["logProb_x", 1], mv["logProb_x", 2]) + + # logProb from a modelValues to a model uses the model to find the logProb elements + m3$logProb_mu <- 0 + nimCopy(mv, m3, nodes = "mu", row = 2, logProb = TRUE) + expect_equal(m3$logProb_mu, m$logProb_mu) + + expect_error(nimCopy(m, m3, nodes = "x", nodesTo = "mu"), "6 values to copy, but to has 1") + expect_error(nimCopy(mv, m3, nodes = "x"), "row is needed") +})