Home > Enterprise >  Eigen3 (cpp) select column given mask and sum where true
Eigen3 (cpp) select column given mask and sum where true

Time:12-02

I have a Eigen::Matrix2Xf where row are X and Y positions and cols act as list index

I would like to have the sum of the columns (rowwise) where some column condition is true, here some example code:

Eigen::Vector2f computeStuff(Eigen::Matrix2Xf & values, const float max_norm){

    const auto mask = values.colwise().norm().array() < max_norm;

    return mask.select(values.colwise(), Eigen::Vector2f::Zero()).rowwise().sum();
    
}

But this code does not compile complaining about the types of the if/else matrices, what is the correct (and computationally faster) way to do it?

Also I know that there are similar question with an answer, but they create a new Eigen::Matrix2Xf with the filtered values given the mask, this code is meant to run inside a #pragma omp parallel for so the basic idea is to do not create a new matrix for maintaining cache coherency

Thanks

CodePudding user response:

The main problem with your code is that .select( ... ) needs one of its arguments to have the same shape as the mask -- at least I don't know how to use .select() otherwise.

In your code mask is a row vector but values is a 2 by x matrix. One way to handle this is to just replicate the mask into a matrix:

#include <Eigen/Dense>
#include <iostream>

Eigen::Vector2f computeStuff(Eigen::Matrix2Xf& values, const float max_norm) {

    auto mask = (values.colwise().norm().array() < max_norm).replicate(2, 1);
    return mask.select(values, 0).rowwise().sum();
}

int main() {

    Eigen::Matrix2Xf mat(2,4);
    mat << 1, 4, 3, 2, 
           1, 2, 4, 3;

    auto val = computeStuff(mat, 5);

    std::cout << val;

    return 0;
}

In the above mask will be:

1 1 0 1
1 1 0 1

the row 1 1 0 1 duplicated once. Then mask.select(values, 0) will select

1 4 0 2
1 2 0 3

so the result will be

7
6

which i think is what you want, if I am understanding the question.

  • Related