c++ - Eigen Indices of Dense Matrix meeting Condition

Question

I'm looking to get the row/column indices from a dense matrix which meet a condition. In my case the result is likely to be very sparse. For example, with matrix

1 5 2
7 6 3
2 3 8

I'd like to get the indicies where the coeff is greater than 4, or

(0,1), (1,0), (1,1), (2,2)

My initial thoughts include using either select or coefficientwise operations to build a bool/int matrix

0 1 0
1 1 0
0 0 1

and then convert it to a sparse matrix

(0,1,1), (1,0,1), (1,1,1), (2,2,1)

then remove the coeff values

(0,1), (1,0), (1,1), (2,2)

However, that requires two passes over matrices the size of the original matrix, which may be very large.

Alternatively a naive double loop over the original matrix akin to pseudocode

for (int i; i < mat.cols(); i++) {
    for (int j; j < mat.rows(); j++) {
        if(cond(mat(j, i))) {
            add_index(i, j, index_list)
        }
    }
}

but that only takes the compilers optimizations and none of Eigen's optimizations/vectorizations.

Is there a more efficient way that I am missing? In my case the conditions are simple comparisons.

Thanks for any help

Answer 1

There is not much to vectorize here, but to avoid writing the two loops (which might better be reversed for a row-major storage), the best tool would be a visitor:

mat.visit(some_visitor);

Sadly, the visitor cannot be a simple lambda as visit calls the method init(val,0,0) for the first element. This is useful for reduction-like visitors but not always. In order to make visit accepts a simple lambda, you can use the following helper:

template<typename Func>
struct lambda_as_visitor_wrapper : Func {
    lambda_as_visitor_wrapper(const Func& f) : Func(f) {}
    template<typename S,typename I>
    void init(const S& v, I i, I j) { return Func::operator()(v,i,j); }
};

template<typename Mat, typename Func>
void visit_lambda(const Mat& m, const Func& f)
{
    lambda_as_visitor_wrapper<Func> visitor(f);
    m.visit(visitor);
}

In your case, you can then write:

int main() {
    int n = 5;
    double th = 0.5;
    Eigen::MatrixXd M = Eigen::MatrixXd::Random(n,n);

    std::vector<std::pair<int,int>> indices;
    visit_lambda(M,
        [&indices,th](double v, int i, int j) {
            if(v>th)
                indices.push_back(std::make_pair(i,j));
        });


    std::cout << M << "

";

    for(auto p:indices)
        std::cout << '(' << p.first << ',' << p.second << ") ";
    std::cout << '
';

    return 0;
}