I expanded ggael's answer to Array
types, vertical concatenation, and more than two arguments:
#include <iostream>
#include <Eigen/Core>
namespace EigenCustom
{
using namespace Eigen;
constexpr Index dynamicOrSum( const Index& a, const Index& b ){
return a == Dynamic || b == Dynamic ? Dynamic : a + b;
}
enum class Direction { horizontal, vertical };
template<Direction direction, typename Arg1, typename Arg2>
struct ConcatHelper {
static_assert( std::is_same_v<
typename Arg1::Scalar, typename Arg2::Scalar
> );
using Scalar = typename Arg1::Scalar;
using D = Direction;
static constexpr Index
RowsAtCompileTime { direction == D::horizontal ?
Arg1::RowsAtCompileTime :
dynamicOrSum( Arg1::RowsAtCompileTime, Arg2::RowsAtCompileTime )
},
ColsAtCompileTime { direction == D::horizontal ?
dynamicOrSum( Arg1::ColsAtCompileTime, Arg2::ColsAtCompileTime ) :
Arg1::ColsAtCompileTime
},
MaxRowsAtCompileTime { direction == D::horizontal ?
Arg1::MaxRowsAtCompileTime :
dynamicOrSum( Arg1::MaxRowsAtCompileTime, Arg2::MaxRowsAtCompileTime )
},
MaxColsAtCompileTime { direction == D::horizontal ?
dynamicOrSum( Arg1::MaxColsAtCompileTime, Arg2::MaxColsAtCompileTime ) :
Arg1::MaxColsAtCompileTime
};
static_assert(
(std::is_base_of_v<MatrixBase<Arg1>, Arg1> &&
std::is_base_of_v<MatrixBase<Arg2>, Arg2> ) ||
(std::is_base_of_v<ArrayBase<Arg1>, Arg1> &&
std::is_base_of_v<ArrayBase<Arg2>, Arg2> )
);
using DenseType = std::conditional_t<
std::is_base_of_v<MatrixBase<Arg1>, Arg1>,
Matrix<
Scalar, RowsAtCompileTime, ColsAtCompileTime,
ColMajor, MaxRowsAtCompileTime, MaxColsAtCompileTime
>,
Array<
Scalar, RowsAtCompileTime, ColsAtCompileTime,
ColMajor, MaxRowsAtCompileTime, MaxColsAtCompileTime
>
>;
};
template<Direction direction, typename Arg1, typename Arg2>
class ConcatFunctor
{
using Scalar = typename ConcatHelper<direction, Arg1, Arg2>::Scalar;
const typename Arg1::Nested m_mat1;
const typename Arg2::Nested m_mat2;
public:
ConcatFunctor(const Arg1& arg1, const Arg2& arg2)
: m_mat1(arg1), m_mat2(arg2)
{}
const Scalar operator() (Index row, Index col) const {
if constexpr (direction == Direction::horizontal){
if (col < m_mat1.cols())
return m_mat1(row,col);
return m_mat2(row, col - m_mat1.cols());
} else {
if (row < m_mat1.rows())
return m_mat1(row,col);
return m_mat2(row - m_mat1.rows(), col);
}
}
};
template<Direction direction, typename Arg1, typename Arg2>
using ConcatReturnType = CwiseNullaryOp<
ConcatFunctor<direction,Arg1,Arg2>,
typename ConcatHelper<direction,Arg1,Arg2>::DenseType
>;
template<Direction direction, typename Arg1, typename Arg2>
ConcatReturnType<direction, Arg1, Arg2>
concat(
const Eigen::DenseBase<Arg1>& arg1,
const Eigen::DenseBase<Arg2>& arg2
){
using DenseType = typename ConcatHelper<direction,Arg1,Arg2>::DenseType;
using D = Direction;
return DenseType::NullaryExpr(
direction == D::horizontal ? arg1.rows() : arg1.rows() + arg2.rows(),
direction == D::horizontal ? arg1.cols() + arg2.cols() : arg1.cols(),
ConcatFunctor<direction,Arg1,Arg2>( arg1.derived(), arg2.derived() )
);
}
template<Direction direction, typename Arg1, typename Arg2, typename ... Ts>
decltype(auto)
concat(
const Eigen::DenseBase<Arg1>& arg1,
const Eigen::DenseBase<Arg2>& arg2,
Ts&& ... rest
){
return concat<direction>(
concat<direction>(arg1, arg2),
std::forward<Ts>(rest) ...
);
}
template<typename Arg1, typename Arg2, typename ... Ts>
decltype(auto)
concat_horizontal(
const Eigen::DenseBase<Arg1>& arg1,
const Eigen::DenseBase<Arg2>& arg2,
Ts&& ... rest
){
return concat<Direction::horizontal>(
arg1, arg2, std::forward<Ts>(rest) ...
);
}
template<typename Arg1, typename Arg2, typename ... Ts>
decltype(auto)
concat_vertical(
const Eigen::DenseBase<Arg1>& arg1,
const Eigen::DenseBase<Arg2>& arg2,
Ts&& ... rest
){
return concat<Direction::vertical>(
arg1, arg2, std::forward<Ts>(rest) ...
);
}
} // namespace EigenCustom
int main()
{
using namespace Eigen;
using namespace EigenCustom;
MatrixXd mat(3, 3);
mat << 0, 1, 2, 3, 4, 5, 6, 7, 8;
auto example1 = concat_horizontal(mat,2*mat);
std::cout << "example1:\n" << example1 << '\n';
auto example2 = concat_horizontal(VectorXd::Ones(3),mat);
std::cout << "example2:\n" << example2 << '\n';
auto example3 = concat_vertical(mat,RowVectorXd::Zero(3));
std::cout << "example3:\n" << example3 << '\n';
ArrayXXi arr (2,2);
arr << 0, 1, 2, 3;
auto example4 = concat_vertical(arr,Array2i{4,5}.transpose());
std::cout << "example4:\n" << example4 << '\n';
/* concatenating more than two arguments */
auto example5 = concat_horizontal(mat, mat, mat);
std::cout << "example5:\n" << example5 << '\n';
using RowArray2i = Array<int, 1, 2>;
auto example6 = concat_vertical( arr, RowArray2i::Zero(), RowArray2i::Ones() );
std::cout << "example6:\n" << example6 << '\n';
return 0;
}