Skip to content

Make rmnorm_inv_ld match rmnorm_chol and fix recycling rule #1586

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
paciorek merged 1 commit into from
Aug 14, 2025
Merged

Make rmnorm_inv_ld match rmnorm_chol and fix recycling rule #1586

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
62 changes: 18 additions & 44 deletions packages/nimble/inst/CppCode/dists.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1409,49 +1409,23 @@ SEXP C_dmnorm_inv_ld(SEXP x, SEXP mean, SEXP mat, SEXP inv_ld, SEXP prec_param,
void rmnorm_inv_ld(double *ans, double* mean,
double *mat, double* inv_ld, int n, int prec_param) {
// inv_ld: length n*n+1, first n*n elements are precision matrix (column-major), last is log(det(cov))
// last element of inv_ld is not used.
// Generate standard normal
bool return_NaN = ISNAN_ANY(mean, n);
if(!prec_param) return_NaN |= ISNAN_ANY(inv_ld, n*n) || (!R_FINITE_ANY(inv_ld, n*n));
else return_NaN |= ISNAN(inv_ld[n*n] || ISNAN_ANY(mat, n*n)) ||
(!R_FINITE(inv_ld[n*n]) || !R_FINITE_ANY(mat, n*n));
if (return_NaN) {
for(int j = 0; j < n; j++)
ans[j] = R_NaN;
return;
}

for(int i = 0; i < n; i++)
ans[i] = norm_rand();

// Cholesky decomposition of precision matrix
double* chol_prec = new double[n*n];
if(!prec_param) {
for(int i = 0; i < n*n; i++)
chol_prec[i] = inv_ld[i];
} else {
for(int i = 0; i < n*n; i++)
chol_prec[i] = mat[i];
}
char uplo = 'U';
// Actually inv_ld is not used at all.
// We retain it for now in case we use it in the future.

// Make Cholesky decomposition of covariance or precision.
// We do not check for NaNs here because they will be checked
// in rmnorm_chol.
// dpotrf will simply propagate NaNs.
// (According to chatGPT, the info code is not a reliable way to
// check if NaNs occurred.)
double* chol = new double[n*n];
for(int i = 0; i < n*n; i++)
chol[i] = mat[i]; // cov if prec_param is FALSE, prec if TRUE
char uplo('U');
int info(0);
F77_CALL(dpotrf)(&uplo, &n, chol_prec, &n, &info FCONE);
if (info != 0) {
Rf_error("Error in Cholesky decomposition in rnorm_inv_ld: dpotrf returned info != %d", info);
}
// Solve U^T y = z (z = ans), then U x = y
char transPrec = 'N';
char diag = 'N';
int incx = 1;
int lda(n);
F77_CALL(dtrsv)(&uplo, &transPrec, &diag, &n,
chol_prec, &lda, ans, &incx FCONE FCONE FCONE);

// Add mean
for(int i = 0; i < n; i++)
ans[i] += mean[i];

delete [] chol_prec;
F77_CALL(dpotrf)(&uplo, &n, chol, &n, &info FCONE);
rmnorm_chol(ans, mean, chol, n, prec_param);
delete [] chol;
}

// also substantial editing
Expand Down Expand Up @@ -1496,8 +1470,8 @@ if(!Rf_isReal(mean))
GetRNGstate();

SEXP ans;
PROTECT(ans = Rf_allocVector(REALSXP, n_mean));
rmnorm_inv_ld(REAL(ans), full_mean, c_mat, c_inv_ld, n_mean, c_prec_param);
PROTECT(ans = Rf_allocVector(REALSXP, n_values));
rmnorm_inv_ld(REAL(ans), full_mean, c_mat, c_inv_ld, n_values, c_prec_param);

PutRNGstate();
if(n_mean < n_values)
Expand Down
68 changes: 47 additions & 21 deletions packages/nimble/inst/CppCode/nimDists.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,21 +295,34 @@ void nimArr_rmnorm_chol(NimArr<1, double> &ans, NimArr<1, double> &mean, NimArr<
NimArr<2, double> cholCopy;
double *ansPtr, *meanPtr, *cholPtr;

int n = mean.size();
int n_mean = mean.size();
int n_values = chol.dimSize(0);
if(!ans.isMap()) {
ans.setSize(n);
ans.setSize(n_values);
} else {
if(ans.size() != n) {
_nimble_global_output<<"Error in nimArr_rmnorm_chol: answer size ("<< ans.size() <<") does not match mean size ("<<n<<").\n";
if(ans.size() != n_values) {
_nimble_global_output<<"Error in nimArr_rmnorm_chol: answer size ("<< ans.size() <<") does not match mean size ("<<n_values<<").\n";
nimble_print_to_R(_nimble_global_output);
}
}

// recycling rule for the mean.
if(n_mean < n_values) {
meanPtr = new double[n_values];
int i_mean = 0;
for(int i = 0; i < n_values; i++) {
meanPtr[i] = mean[i_mean++];
if(i_mean == n_mean) i_mean = 0;
}
} else meanPtr = nimArrCopyIfNeeded<1, double>(mean, meanCopy).getPtr();

ansPtr = nimArrCopyIfNeeded<1, double>(ans, ansCopy).getPtr();
meanPtr = nimArrCopyIfNeeded<1, double>(mean, meanCopy).getPtr();
cholPtr = nimArrCopyIfNeeded<2, double>(chol, cholCopy).getPtr();
rmnorm_chol(ansPtr, meanPtr, cholPtr, n, prec_param);
rmnorm_chol(ansPtr, meanPtr, cholPtr, n_values, prec_param);

if(ansPtr != ans.getPtr()) {ans = ansCopy;}
if(n_mean < n_values)
delete [] meanPtr;
// if(ans.isMap()) {
// ans = ansCopy;
// }
Expand Down Expand Up @@ -370,31 +383,44 @@ void nimArr_rmnorm_inv_ld(NimArr<1, double> &ans, NimArr<1, double> &mean,
NimArr<2, double> matCopy;
double *ansPtr, *meanPtr, *matPtr, *inv_ldPtr;

int n = mean.size();
if(!ans.isMap()) {
ans.setSize(n);
} else {
if(ans.size() != n) {
_nimble_global_output<<"Error in nimArr_rmnorm_inv_ld: answer size ("<< ans.size() <<") does not match mean size ("<<n<<").\n";
nimble_print_to_R(_nimble_global_output);
int n_mean = mean.size();
int n_values = mat.dimSize(0);
if(!ans.isMap()) {
ans.setSize(n_values);
} else {
if(ans.size() != n_values) {
_nimble_global_output<<"Error in nimArr_rmnorm_inv_ld: answer size ("<< ans.size() <<") does not match mean size ("<<n_values<<").\n";
nimble_print_to_R(_nimble_global_output);
}
}
if(inv_ld.size() != n*n + 1) {
_nimble_global_output<<"Error in nimArr_rmnorm_inv_ld: inv_ld size ("<<inv_ld.size()<<") does not match n*n+1 ("<<(n*n+1)<<").\n";
nimble_print_to_R(_nimble_global_output);
}
if(mat.dim()[0] != n || mat.dim()[1] != n) {
_nimble_global_output<<"Error in nimArr_rmnorm_inv_ld: mat does not match size of mean.\n";
// Actually inv_ld is not used, so don't bother checking it,
// and thus allow to not form it if not needed otherwise.
// if(inv_ld.size() != n*n + 1) {
// _nimble_global_output<<"Error in nimArr_rmnorm_inv_ld: inv_ld size ("<<inv_ld.size()<<") does not match n*n+1 ("<<(n*n+1)<<").\n";
// nimble_print_to_R(_nimble_global_output);
// }
if(mat.dim()[0] != mat.dim()[1]) {
_nimble_global_output<<"Error in nimArr_rmnorm_inv_ld: mat is not square.\n";
nimble_print_to_R(_nimble_global_output);
}
if(n_mean < n_values) {
meanPtr = new double[n_values];
int i_mean = 0;
for(int i = 0; i < n_values; i++) {
meanPtr[i] = mean[i_mean++];
if(i_mean == n_mean) i_mean = 0;
}
} else meanPtr = nimArrCopyIfNeeded<1, double>(mean, meanCopy).getPtr();

ansPtr = nimArrCopyIfNeeded<1, double>(ans, ansCopy).getPtr();
meanPtr = nimArrCopyIfNeeded<1, double>(mean, meanCopy).getPtr();
matPtr = nimArrCopyIfNeeded<2, double>(mat, matCopy).getPtr();
inv_ldPtr = nimArrCopyIfNeeded<1, double>(inv_ld, inv_ldCopy).getPtr();

rmnorm_inv_ld(ansPtr, meanPtr, matPtr, inv_ldPtr, n, prec_param);
rmnorm_inv_ld(ansPtr, meanPtr, matPtr, inv_ldPtr, n_values, prec_param);

if(ansPtr != ans.getPtr()) {ans = ansCopy;}
if(n_mean < n_values)
delete [] meanPtr;
// if(ans.isMap()) {
// ans = ansCopy;
// }
Expand Down
94 changes: 93 additions & 1 deletion packages/nimble/tests/testthat/test-ADdmnorm.R
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,98 @@ dmnormAD_test2b <- make_AD_test2(
)
dmnormAD_test2b_out <- test_AD2(dmnormAD_test2b)

## Check that rmnorm_inv_ld works and draws the same numbers as from rmnorm_chol
test_that("rmnorm_inv_ld works and matches rmnorm_chol", {
set.seed(1)
cov <- crossprod(matrix(rnorm(25), 5))
PDild_cov <- PDinverse_logdet(cov)
chol_cov <- chol(cov)

prec <- solve(cov)
PDild_prec <- PDinverse_logdet(prec)
chol_prec <- chol(prec)

mu <- c(0.2,0.1,0.3,0.15,0.25)

# PD_ild (PDinverse_logdet) will be ignored for method 2, rmnorm_chol
nf <- nimbleFunction(
run = function(mu=double(1), mat=double(2), PDild=double(1),
prec_param=double(0), method = integer(0)) {
if(method == 1)
ans <- rmnorm_inv_ld(1, mu, mat=mat, inv_ld = PDild, prec_param=prec_param)
if(method == 2)
ans <- rmnorm_chol(1, mu, cholesky = mat, prec_param=prec_param)
return(ans)
returnType(double(1))
}
)
cnf <- compileNimble(nf)

# uncompiled
set.seed(1); x_PLD_0 <- nf(mu, cov, PDild_cov, 0, 1)
set.seed(1); x_chol_0 <- nf(mu, chol_cov, c(0), 0, 2)
set.seed(1); x_PLD_1 <- nf(mu, prec, PDild_prec, 1, 1)
set.seed(1); x_chol_1 <- nf(mu, chol_prec, c(0), 1, 2)

# compiled
set.seed(1); Cx_PLD_0 <- cnf(mu, cov, PDild_cov, 0, 1)
set.seed(1); Cx_chol_0 <- cnf(mu, chol_cov, c(0), 0, 2)
set.seed(1); Cx_PLD_1 <- cnf(mu, prec, PDild_prec, 1, 1)
set.seed(1); Cx_chol_1 <- cnf(mu, chol_prec, c(0), 1, 2)

# look all together
#rbind(x_PLD_0, x_chol_0, x_PLD_1, x_chol_1,
# Cx_PLD_0, Cx_chol_0, Cx_PLD_1, Cx_chol_1)

expect_equal(x_PLD_0, x_chol_0)
expect_equal(x_PLD_1, x_chol_1)
expect_equal(Cx_PLD_0, Cx_chol_0)
expect_equal(Cx_PLD_1, Cx_chol_1)
expect_equal(x_PLD_0, Cx_PLD_0)
expect_equal(x_PLD_1, Cx_PLD_1)

# Check use of recycling rule for the mean
mu <- c(0, 100)

set.seed(1); x_PLD_0 <- nf(mu, cov, PDild_cov, 0, 1)
set.seed(1); x_chol_0 <- nf(mu, chol_cov, c(0), 0, 2)
set.seed(1); x_PLD_1 <- nf(mu, prec, PDild_prec, 1, 1)
set.seed(1); x_chol_1 <- nf(mu, chol_prec, c(0), 1, 2)

# compiled
set.seed(1); Cx_PLD_0 <- cnf(mu, cov, PDild_cov, 0, 1)
set.seed(1); Cx_chol_0 <- cnf(mu, chol_cov, c(0), 0, 2)
set.seed(1); Cx_PLD_1 <- cnf(mu, prec, PDild_prec, 1, 1)
set.seed(1); Cx_chol_1 <- cnf(mu, chol_prec, c(0), 1, 2)

expect_equal(x_PLD_0, x_chol_0)
expect_equal(x_PLD_1, x_chol_1)
expect_equal(Cx_PLD_0, Cx_chol_0)
expect_equal(Cx_PLD_1, Cx_chol_1)
expect_equal(x_PLD_0, Cx_PLD_0)
expect_equal(x_PLD_1, Cx_PLD_1)

mu <- c(0, 10, 20, 30, 40, 50, 60) # too long, only first 5 should be used

set.seed(1); x_PLD_0 <- nf(mu, cov, PDild_cov, 0, 1)
set.seed(1); x_chol_0 <- nf(mu, chol_cov, c(0), 0, 2)
set.seed(1); x_PLD_1 <- nf(mu, prec, PDild_prec, 1, 1)
set.seed(1); x_chol_1 <- nf(mu, chol_prec, c(0), 1, 2)

# compiled
set.seed(1); Cx_PLD_0 <- cnf(mu, cov, PDild_cov, 0, 1)
set.seed(1); Cx_chol_0 <- cnf(mu, chol_cov, c(0), 0, 2)
set.seed(1); Cx_PLD_1 <- cnf(mu, prec, PDild_prec, 1, 1)
set.seed(1); Cx_chol_1 <- cnf(mu, chol_prec, c(0), 1, 2)

expect_equal(x_PLD_0, x_chol_0)
expect_equal(x_PLD_1, x_chol_1)
expect_equal(Cx_PLD_0, Cx_chol_0)
expect_equal(Cx_PLD_1, Cx_chol_1)
expect_equal(x_PLD_0, Cx_PLD_0)
expect_equal(x_PLD_1, Cx_PLD_1)
})

test_that("non-assignment of dmnormAD if cholesky param used", {
code <- nimbleCode({
y[1:3] ~ dmnorm(mu[1:3], cholesky = L[1:3,1:3], prec_param = 0)
Expand Down Expand Up @@ -418,7 +510,7 @@ test_that('Wishart-dmnorm conjugacy', {
m <- nimbleModel(code, data = data)
conf <- configureMCMC(m)
expect_identical(conf$getSamplers()[[1]]$name, "conjugate_dinvwish_dmnormAD_identity")
})
})

# Run various NIMBLE tests that use dmnorm with dmnormAD instead.

Expand Down
Loading