Support for multivariate distribution calls

nCompiler/R/Rcpp_nCompiler_plugin.R

@@ -45,7 +45,8 @@ make_nCompiler_plugin <- function(nCompiler_pluginEnv) {
                                  else
                                    "")
     result$env$PKG_CXXFLAGS <- "-std=c++17"
-    result$env$PKG_LIBS <- get_nCompLocal_PKG_LIBS_entry()
+    result$env$PKG_LIBS <- paste(get_nCompLocal_PKG_LIBS_entry(),
+                                "$(LAPACK_LIBS) $(BLAS_LIBS)")
     if(isTRUE(nCompiler_pluginEnv$uses_TBB))
       result$env <- setEnvTBB(result$env)
     ## Makevars doesn't work for sourceCpp plugins
@@ -73,7 +74,8 @@ make_nCompiler_Eigen_plugin <- function(nCompiler_pluginEnv) {
                                      else
                                        "")
     result$env$PKG_CXXFLAGS <- "-std=c++17"
-    result$env$PKG_LIBS <- get_nCompLocal_PKG_LIBS_entry()
+    result$env$PKG_LIBS <- paste(get_nCompLocal_PKG_LIBS_entry(),
+                                "$(LAPACK_LIBS) $(BLAS_LIBS)")
     if(isTRUE(get_nOption('compilerOptions')$throwEigenErrors)) {
       # replace include directives to enable Eigen errors
       #preamble = system.file(file.path('include', 'nCompiler',

nCompiler/R/compile_aaa_operatorLists.R

@@ -1136,6 +1136,111 @@ updateOperatorDef('rt_nonstandard', 'matchDef', val = function(n, df = 1, mu = 0
 updateOperatorDef('runif', 'matchDef', val = function(n, min = 0, max = 1) {})
 updateOperatorDef('rweibull', 'matchDef', val = function(n, shape, scale = 1) {})
 
+assignOperatorDef(
+  c("dmnorm_chol"),
+  list(
+    matchDef = function(x, mean, cholesky, prec_param = TRUE, log = FALSE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist'
+    )
+  )
+)
+
+assignOperatorDef(
+  c("rmnorm_chol"),
+  list(
+    matchDef = function(n = 1, mean, cholesky, prec_param = TRUE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist',
+      return_nDim = 1
+    )
+  )
+)
+
+assignOperatorDef(
+  c("dmvt_chol"),
+  list(
+    matchDef = function(x, mu, cholesky, df, prec_param = TRUE, log = FALSE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist'
+    )
+  )
+)
+
+assignOperatorDef(
+  c("rmvt_chol"),
+  list(
+    matchDef = function(n = 1, mu, cholesky, df, prec_param = TRUE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist',
+      return_nDim = 1
+    )
+  )
+)
+
+assignOperatorDef(
+  c("dwish_chol"),
+  list(
+    matchDef = function(x, cholesky, df, scale_param = TRUE, log = FALSE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist'
+    )
+  )
+)
+
+assignOperatorDef(
+  c("rwish_chol"),
+  list(
+    matchDef = function(n = 1, cholesky, df, scale_param = TRUE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist',
+      return_nDim = 2
+    )
+  )
+)
+
+assignOperatorDef(
+  c("ddirch"),
+  list(
+    matchDef = function(x, alpha, log = FALSE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist'
+    )
+  )
+)
+
+assignOperatorDef(
+  c("rdirch"),
+  list(
+    matchDef = function(n = 1, alpha) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist',
+      return_nDim = 1
+    )
+  )
+)
+
+assignOperatorDef(
+  c("dmulti"),
+  list(
+    matchDef = function(x, size, prob, log = FALSE) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist'
+    )
+  )
+)
+
+assignOperatorDef(
+  c("rmulti"),
+  list(
+    matchDef = function(n = 1, size, prob) {},
+    labelAbstractTypes = list(
+      handler = 'MVDist',
+      return_nDim = 1
+    )
+  )
+)
+
 assignOperatorDef(
   c('nSwitch'),
   list(

nCompiler/R/compile_labelAbstractTypes.R

@@ -568,6 +568,20 @@ inLabelAbstractTypesEnv(
   }
 )
 
+inLabelAbstractTypesEnv(
+  # Used for both multivariate density functions (e.g. dmnorm_chol), which
+  # return a scalar (nDim = 0, the default via setReturn_nDim), and their
+  # random-generation counterparts (e.g. rmnorm_chol), which return a
+  # vector -- set via `labelAbstractTypes = list(handler = 'MVDist',
+  # return_nDim = 1)` on the latter's operatorDef.
+  MVDist <- function(code, symTab, auxEnv, handlingInfo) {
+    inserts <- recurse_labelAbstractTypes(code, symTab, auxEnv, handlingInfo)
+    code$type <- symbolBasic$new(type = 'double',
+                                  nDim = setReturn_nDim(handlingInfo, 0))
+    invisible(inserts)
+  }
+)
+
 inLabelAbstractTypesEnv(
   nChol <- function(code, symTab, auxEnv, handlingInfo) {
     inserts <- recurse_labelAbstractTypes(code, symTab, auxEnv, handlingInfo)

nCompiler/R/cppDefs_nClass.R

@@ -12,11 +12,13 @@ nClassBaseClass_init_impl <- function(cppDef) {
   cppDef$Hpreamble <- c(cppDef$Hpreamble,
                         "#define NCOMPILER_USES_EIGEN",
                         "#define NCOMPILER_USES_NCPPVEC",
+                        "#define NCOMPILER_USES_MVDISTS",
                         "#define USES_NCOMPILER")
   cppDef$CPPpreamble <- pluginIncludes
   cppDef$CPPpreamble <- c(cppDef$CPPpreamble,
                         "#define NCOMPILER_USES_EIGEN",
                         "#define NCOMPILER_USES_NCPPVEC",
+                        "#define NCOMPILER_USES_MVDISTS",
                         "#define USES_NCOMPILER")
 
   cppDef$Hincludes <- c(cppDef$Hincludes,

nCompiler/R/cppDefs_nFunction.R

@@ -9,16 +9,21 @@ cpp_nFunctionClass_init_impl <- function(cppDef) {
   cppDef$Hpreamble <- c(cppDef$Hpreamble,
                         "#define NCOMPILER_USES_EIGEN",
                         "#define NCOMPILER_USES_NCPPVEC",
+                        "#define NCOMPILER_USES_MVDISTS",
                         "#define USES_NCOMPILER")
   ## handler nCppVec in labelAbstractTypes does record in auxEnv if an
   ## explicit call to nCppVec() was uses. That is the beginning of a smarter
   ## system for determining what #include (via #define) components are
   ## really needed. But I punted on further extension for now and
   ## simply tacked NCPPVEC onto the universal set of includes for now.
+  ## NCOMPILER_USES_MVDISTS (multivariate distributions, e.g. dmnorm_chol)
+  ## is likewise just tacked on universally for now; a future pass could
+  ## set it only when generated code actually calls one.
   cppDef$CPPpreamble <- pluginIncludes
   cppDef$CPPpreamble <- c(cppDef$CPPpreamble,
                           "#define NCOMPILER_USES_EIGEN",
                           "#define NCOMPILER_USES_NCPPVEC",
+                          "#define NCOMPILER_USES_MVDISTS",
                           "#define USES_NCOMPILER")
   cppDef$Hincludes <- c(cppDef$Hincludes)#,
   ##                      nCompilerIncludeFile("nCompiler_omnibus_first_h.h"))

nCompiler/inst/include/nCompiler/dists/ddirch.h

@@ -0,0 +1,54 @@
+// The ddirch (Dirichlet) density kernel (raw double* arguments), ported
+// from nimble's dists.cpp.
+//
+// Header-only and self-contained for the same reason as dmnorm_chol.h: the
+// package DLL and code nCompile generates and compiles on a user's machine
+// are separate shared objects that cannot reliably link against each other,
+// so each side compiles its own copy of this source instead.
+//
+// Unlike the other dist kernels ported so far, ddirch never mutates x or
+// alpha (no BLAS/LAPACK, no in-place overwrite), so there's no
+// overwrite_inputs parameter here, and the Eigen::Tensor wrapper
+// (ddirch_tensor.h) can use the no-copy reference path for both arguments.
+//
+// See also ddirch_tensor.h for an Eigen::Tensor-argument wrapper around
+// this kernel, for use from nCompile-generated code.
+
+#ifndef _NCOMPILER_DISTS_DDIRCH
+#define _NCOMPILER_DISTS_DDIRCH
+
+#include <cfloat> // DBL_EPSILON
+#include <Rmath.h>
+#include "dists_utils.h"
+
+inline double ddirch(double* x, double* alpha, int K, int give_log) {
+  double sumAlpha(0.0);
+  double sumX(0.0);
+  double dens(0.0);
+
+  if (R_IsNA_ANY(x, K) || R_IsNA_ANY(alpha, K))
+    return NA_REAL;
+  if (R_IsNaN_ANY(x, K) || R_IsNaN_ANY(alpha, K))
+    return R_NaN;
+
+  for(int i = 0; i < K; i++) {
+    if(alpha[i] <= 0.0) ML_ERR_return_NAN;
+    if(x[i] < 0.0 || x[i] > 1.0) return give_log ? R_NegInf : 0.0;
+    // Rf_lgammafn (not the "lgammafn" convenience macro) is used here
+    // deliberately: Rcpp's sugar layer #undefs that macro (Rcpp/sugar/
+    // undoRmath.h, to avoid polluting the global namespace with names like
+    // beta/choose/gamma), so this header needs to work whether or not it's
+    // included after Rcpp.h has already run.
+    dens += (alpha[i]-1) * log(x[i]) - Rf_lgammafn(alpha[i]);
+    sumAlpha += alpha[i];
+    sumX += x[i];
+  }
+  if(sumX > 1.0 + 10*DBL_EPSILON || sumX < 1.0 - 10*DBL_EPSILON) {
+    return give_log ? R_NegInf : 0.0;
+  }
+
+  dens += Rf_lgammafn(sumAlpha);
+  return give_log ? dens : exp(dens);
+}
+
+#endif // _NCOMPILER_DISTS_DDIRCH

nCompiler/inst/include/nCompiler/dists/ddirch_tensor.h

@@ -0,0 +1,36 @@
+// Eigen::Tensor-argument wrapper around the ddirch density kernel (see
+// ddirch.h). Templated code like this can only ever be header-only -- the
+// compiler must see the definition wherever it's instantiated -- so this
+// lives in inst/include regardless of the package-DLL/on-the-fly-
+// compilation split that motivates ddirch.h itself.
+
+#ifndef _NCOMPILER_DISTS_DDIRCH_TENSOR
+#define _NCOMPILER_DISTS_DDIRCH_TENSOR
+
+#include <unsupported/Eigen/CXX11/Tensor>
+#include "dists_tensor_utils.h"
+#include "ddirch.h"
+
+// Eigen::Tensor wrapper for ddirch.
+//
+// Unlike dmnorm_chol/dmvt_chol/dwish_chol, ddirch never mutates x or alpha
+// (no BLAS/LAPACK, no in-place overwrite), so both arguments can use
+// asDenseTensor's no-copy reference path when they're already concrete
+// tensors, and only get materialized when the call site passed a lazy
+// expression.
+template<typename TensorExprX, typename TensorExprAlpha>
+double ddirch(const TensorExprX &x, const TensorExprAlpha &alpha, int give_log) {
+  const auto &xEval = asDenseTensor<1>(x);
+  const auto &alphaEval = asDenseTensor<1>(alpha);
+
+  int K = static_cast<int>(alphaEval.dimension(0));
+
+  // xEval/alphaEval may alias the caller's own tensors (when no copy was
+  // needed above); ddirch's double* signature is non-const only because it
+  // predates const-correctness conventions here, but it never writes
+  // through x or alpha, so this cast is safe.
+  return ddirch(const_cast<double*>(xEval.data()),
+                const_cast<double*>(alphaEval.data()), K, give_log);
+}
+
+#endif // _NCOMPILER_DISTS_DDIRCH_TENSOR

nCompiler/inst/include/nCompiler/dists/dists_post_Rcpp.h

@@ -0,0 +1,33 @@
+// Umbrella header for nCompiler's distribution kernels, for use from
+// nCompile-generated code. Pulled into nCompiler_omnibus.h under
+// #ifdef NCOMPILER_USES_MVDISTS.
+//
+// Individual dist headers are self-contained (each brings its own Eigen
+// includes, etc.), so this file's only job is to gather them in one place.
+// As more dist kernels are added, they just get an #include line here.
+
+#ifndef _NCOMPILER_DISTS_POST_RCPP
+#define _NCOMPILER_DISTS_POST_RCPP
+
+#include <nCompiler/dists/dmnorm_chol.h>
+#include <nCompiler/dists/dmnorm_chol_tensor.h>
+#include <nCompiler/dists/rmnorm_chol.h>
+#include <nCompiler/dists/rmnorm_chol_tensor.h>
+#include <nCompiler/dists/dmvt_chol.h>
+#include <nCompiler/dists/dmvt_chol_tensor.h>
+#include <nCompiler/dists/rmvt_chol.h>
+#include <nCompiler/dists/rmvt_chol_tensor.h>
+#include <nCompiler/dists/dwish_chol.h>
+#include <nCompiler/dists/dwish_chol_tensor.h>
+#include <nCompiler/dists/rwish_chol.h>
+#include <nCompiler/dists/rwish_chol_tensor.h>
+#include <nCompiler/dists/ddirch.h>
+#include <nCompiler/dists/ddirch_tensor.h>
+#include <nCompiler/dists/rdirch.h>
+#include <nCompiler/dists/rdirch_tensor.h>
+#include <nCompiler/dists/dmulti.h>
+#include <nCompiler/dists/dmulti_tensor.h>
+#include <nCompiler/dists/rmulti.h>
+#include <nCompiler/dists/rmulti_tensor.h>
+
+#endif // _NCOMPILER_DISTS_POST_RCPP

nCompiler/inst/include/nCompiler/dists/dists_tensor_utils.h

@@ -0,0 +1,62 @@
+// Shared Eigen::Tensor materialization helpers for nCompiler's dist
+// wrappers (e.g. dmnorm_chol_tensor.h, dmvt_chol_tensor.h). Split out into
+// its own header, rather than being duplicated in each *_tensor.h, because
+// multiple such wrappers get included together (via dists_post_Rcpp.h) and
+// would otherwise redefine the same templates in one translation unit.
+
+#ifndef _NCOMPILER_DISTS_TENSOR_UTILS
+#define _NCOMPILER_DISTS_TENSOR_UTILS
+
+#include <unsupported/Eigen/CXX11/Tensor>
+#include <type_traits>
+
+// True only when TensorExpr is exactly Eigen::Tensor<double, Rank>, i.e. an
+// already-materialized, densely-packed tensor rather than a lazy expression
+// (a block, a cast, an arithmetic chain, a TensorMap, etc).
+template<typename TensorExpr, int Rank>
+struct is_dense_double_tensor
+  : std::is_same<TensorExpr, Eigen::Tensor<double, Rank>> {};
+
+// Returns x itself (no copy) when it is already a concrete
+// Eigen::Tensor<double, Rank>; otherwise assigns it into a fresh
+// Eigen::Tensor<double, Rank>, which forces Eigen to evaluate the expression
+// into a contiguous buffer we can call .data() on. The .cast<double>() is
+// required here, not just a safety net: Eigen::Tensor's assignment/
+// construction from another tensor expression requires matching Scalar
+// types, so a non-double input (e.g. Eigen::Tensor<int, Rank>, or a bool
+// expression from a comparison op) would otherwise fail to compile.
+template<int Rank, typename TensorExpr>
+decltype(auto) asDenseTensor(const TensorExpr &x) {
+  if constexpr (is_dense_double_tensor<TensorExpr, Rank>::value) {
+    return (x);
+  } else {
+    return Eigen::Tensor<double, Rank>(x.template cast<double>());
+  }
+}
+
+// Applies nimble's/R's recycling rule for a "mean"-like argument that may
+// be shorter than the target length n: if v has fewer than n elements, its
+// elements are reused cyclically to fill out length n (matching the
+// full_mean/full_mu construction in dists.cpp's C_dmnorm_chol,
+// C_rmnorm_chol, C_dmvt_chol, and C_rmvt_chol). If v already has at least n
+// elements, it's returned unchanged (extra elements beyond n are simply
+// never read by the kernels that consume it, same as the original code).
+//
+// Unlike asDenseTensor, this always returns by value: whether recycling is
+// needed is a runtime property (v's length vs. n), not something knowable
+// from the type system, so there's no way to give the "no recycling
+// needed" case a reference-typed fast path here without splitting the
+// return type per branch.
+inline Eigen::Tensor<double, 1> recycleToLength(const Eigen::Tensor<double, 1> &v, int n) {
+  int m = static_cast<int>(v.dimension(0));
+  if (m >= n) return v;
+  Eigen::Tensor<double, 1> out(n);
+  int j = 0;
+  for (int i = 0; i < n; i++) {
+    out(i) = v(j++);
+    if (j == m) j = 0;
+  }
+  return out;
+}
+
+#endif // _NCOMPILER_DISTS_TENSOR_UTILS