RcppArmadillo: Lambda expression with each_slice - c++

Question

I have a three dimensional array with positive definite matrices and I would like to obtain an array of the same size with the Cholesky factors of all matrices. I am using the Armadillo library and the cube type, for which there is the convenient function each_slice which I'm trying to use. But I am not getting the lambda expression to work correctly, so hopefully someone can help me and point out my mistake.
Here is a minimal example:
// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>
// [[Rcpp::export]]
arma::cube chol_array(arma::cube Sigma) {
arma::cube Sigma_chol = Sigma;
Sigma_chol.each_slice([](arma::mat X) {return arma::chol(X);});
return Sigma_chol;
}
// [[Rcpp::export]]
arma::cube chol_array2(arma::cube Sigma) {
arma::cube Sigma_chol(size(Sigma));
for (arma::uword i = 0; i < Sigma.n_slices; i++) {
Sigma_chol.slice(i) = arma::chol(Sigma.slice(i));
}
return Sigma_chol;
}
/*** R
Sigma <- array(crossprod(matrix(rnorm(9), 3, 3)), dim = c(3, 3, 2))
chol_array(Sigma)
chol_array2(Sigma)
*/
The function chol_array2 does the job, but chol_array just returns the original matrices. What am I missing?

Answer 1

The issue here is the lack of references in the .each_slice() call. Armadillo's use of lambda expressions require references to update the object and not a return statement. In particular, we have:
For form 3:
apply the given lambda_function to each slice; the function must accept a reference to a Mat object with the same element type as the underlying cube
So, change:
Sigma_chol.each_slice([](arma::mat X) {return arma::chol(X);});
to:
Sigma_chol.each_slice([](arma::mat& X) {X = arma::chol(X);});
Fixed Code
// [[Rcpp::depends(RcppArmadillo)]]
#include <RcppArmadillo.h>
// Enable lambda expressions....
// [[Rcpp::plugins(cpp11)]]
// [[Rcpp::export]]
arma::cube chol_array(arma::cube Sigma) {
arma::cube Sigma_chol = Sigma;
// NOTE: the '&' and saving _back_ into the object are crucial
Sigma_chol.each_slice( [](arma::mat& X) { X = arma::chol(X); } );
return Sigma_chol;
}
Test code
set.seed(1113)
Sigma = array(crossprod(matrix(rnorm(9), 3, 3)), dim = c(3, 3, 2))
all.equal(chol_array(Sigma), chol_array2(Sigma))
# [1] TRUE