Rcpp: retrieve and replace off-diagonal values of a square matrix - r

Question

Using Rcpp/Armadillo, how can I efficiently extract/replace off-diagonal values of a square matrix? In R it can be achieved using: old_values = A[row(A) == (col(A) - k)]; A[row(A) == (col(A) - k)] = new_values. Using Armadillo a for-loop (see below) can be used to achive this goal. But is there a simpler way to write the code? As I need to do this operation for all k of a big matrix (>10000 rows, > 10000 columns), efficiency is better to be considered. Here is a reproducible example:
A = matrix(1:25, 5, 5)
A[row(A) == (col(A) - 3)] # extract the 3rd off-diagnal values
A[row(A) == (col(A) - 2)] = -5 of # replace the 2nd off-diagnal values with -5
The cpp code using a for loop:
arma::vec retrieve_off_diag_values( arma::mat A, unsigned k )
{
unsigned n_cols = A.n_cols;
arma::vec off_diag_values(n_cols - k);
for( unsigned i=0; i <(n_cols - k); i++ )
{
off_diag_values(i) = A(i, i+k);
}
return off_diag_values;
}

Answer 1

To extract values from a specified diagonal into a vector, where k < 0 indicates subdiagonals, k = 0 indicates the main diagonal, and k > 0 indicates superdiagonals:
#include<RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::export]]
arma::vec diag_get(const arma::mat& X, int k) // note the 'const' and '&'
{
return X.diag(k);
}
To set values on a diagonal to a specific value:
// [[Rcpp::export]]
void diag_fill(arma::mat& X, int k, double value) // note the '&' character
{
X.diag(k).fill(value);
}
To change instances of a specific value on a diagonal with another value:
// [[Rcpp::export]]
void diag_change(arma::mat& X, int k, double old_value, double new_value)
{
X.diag(k).replace(old_value, new_value);
}

Answer 2

Retrieving off-diagonal
You can use armadillo's .diag() member function with an index k to retrieve off-diagonal.
Cases:
if k == 0 (default), then the main diagonal.
else if k < 0, then a lower triangular diagonal.
else then an upper triangular diagonal.
Example:
#include<RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::export]]
arma::vec offdiag_extract(arma::mat& A, int k) {
return A.diag(k);
}
Test:
A = matrix(1:25, 5, 5)
offdiag_extract(A, 3)
# [,1]
# [1,] 16
# [2,] 22
Replacing off-diagonal
Edit: This section has been updated thanks to #mtall's point regarding behavior provided by other member functions.
However, .diag()= can only be used to save into the main diagonal. To ensure the replacement being viable for all diagonals, you would need to chain the .diag() member function with .fill(value), e.g.
#include<RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::export]]
arma::mat offdiag_fill_arma(arma::mat& A, int k, double replace_value) {
A.diag(k).fill(replace_value);
return A;
}
Test:
offdiag_fill_arma(A, 2, 4)
# [,1] [,2] [,3] [,4] [,5]
# [1,] 1 6 4 16 21
# [2,] 2 7 12 4 22
# [3,] 3 8 13 18 4
# [4,] 4 9 14 19 24
# [5,] 5 10 15 20 25
Implementing off-diagonal replacement
In short, one could implement off-diagonal replacement using a single for loop with the appropriate k offset.
#include<RcppArmadillo.h>
// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::export]]
arma::mat offdiag_replace(arma::mat& A, int k, double replace_val = -5) {
// Determine whether to go over upper or lower diagonal
unsigned int row_offset = (k < 0) ? -k : 0;
unsigned int col_offset = (k > 0) ? k : 0;
// Compute total number of elements
unsigned int N = std::min(A.n_rows - row_offset, A.n_cols - col_offset);
// Loop over diagonal
for(unsigned int i = 0; i < N; ++i) {
unsigned int row = i + row_offset;
unsigned int col = i + col_offset;
// Disregard bounds checks with .at()
A.at(row,col) = replace_val;
}
return A;
}
Test:
offdiag_replace(A, 2, 4)
# [,1] [,2] [,3] [,4] [,5]
# [1,] 1 6 4 16 21
# [2,] 2 7 12 4 22
# [3,] 3 8 13 18 4
# [4,] 4 9 14 19 24
# [5,] 5 10 15 20 25