eigen：设置/返回整列的表达式或函数

Question

我的代码中有几个实例，其中我有一个基于 1xN 数组系数的条件，并且需要根据这些条件设置 MxN 数组的整列。在我的例子中，N 是Eigen::Dynamic，M 的范围是 2 到 4，但在每个实例中都是编译时常量。

这是一个简单的函数来说明我的意思，a 和 b 是构成条件的 1xN 数组，c 是带有附加数据的 2xN 数组，res 是外参数，其列始终设置为一个整体：

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

using namespace Eigen;

template<Index nRows>
using ArrayNXd = Array<double, nRows, Dynamic>;

using Array1Xd = ArrayNXd<1>;
using Array2Xd = ArrayNXd<2>;
using Array3Xd = ArrayNXd<3>;

void asFunction(
    Array3Xd& res,
    const Array1Xd& a, const Array1Xd& b, const Array2Xd& c
){
    for (Index col{0}; col<a.cols(); ++col){
        if ( a[col] > b[col] )
            res.col(col) = Array3d{
                 a[col] + b[col],
                (a[col] + b[col]) * c(0, col),
                (a[col] - b[col]) * c(1, col)
            };
        else
            res.col(col) = Array3d{
                 a[col] - b[col],
                 a[col] + b[col],
                (a[col] + b[col]) * (a[col] - b[col])
            };
    }
}


int main(){
    Array1Xd a (3), b(3);
    Array2Xd c (2, 3);
    
    a << 1, 2, 3;
    b << 0, 1, 2;
    c <<
        0, 1, 2,
        1, 2, 3;

    Array3Xd res (3,3);
    
    asFunction(res, a, b, c);

    std::cout << "as function:\n" << res << "\n";

    return 0;
}

在我的代码的性能关键部分*中使用了与此类似的函数，我觉得我将性能留在了桌面上，因为使用带有 Eigen 类型的循环通常不是最佳解决方案。

*是的，我对其进行了分析。

我编写了与NullaryExpr 相同的函数，但速度有点慢。我想这是有道理的，考虑到条件的额外评估和每一行的分支：

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

using namespace Eigen;

template<Index nRows>
using ArrayNXd = Array<double, nRows, Dynamic>;

using Array1Xd = ArrayNXd<1>;
using Array2Xd = ArrayNXd<2>;
using Array3Xd = ArrayNXd<3>;

class MyFunctor
{
public:
    using Scalar = double;

    static constexpr Index
        RowsAtCompileTime { 3 },
        MaxRowsAtCompileTime { 3 },
        ColsAtCompileTime { Dynamic },
        MaxColsAtCompileTime { Dynamic };

    using DenseType = Array<
        Scalar  ,    RowsAtCompileTime,    ColsAtCompileTime,
        ColMajor, MaxRowsAtCompileTime, MaxColsAtCompileTime
    >;

private:
    typename Array1Xd::Nested m_a;
    typename Array1Xd::Nested m_b;
    typename Array2Xd::Nested m_c;

public:
    MyFunctor(
        const Array1Xd& a,
        const Array1Xd& b,
        const Array2Xd& c
    ) : m_a {a}, m_b {b}, m_c{c}
    {}

    bool cond(Index col) const {
        return m_a[col] > m_b[col];
    }

    Scalar func1(Index col) const {
        return m_a[col] + m_b[col];
    }

    Scalar func2(Index col) const {
        return m_a[col] - m_b[col];
    }

    Scalar func3(Index row, Index col) const {
        switch(row){
            case 0: return func1(col);
            case 1: return func1(col) * m_c(0, col);
            case 2: return func2(col) * m_c(1, col);
            default: __builtin_unreachable();
        }
    }

    Scalar func4(Index row, Index col) const {
        switch (row){
            case 0: return func2(col);
            case 1: return func1(col);
            case 2: return func1(col) / func2(col);
            default: __builtin_unreachable();
        }
    }

    Scalar operator() (Index row, Index col) const {
        if ( cond(col) )
            return func3(row, col);
        else
            return func4(row, col);
    }
};

using MyReturnType = Eigen::CwiseNullaryOp<
    MyFunctor, typename MyFunctor::DenseType
>;


MyReturnType asFunctor(
    const Array1Xd& a,
    const Array1Xd& b,
    const Array2Xd& c
){
    using DenseType = typename MyFunctor::DenseType;
    return DenseType::NullaryExpr(
        3, a.cols(),
        MyFunctor(a, b, c)
    );
}


int main(){
    Array1Xd a (3), b(3);
    Array2Xd c (2, 3);
    
    a << 1, 2, 3;
    b << 0, 1, 2;
    c <<
        0, 1, 2,
        1, 2, 3;

    std::cout << "as functor:\n" << asFunctor(a,b,c) << "\n";

    return 0;
}

我的问题是：有没有更有效的方法来实现类似于上面的逻辑（评估矩阵每一列的标量条件，根据条件返回整列的值）使用eigen 库？

注意：使用表达式会稍微好一些，因为我不需要担心内存分配、输出参数等，并且可以在编写代码时考虑到标量，这使得它更容易理解.

编辑：注2：我也尝试使用<Condition>.template replicate<nRows,1>().select(..., ...)，但它更慢且更难阅读。

Answer 1

所以我只看了这段代码

    for (Index col{0}; col<a.cols(); ++col){
        if ( a[col] > b[col] )
            res.col(col) = Array3d{
                 a[col] + b[col],
                (a[col] + b[col]) * c(0, col),
                (a[col] - b[col]) * c(1, col)
            };
        else
            res.col(col) = Array3d{
                 a[col] - b[col],
                 a[col] + b[col],
                (a[col] + b[col]) * (a[col] - b[col])
            };
    }

我怀疑，但不能证明，每次调用它们时 a[col] 和 b[col] 都会被访问。您可能想尝试为您重用的值制作简短的临时文件。例如： 所以我只看了这段代码

    for (Index col{0}; col<a.cols(); ++col){
        auto acol=a[col];
        auto bcol=b[col];
        auto apb=acol+bcol;
        auto amb=acol-bcol;
        if ( acol > bcol )
            res.col(col) = Array3d{
                 apb,
                (apb) * c(0, col),
                (amb) * c(1, col)
            };
        else
            res.col(col) = Array3d{
                 amb,
                 apb,
                (apb) * (amb)
            };
    }

是的，我知道这不是您想要的。也许对你有帮助

Answer 2

您可以使用 Eigen 的 select 方法，但它仅适用于标量，因此您必须沿一维循环。

const auto condition = a > b;
res.row(0) = condition.select(a + b /*true*/, a - b /*false*/);
res.row(1) = condition.select((a + b) * c.row(0), a + b);
res.row(2) = condition.select((a - b) * c.row(1), (a + b) * (a - b));

请注意，如果您转置所有数组，您可能会更快。然后迭代逐列进行，因为 Eigen 是列主要的，所以向量化得更好。