Search code examples
c++eigeneigen3

Eigen unaryExpr(): get index of element


I know that in Eigen I can use unaryExpr() to apply a custom function to my Eigen Matrices and Vectors, e.g.

Eigen::VectorXd vec(4);
vec.unaryExpr([](double d) {
    return d * cos(d);
});

to apply custom functions to a Eigen::VectorXd. But is there also a way to get the position of the current element in my Vector? I want to be able to do something like this:

Eigen::VectorXd vec(4);
vec.unaryExpr([](double d, int index) {
     return index * d;
 });

which would for example multiply each entry in the vector by it's position.


Solution

  • Besides using a nullary expression as suggested by @ggael, you can also use a binary expression together with LinSpaced(size, 0, size-1):

    VectorXd v; // some input vector
    
    v = v.binaryExpr(VectorXd::LinSpaced(v.size(), 0, v.size()-1),
                     [](double d, double i){return d*i;});
    // or much simpler and more efficient in this case:
    v = v.cwiseProduct(VectorXd::LinSpaced(v.size(), 0, v.size()-1));
    

    On sufficiently recent versions of Eigen, LinSpaced should get vectorized (although there are some border cases regarding the last element). binaryExpr only gets vectorized if the passed functor is vectorized, of course.

    N.B.: If you are doing mostly element-wise operations, consider using ArrayXd instead of VectorXd.