c++ - 你如何在 Eigen 中执行 numpy.repeat？

Question

Eigen 有一个类似于 numpy.repeat 的 replicate 方法，但它不支持重复可变次数。例如:

np.repeat(np.array([[ 0.,  1.,  2.], [ 3.,  4., 5.]]), [1, 2], axis=0)

给予

array([[ 0.,  1.,  2.],
       [ 3.,  4.,  5.],
       [ 3.,  4.,  5.]])

如何在 Eigen 中重现此行为？

Answer 1

这是我的尝试:

namespace Eigen {
    template <typename T>
    using ArrayXX = Array<T, Dynamic, Dynamic>;
}

template<typename A>
Eigen::ArrayXX<typename A::Scalar> repeat(const A& a, const Eigen::Ref<const Eigen::ArrayXi>& repeats, const int axis) {
    typedef typename A::Scalar T;
    if (axis==0) {
        eigen_assert(a.rows() == repeats.size());
        const int new_rows = repeats.sum();
        Eigen::ArrayXX<T> repeated_array (new_rows, a.cols());
        int j = 0;
        for (int i = 0; i < repeats.size(); ++i) {
            const int k = repeats(i);
            repeated_array.middleRows(j, k) = a.row(i).colwise().replicate(k);
            j += k;
        }
        return repeated_array;
    } else {
        eigen_assert(a.cols() == repeats.size());
        const int new_cols = repeats.sum();
        Eigen::ArrayXX<T> repeated_array (a.rows(), new_cols);
        int j = 0;
        for (int i = 0; i < repeats.size(); ++i) {
            const int k = repeats(i);
            repeated_array.middleCols(j, k) = a.col(i).rowwise().replicate(k);
            j += k;
        }
        return repeated_array;
    }
}