Khue Khue - 28 days ago 13
C++ Question

How to vectorize in Eigen/C++: set columns under condition

I have a function

compute()
that takes as input
a matrix
C
and output a matrix
X
of the same dimension, and does the following:


  • For each column
    C_j
    of C, if the sum of its positive elements are less than or equal to 1, then the column
    X_j
    of
    X
    is the same as
    C_j
    , except that the non-positive elements are set to zeros, i.e.
    X_j = C_j(C_j > 0)
    (in Matlab language).

  • Otherwise
    X_j = f(C_j)
    where
    f
    is some function. This part is not concerned in this question.



What I have tried so far.

MatrixXd compute(MatrixXd &C)
{
// The matrix to return
MatrixXd X(C.rows(), C.cols());

// Thresholding C: if negative then set to 0
MatrixXd P = (C.array() > 0).select(C, 0);

// Compute the sum of each column of P (i.e. the
// sum of positive elements in C
MatrixXd S = P.colwise().sum();

// Now want to set all column X_j to P_j whenever S_j <= 1
// I don't know how to vectorize this code
// It's easy in Matlab: X(:, S <= 1) = P(:, S <= 1);
for(int j = 0; j < S.cols(); j++){
if(S(j) <= 1)
X.col(j) = P.col(j);
}
return X;
}


For each column of
C

Answer

The shortest solution to your problem I found is this:

MatrixXd compute(const MatrixXd& C)
{
    MatrixXd P = C.cwiseMax(0.0); // cheaper than (...).select(C,0)
    RowVectorXd S = P.colwise().sum(); // colwise().sum() returns only one row
    MatrixXd X = (S.array()<=1).replicate(C.rows(), 1).select(P, 0.0);
    //      compare by 1 --^             ^               ^  
    // replicate result of comparison ---/               |
    //        select P or 0 depending on comparison  ----/
    return X;
}

Notice that neither comparisons nor select are currently (version 3.3rc2) vectorized in Eigen (but they eventually shall be, for any progress on that see http://eigen.tuxfamily.org/bz/show_bug.cgi?id=97). That means that your current loop-implementation could actually be faster (which likely also depends on several other factors, like the size of your input).

Comments