Iterating over arma::mat and retrieving element locations

Question

I'm trying to set the values of an arma::mat element-wise and the value of each element depends on the multi-index (row, column) of each element. Is there a way to retrieve the current location of an element during iteration?

Basically, I'd like to be able to do something like in the sparse matrix iterator, where it.col() and it.row() allow to retrieve the current element's location. For illustration, the example given in the arma::sp_mat iterator documentation) is:

sp_mat X = sprandu<sp_mat>(1000, 2000, 0.1);

sp_mat::const_iterator it     = X.begin();
sp_mat::const_iterator it_end = X.end();

for (; it != it_end; ++it) {
  cout << "val: " << (*it)    << endl;
  cout << "row: " << it.row() << endl;  // only available for arma::sp_mat, not arma::mat
  cout << "col: " << it.col() << endl;  // only available for arma::sp_mat, not arma::mat
}

Of course there are a number of workarounds to get element locations for arma::mat iteration, the most straight-forward ones perhaps being:

• Use nested for-loops over the row and column sizes.
• Use a single for loop and, using the matrix size, transform the iteration number to a row and column index.
• Some form of "zipped" iteration with an object containing or computing the corresponding indices.

However, these seem rather hacky and error-prone to me, because they require to work with the matrix sizes or even do manual index juggling. I'm looking for a cleaner (and perhaps internally optimised) solution. It feels to me like there should be a way to achieve this ...

Apart from the solution used for arma::sp_mat, other such "nice" solution for me would be using .imbue or .for_each but with a functor that accepts not only the element's current value but also its location as additional argument; this doesn't seem to be possible currently.

Answer 1

You seem to have already answered your question yourself. I wish armadillo provided us with a .imbue method overload that received a functor with as many arguments as the dimension of the armadillo object, but currently it only accepts a functor without arguments. Then the probably simplest option (in my opinion) is using a lambda capturing the necessary information, such as the code below

    arma::umat m(3, 3);
    {
        int i = 0;
        m.imbue([&i, num_rows = m.n_rows, num_cols = m.n_cols]() {
            arma::uvec sub = arma::ind2sub(arma::SizeMat{num_rows, num_cols}, i++);
            return 10 * (sub[0] + 1) + sub[1];
        });
    }

In this example each element is computed as 10 times its row index plus its column index. I capture ì here to be the linear indexing and put it as well as the lambda inside curly brackets to delimit its scope.

I also wish I could write something like auto [row_idx, col_idx] = arma::ind2sub( ... ), but unfortunately what ind2sub returns does not work with structured binding.

You can also capture m by const reference and use arma::size(m) as the first argument of arma::ind2sub, if you prefer.

Answer 2

Looking at the armadillo source code, row_col_iterator provides row and column indices of each element. This works like the sparse matrix iterator, but doesn't skip zeros. Adapting your code:

mat X(10,10,fill::randu);

mat::const_row_col_iterator it     = X.begin_row_col();
mat::const_row_col_iterator it_end = X.end_row_col();

for (; it != it_end; ++it) {
  cout << "val: " << (*it)    << endl;
  cout << "row: " << it.row() << endl;
  cout << "col: " << it.col() << endl;
}