mimic MatLab’s ability to index a matrix through two vectors of indices referencing the rows and columns to be picked

The goal here is to mimic MatLab’s ability to index a matrix through two vectors of indices referencing the rows and columns to be picked respectively, like this:

A =
7 9 -5 -3
-2 -6 1 0
6 -3 0 9
6 6 3 9
A([1 2 1], [3 2 1 0 0 2]) =
0 1 -6 -2 -2 1
9 0 -3 6 6 0
0 1 -6 -2 -2 1

To this end, let us first write a nullary-functor storing references to the input matrix and to the two arrays of indices, and implementing the required operator()(i,j):

template<class ArgType, class RowIndexType, class ColIndexType>
class indexing_functor {
const ArgType &m_arg;
const RowIndexType &m_rowIndices;
const ColIndexType &m_colIndices;
public:
typedef Matrix<typename ArgType::Scalar,
RowIndexType::SizeAtCompileTime,
ColIndexType::SizeAtCompileTime,
ArgType::Flags&RowMajorBit?RowMajor:ColMajor,
RowIndexType::MaxSizeAtCompileTime,
ColIndexType::MaxSizeAtCompileTime> MatrixType;

indexing_functor(const ArgType& arg, const RowIndexType& row_indices, const ColIndexType& col_indices)
: m_arg(arg), m_rowIndices(row_indices), m_colIndices(col_indices)
{}

const typename ArgType::Scalar& operator() (Index row, Index col) const {
return m_arg(m_rowIndices[row], m_colIndices[col]);
}
};

Then, let’s create an indexing(A,rows,cols) function creating the nullary expression:

template <class ArgType, class RowIndexType, class ColIndexType>
CwiseNullaryOp<indexing_functor<ArgType,RowIndexType,ColIndexType>, typename indexing_functor<ArgType,RowIndexType,ColIndexType>::MatrixType>
indexing(const Eigen::MatrixBase<ArgType>& arg, const RowIndexType& row_indices, const ColIndexType& col_indices)
{
typedef indexing_functor<ArgType,RowIndexType,ColIndexType> Func;
typedef typename Func::MatrixType MatrixType;
return MatrixType::NullaryExpr(row_indices.size(), col_indices.size(), Func(arg.derived(), row_indices, col_indices));
}

Finally, here is an example of how this function can be used:

Eigen::MatrixXi A = Eigen::MatrixXi::Random(4,4);
Array3i ri(1,2,1);
ArrayXi ci(6); ci << 3,2,1,0,0,2;
Eigen::MatrixXi B = indexing(A, ri, ci);
std::cout << "A =" << std::endl;
std::cout << A << std::endl << std::endl;
std::cout << "A([" << ri.transpose() << "], [" << ci.transpose() << "]) =" << std::endl;
std::cout << B << std::endl;
This straightforward implementation is already quite powerful as the row or column index arrays can also be expressions to perform offsetting, modulo, striding, reverse, etc.

B = indexing(A, ri+1, ci);
std::cout << "A(ri+1,ci) =" << std::endl;
std::cout << B << std::endl << std::endl;
#if __cplusplus >= 201103L
B = indexing(A, ArrayXi::LinSpaced(13,0,12).unaryExpr([](int x){return x%4;}), ArrayXi::LinSpaced(4,0,3));
std::cout << "A(ArrayXi::LinSpaced(13,0,12).unaryExpr([](int x){return x%4;}), ArrayXi::LinSpaced(4,0,3)) =" << std::endl;
std::cout << B << std::endl << std::endl;
#endif

and the output is:

A(ri+1,ci) =
9 0 -3 6 6 0
9 3 6 6 6 3
9 0 -3 6 6 0
A(ArrayXi::LinSpaced(13,0,12).unaryExpr([](int x){return x%4;}), ArrayXi::LinSpaced(4,0,3)) =
7 9 -5 -3
-2 -6 1 0
6 -3 0 9
6 6 3 9
7 9 -5 -3
-2 -6 1 0
6 -3 0 9
6 6 3 9
7 9 -5 -3
-2 -6 1 0
6 -3 0 9
6 6 3 9
7 9 -5 -3

Original from:  http://eigen.tuxfamily.org/dox-devel/TopicCustomizing_NullaryExpr.html#title1

Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s