Search code examples
c++eigentemplate-meta-programming

Template Metaprogramming Eigen Expressions


Let's say I have some (short) vector, the length of which I know at compile time, and another, longer vector, the length of which I don't know at compile time. I could write something like this:

template<int N>
Eigen::ArrayXd do_transformation(Eigen::Array<double,N,1> short_vec, Eigen::ArrayXd long_vec){
    Eigen::ArrayXd return_vector(long_vec.size());
    for(int i=0; i<N; i++){
      return_vector+=short_vec(i)*long_vec.pow(2-i);
    }
    return return_vector;
}

Is there a way I can construct that sum using expression templates without having to write out:

template<>
Eigen::ArrayXd do_transformation<1>(Eigen::Array<double,1,1> short_vec, Eigen::ArrayXd long_vec){

    return short_vec(0)*long_vec.pow(2);
}

template<>
Eigen::ArrayXd do_transformation<2>(Eigen::Array<double,2,1> short_vec, Eigen::ArrayXd long_vec){

    return short_vec(0)*long_vec.pow(2)+short_vec(1)*long_vec.pow(1);
}

for every value of N?

Ideally this could be done in c++11. What would be really awesome would be for the function to return some kind of Eigen expression, so that I could do something like:

long_vec+do_transformation(short_vec,long_vec) and have Eigen direct the compiler to generate code that only traverses the vectors once.


Solution

  • In c++11 you can do a good old recursion:

    template <int I>
    struct do_transformation_impl {
        template<int M>
        static auto run(const Array<double,M,1> &short_vec, const ArrayXd &long_vec)
        -> decltype(short_vec(I)*long_vec.pow(2-I) + do_transformation_impl<I-1>::run(short_vec,long_vec))
        {
            return short_vec(I)*long_vec.pow(2-I) + do_transformation_impl<I-1>::run(short_vec,long_vec);
        }
    };
    
    template <>
    struct do_transformation_impl<0> {
        template<int M>
        static auto run(const Array<double,M,1> &short_vec, const ArrayXd &long_vec)
        -> decltype(short_vec(0)*long_vec.pow(2))
        {
            return short_vec(0)*long_vec.pow(2);
        }
    };
    
    template<int N>
    auto do_transformation(const Array<double,N,1> &short_vec, const ArrayXd &long_vec)
    -> decltype(do_transformation_impl<N-1>::run(short_vec,long_vec))
    {
        return do_transformation_impl<N-1>::run(short_vec,long_vec);
    }
    

    Demo: https://godbolt.org/z/4xPdVm

    After some adjustments of max66's answer, both solution yield to the same code: https://godbolt.org/z/Ha5qMa