R: Changing diagonal of sparse matrix is very slow

Question

I hava a sparse matrix with zeros on the main diagonal that I want to change to ones, but compared to a QR-decomposition the operation is very very slow:

mat <- matrix(c(0,1,1,1,0,1,1,1,0),ncol=3)
mat1 <- Matrix::bdiag(mat,mat,mat)
mat2 <- Matrix::bdiag(mat,mat,mat)
identity_mat <- Matrix::Diagonal(9)

microbenchmark::microbenchmark(
  qr(mat1),
  Matrix::diag(mat2) <- 1,
  mat1 + identity_mat
)

results in

Unit: microseconds
                    expr      min        lq       mean    median        uq      max neval
                qr(mat1)   55.825   69.0080   79.16561   72.9365   85.6095  149.676   100
 Matrix::diag(mat2) <- 1  302.172  326.2365  379.60509  364.1985  401.8005  756.477   100
     mat1 + identity_mat 1714.578 1762.8665 2006.50270 1974.4125 2073.1795 6671.644   100

How can I set the diagonal to ones faster?

Answer 1

Here's a solution using RcppArmadillo:

Rcpp::cppFunction(code="
arma::sp_mat set_unit_diagonal(arma::sp_mat& A) {
      A.diag().ones();
    return A;
}", depends="RcppArmadillo")

mat <- matrix(c(0,1,1,1,0,1,1,1,0),ncol=3)
mat1 <- Matrix::bdiag(mat,mat,mat)
mat2 <- Matrix::bdiag(mat,mat,mat)
mat3 <- Matrix::bdiag(mat,mat,mat)
identity_mat <- Matrix::Diagonal(9)

The solution is very quick:

Unit: microseconds
                    expr      min        lq       mean   median        uq       max neval
                qr(mat1)   52.374   62.4805   69.12798   65.747   75.4455   110.630   100
 Matrix::diag(mat2) <- 1  272.026  289.9080  323.98992  300.557  358.5355   419.486   100
     mat1 + identity_mat 1543.191 1620.2835 1913.68513 1637.990 1970.8930 13572.716   100
 set_unit_diagonal(mat3)    7.426   11.9100   14.00686   13.970   15.5955    26.484   100