MCMC Sampling of partially-observed multivariate normal nodes (#1612)

* added getMixedDataNodeNames method for models

* removed comments from NEWS

* partial_mvn (#1543)

`partial_mvn` sampler, to have the code here for posterity.

* fixed unlist for graph node ids

* updates

* updates

* updates

* updates

* removed extraneous tests of partial_mvn sampler

* updates per comments on PR

* partial_mvn sampler can assign barker sampler

* Use Cholesky of Sigma22 in partial_mvn_pp.

* fixed dimension mistake

* Use Cholesky in sampler_partial_mvn_pp.

* Add more testing for sampler_partial_mvn.

* fix spacing again

* remove extra tests, and fix use of expect_success

* added caching of sigma and mu in partial_mvn_pp sampler

* correction in caching calculations

* Add partial_mvn_pp test.

---------

Co-authored-by: Christopher Paciorek <paciorek@stat.berkeley.edu>

Author: danielturek

packages/RELEASE_INSTRUCTIONS

@@ -85,11 +85,13 @@ R CMD check --as-cran nimble_${VERSION}.tar.gz
 
 Officially this should be done on a development build of R, so this is just a first pass and we check with r-devel via winbuilder or r-hub below.
 
-Recently I've been getting this warning on my Ubuntu machine:
+Recently I've been getting warnings like this on my Ubuntu machine:
 
 ```
 Compilation used the following non-portable flag(s):
   ‘-Wdate-time’ ‘-Werror=format-security’ ‘-Wformat’
+Compilation used the following non-portable flag(s):
+  ‘-mno-omit-leaf-frame-pointer’
 ```
 
 Seems like it might just be my local system as it doesn't show up on winbuilder or r-hub checking.

packages/nimble/R/BUGS_model.R

@@ -499,6 +499,22 @@ Details: Multiple logical input arguments may be used simultaneously.  For examp
                                       }
                                       return(ans)
                                   },
+                                  getMixedDataNodeNames = function(returnType = 'names') {
+                                      multivariateStochBool <- sapply(modelDef$declInfo,
+                                                                      function(di) di$type == 'stoch' && grepl(':', deparse(di$targetExpr)))
+                                      multivariateStochIDsList <- lapply(modelDef$declInfo[multivariateStochBool], `[[`, 'graphIDs')
+                                      multivariateStochIDs <- unlist(multivariateStochIDsList)
+                                      if(length(multivariateStochIDs) == 0)   multivariateStochIDs <- numeric()   ## length=0 case, make a numeric vector
+                                      isDataResult <- isDataFromGraphID(multivariateStochIDs, includeMixed = TRUE)    ## values are in {0, 1, 2}, where 0=FALSE, 1=TRUE, 2=MIXED
+                                      mixedDataIDs <- multivariateStochIDs[isDataResult == 2]
+                                      if(returnType == 'ids')   return(mixedDataIDs)
+                                      if(returnType == 'names') return(modelDef$maps$graphID_2_nodeName[mixedDataIDs])
+                                      stop('returnType argument to getMixedDataNodeNames was invalid')
+                                  },
+                                  isMixedData = function(nodeNames) {
+                                      ids <- expandNodeNames(nodeNames, returnType = 'ids')
+                                      return(isDataFromGraphID(ids, includeMixed = TRUE) == 2)
+                                  },
                                   safeUpdateValidValues = function(validValues, idsVec_only, idsVec_exclude) {
                                       if(!missing(idsVec_only) && !missing(idsVec_exclude)) stop()
                                       if(!missing(idsVec_only)) {
@@ -787,15 +803,23 @@ Details: The variable or node names specified is expanded into a vector of model
                                   		return(isDataFromGraphID(g_id))
                                   },
 
-                                  isDataFromGraphID = function(g_id){
-                                      ## returns TRUE if any elements are flagged as data
+                                  isDataFromGraphID = function(g_id, includeMixed = FALSE) {
+                                      ## default behaviour: returns TRUE if any elements are flagged as data
+                                      ## when includeMixed=TRUE: 0 for FALSE, 1 for TRUE, or 2 for MIXED
                                       nodeNames <- modelDef$maps$graphID_2_nodeName[g_id]
-                                  	ret <- unlist(lapply(as.list(nodeNames),
-                                                  function(nn)
-                                                    return(any(eval(parse(text=nn, keep.source = FALSE)[[1]],
-                                                                                envir=isDataEnv)))))
-                                    if(is.null(ret))   ret <- logical(0)
-                                    return(ret)
+                                      if(includeMixed) {
+                                          f <- function(nn) {
+                                              vals <- eval(parse(text=nn, keep.source = FALSE)[[1]], envir=isDataEnv)
+                                              if(!any(vals))   return(0)
+                                              if(all(vals))    return(1)
+                                              return(2)
+                                          }
+                                      } else {
+                                          f <- function(nn) return(any(eval(parse(text=nn, keep.source = FALSE)[[1]], envir=isDataEnv)))
+                                      }
+                                      ret <- unlist(lapply(as.list(nodeNames), f))
+                                      if(is.null(ret))   ret <- logical(0)
+                                      return(ret)
                                   },
 
                                   getDependenciesList = function(returnNames = TRUE, sort = TRUE) {

packages/nimble/R/MCMC_configuration.R

@@ -245,6 +245,11 @@ print: A logical argument specifying whether to print the monitors and samplers.
             if(missing(nodes)) {
                 nodes <- model$getNodeNames(stochOnly = TRUE, includeData = FALSE, includePredictive = samplePredictiveNodes)
                 # Check of all(model$isStoch(nodes)) is not needed in this case
+                # Check and adds any partially observed nodes
+                mixedDataNodeNames <- model$getMixedDataNodeNames()
+                if(length(mixedDataNodeNames)) {
+                    nodes <- model$topologicallySortNodes(c(nodes, mixedDataNodeNames))
+                }
             } else if(is.null(nodes) || length(nodes)==0) {
                 nodes <- character(0)
             } else   nodes <- filterOutDataNodes(nodes)   ## configureMCMC *never* assigns samplers to data nodes
@@ -431,6 +436,13 @@ For internal use.  Adds default MCMC samplers to the specified nodes.
 
                     ## for multivariate nodes, either add a conjugate sampler, RW_multinomial, or RW_block sampler
                     if(nodeLength > 1) {
+                        if(model$isMixedData(node)) {
+                            if(nodeDist == 'dmnorm') {
+                                thisControlList <- c(controlDefaultsArg, multivariateNodesAsScalars = multivariateNodesAsScalars)
+                                addSampler(target = node, type = 'partial_mvn', control = thisControlList, allowData = TRUE) ; next
+                            }
+                            stop(paste0('The node ', node, ' is partially observed. NIMBLE only handles this case for multivariate normal distibutions.'))
+                        } 
                         if(useConjugacy) {
                             conjugacyResult <- conjugacyResultsAll[[node]]
                             if(!is.null(conjugacyResult)) {
@@ -807,7 +819,9 @@ For internal use only
 
         filterOutDataNodes = function(nodes) {
             nodes <- model$expandNodeNames(nodes)
-            return(nodes[!model$isData(nodes)])
+            # We don't filter out partially observed data nodes
+            gIDs <- model$modelDef$nodeName2GraphIDs(nodes)
+            return(nodes[model$isDataFromGraphID(gIDs, includeMixed=TRUE) != 1])
         },
 
         removeSamplers = function(..., ind, print = FALSE) {

packages/nimble/R/MCMC_samplers.R

@@ -3456,7 +3456,141 @@ sampler_barker <- nimbleFunction(
 )
 
 
-    
+####################################################################
+### partially observed multivariate normal sampler #################
+####################################################################
+
+#' @rdname samplers
+#' @export
+sampler_partial_mvn <- nimbleFunction(
+    name = 'sampler_partial_mvn',
+    contains = sampler_BASE,
+    setup = function(model, mvSaved, target, control) {
+        ## control list extraction
+        multivariateNodesAsScalars <- extractControlElement(control, 'multivariateNodesAsScalars', default = getNimbleOption('MCMCmultivariateNodesAsScalars'))
+        ## node list generation
+        targetAsScalar <- model$expandNodeNames(target, returnScalarComponents = TRUE)
+        isDataBool <- sapply(targetAsScalar, function(targetAsScalar) eval(parse(text = targetAsScalar)[[1]], envir = model$isDataEnv))
+        targetNonDataComponents <- targetAsScalar[!isDataBool]
+        if(length(model$getDependencies(target, self = FALSE, downstream = TRUE, dataOnly = TRUE)) > 0) {
+            predictiveBool <- sapply(targetNonDataComponents, function(n) length(model$getDependencies(n, self = FALSE, downstream = TRUE, dataOnly = TRUE)) == 0)
+            targetNonDataPP <- targetNonDataComponents[ predictiveBool]    ##     predictive
+            targetNonDataNP <- targetNonDataComponents[!predictiveBool]    ## not predictive
+        } else {
+            targetNonDataPP <- targetNonDataComponents                     ## entirely predictive
+            targetNonDataNP <- character()
+        }
+        ## nested function and function list definitions
+        samplerList <- nimbleFunctionList(sampler_BASE)
+        if(length(targetNonDataNP) > 0) {
+            if(multivariateNodesAsScalars) {
+                for(i in seq_along(targetNonDataNP)) {
+                    samplerList[[i]] <- sampler_RW(model, mvSaved, targetNonDataNP[i], control)
+                }
+            } else {
+                if(length(targetNonDataNP) == 1) {
+                    samplerList[[1]] <- sampler_RW(model, mvSaved, targetNonDataNP, control)
+                } else {
+                    if(getNimbleOption('MCMCuseBarkerAsDefaultMV')) {
+                        samplerList[[1]] <- sampler_barker  (model, mvSaved, targetNonDataNP, control)
+                    } else {
+                        samplerList[[1]] <- sampler_RW_block(model, mvSaved, targetNonDataNP, control)
+                    }
+                }
+            }
+        }
+        if(length(targetNonDataPP) > 0) {
+            samplerList[[ length(samplerList)+1 ]] <- sampler_partial_mvn_pp(model, mvSaved, targetNonDataPP)
+        }
+        ## checks
+        if(model$getDistribution(target) != 'dmnorm')   stop('The node ', target, ' is parially observed. NIMBLE only handles this case for multivariate normal distibutions.')
+        if(!model$isMixedData(target))                  stop('The target node ', target, ' is not partially observed.')
+    },
+    run = function() {
+        for(i in seq_along(samplerList)) {
+            samplerList[[i]]$run()
+        }
+    },
+    methods = list(
+        reset = function() {
+            for(i in seq_along(samplerList)) {
+                samplerList[[i]]$reset()
+            }
+        }
+    ) 
+)
+
+sampler_partial_mvn_pp <- nimbleFunction(
+    name = 'sampler_partial_mvn_pp',
+    contains = sampler_BASE,
+    setup = function(model, mvSaved, target) {
+        ## node list generation
+        mvNode <- model$expandNodeNames(target)
+        mvNodeComponents <- model$expandNodeNames(mvNode, returnScalarComponents = TRUE)
+        given <- setdiff(mvNodeComponents, target)
+        calcNodes <- model$getDependencies(target, downstream = TRUE, includePredictive = TRUE)
+        cholNode <- deparse(model$getParamExpr(mvNode, 'cholesky'))
+        meanNode <- deparse(model$getParamExpr(mvNode, 'mean'    ))
+        ## numeric value generation
+        n1 <- length(target)
+        n2 <- length(given)
+        n  <- n1 + n2
+        mu  <- array(0, c(n , 1))
+        mu1 <- array(0, c(n1, 1))
+        mu2 <- array(0, c(n2, 1))
+        Sigma   <- array(0, c(n,  n ))
+        Sigma11 <- array(0, c(n1, n1))
+        Sigma12 <- array(0, c(n1, n2))
+        Sigma21 <- array(0, c(n2, n1))
+        Sigma22 <- array(0, c(n2, n2))
+        tmp <- array(0, c(n2, n1))
+        ind1 <- match(target, mvNodeComponents)
+        ind2 <- match(given,  mvNodeComponents)
+        sgConst <- length(model$getParents(cholNode, self = TRUE, upstream = TRUE, stochOnly = TRUE)) == 0
+        muConst <- length(model$getParents(meanNode, self = TRUE, upstream = TRUE, stochOnly = TRUE)) == 0
+        firstRun <- TRUE
+        ## checks
+        if(length(mvNode) != 1)                         stop('unexpected error in sampler_partial_mvn_pp')
+        if(model$getDistribution(mvNode) != 'dmnorm')   stop('unexpected error in sampler_partial_mvn_pp')
+        if(n != length(mvNodeComponents))               stop('unexpected error in sampler_partial_mvn_pp')
+        if(n1*n2 == 0)                                  stop('unexpected error in sampler_partial_mvn_pp')
+    },
+    run = function() {
+        ## Sigma12 <<- Sigma12 %*% inverse(Sigma22)
+        ## mu1[,1] <<- mu1[,1] + (Sigma12 %*% (values(model,given) - mu2[,1]))[,1]
+        ## Sigma11 <<- Sigma11 - Sigma12 %*% Sigma21
+        if(!sgConst | firstRun) {
+            Sigma <<- model$getParam(mvNode, 'cov' )
+            Sigma11[1:n1,1:n1] <<- Sigma[ind1, ind1]
+            Sigma12[1:n1,1:n2] <<- Sigma[ind1, ind2]
+            Sigma21[1:n2,1:n1] <<- Sigma[ind2, ind1]
+            Sigma22[1:n2,1:n2] <<- Sigma[ind2, ind2]
+            Sigma22 <<- chol(Sigma22)
+            tmp <<- forwardsolve(t(Sigma22), Sigma21)
+            Sigma11 <<- Sigma11 - t(tmp) %*% tmp
+            Sigma11 <<- chol(Sigma11)
+        }
+        if(!muConst | firstRun) {
+            mu[,1] <<- model$getParam(mvNode, 'mean')
+            mu1[,1] <<- mu[ind1,1]
+            mu2[,1] <<- mu[ind2,1]
+        }
+        if(!sgConst | !muConst | firstRun) {
+            mu1[,1] <<- mu1[,1] + (Sigma12 %*% backsolve(Sigma22, forwardsolve(t(Sigma22), values(model,given) - mu2[,1])))[,1]
+        }
+        if(firstRun)   firstRun <<- FALSE
+        values(model, target) <<- rmnorm_chol(1, mu1[,1], Sigma11, prec_param = 0)
+        model$calculate(calcNodes)
+        nimCopy(from = model, to = mvSaved, row = 1, nodes = calcNodes, logProb = TRUE)
+    },
+    methods = list(
+        reset = function() {
+            firstRun <<- TRUE
+        }
+    )
+)
+
+
 #' MCMC Sampling Algorithms
 #'
 #' Details of the MCMC sampling algorithms provided with the NIMBLE MCMC engine; HMC samplers are in the \code{nimbleHMC} package and particle filter samplers are in the \code{nimbleSMC} package. Additional details, including some recommendations for samplers that may perform better than the samplers that NIMBLE assigns by default are provided in Section 7.11 of the User Manual.
@@ -3881,6 +4015,15 @@ sampler_barker <- nimbleFunction(
 #' The posterior_predictive sampler functions by simulating new values for all downstream (dependent) nodes using their conditional distributions, as well as updating the associated model probabilities.  A posterior_predictive sampler will automatically be assigned to all trailing non-data stochastic nodes in a model, or when possible, to any node at a point in the model after which all downstream (dependent) stochastic nodes are non-data.
 #'
 #' The posterior_predictive sampler accepts no control list arguments.
+#' 
+#' @section partial_mvn sampler:
+#'
+#' The partial_mvn sampler is designed to sample multivariate normal distributions that are partially observed.  That is, some dimensions of the target node are observed data values, some dimensions are not data. Sampling is accomplished using either univariate or multivariate random walk Metropolis Hastings of the unobserved dimensions, as determined by the \code{multivariateNodesAsScalars} argument.
+#'
+#' The \code{partial_mvn} sampler accepts the following control list elements:
+#' \itemize{
+#' \item multivariateNodesAsScalars. A logical argument, specifying whether the sampler should sample the unobserved parts of a partially observed node jointly or independently (default = FALSE).
+#' }
 #'
 #' @section RJ_fixed_prior sampler:
 #'
@@ -3896,7 +4039,7 @@ sampler_barker <- nimbleFunction(
 #' 
 #' @name samplers
 #'
-#' @aliases sampler binary categorical prior_samples posterior_predictive RW RW_block RW_multinomial RW_dirichlet RW_wishart RW_llFunction slice AF_slice crossLevel RW_llFunction_block sampler_prior_samples sampler_posterior_predictive sampler_binary sampler_categorical sampler_RW sampler_RW_block sampler_RW_multinomial sampler_RW_dirichlet sampler_RW_wishart sampler_RW_llFunction sampler_slice sampler_AF_slice sampler_crossLevel sampler_RW_llFunction_block CRP CRP_concentration DPmeasure RJ_fixed_prior RJ_indicator RJ_toggled RW_PF RW_PF_block RW_lkj_corr_cholesky sampler_RW_lkj_corr_cholesky RW_block_lkj_corr_cholesky sampler_RW_block_lkj_corr_cholesky sampler_barker barker
+#' @aliases sampler binary categorical prior_samples posterior_predictive RW RW_block RW_multinomial RW_dirichlet RW_wishart RW_llFunction slice AF_slice crossLevel RW_llFunction_block sampler_prior_samples sampler_posterior_predictive sampler_binary sampler_categorical sampler_RW sampler_RW_block sampler_RW_multinomial sampler_RW_dirichlet sampler_RW_wishart sampler_RW_llFunction sampler_slice sampler_AF_slice sampler_crossLevel sampler_RW_llFunction_block CRP CRP_concentration DPmeasure RJ_fixed_prior RJ_indicator RJ_toggled RW_PF RW_PF_block RW_lkj_corr_cholesky sampler_RW_lkj_corr_cholesky RW_block_lkj_corr_cholesky sampler_RW_block_lkj_corr_cholesky sampler_barker barker sampler_partial_mvn partial_mvn
 #'
 #' @examples
 #' ## y[1] ~ dbern() or dbinom():

packages/nimble/R/initializeModel.R

@@ -7,7 +7,7 @@
 #' @author Daniel Turek
 #' @details This nimbleFunction may be used at the beginning of nimble algorithms to perform model initialization.
 #' The intended usage is to specialize an instance of this nimbleFunction in the setup function of an algorithm,
-#' then execute that specialied function at the beginning of the algorithm run function.
+#' then execute that specialized function at the beginning of the algorithm run function.
 #' The specialized function takes no arguments.
 #'
 #' Executing this function ensures that all right-hand-side only nodes have been assigned real values,