Search code examples
c++eigen

Eigen: scale the positive entries in loop-free manner


I have an Eigen::Array<float, N, 1>. I would like to multiply all the positive elements by a positive scalar s chosen so that the positive elements instead sum to a target t. I am able to do this with code like this:

template<int N>
void positive_scale(Eigen::Array<float, N, 1>& A, float t) {
  float pos_sum = 0;
  for (int i = 0; i < N; ++i) {
    if (A(i) > 0) pos_sum += A(i);
  }

  float s = t / pos_sum;
  for (int i = 0; i < N; ++i) {
    if (A(i) > 0) A(i) *= s;
  }
}

I am wondering if there is a more elegant, loop-free way to do this. In python numpy, for instance, we can write this much more compactly:

def positive_scale(A, t):
  A[A > 0] *= t / sum(A[A > 0])

Solution

  • Eigen does not support bools for indices in the way Numpy does, but it has a select method that can be used.

    void positive_scale(Eigen::Ref<Eigen::ArrayXf> inout, float target)
    {
        // sum of positive entries. Fully vectorized
        float pos_sum = inout.cwiseMax(0.f).sum();
        float factor = target / pos_sum;
        // keep old negative values. Replace positive values
        inout = (inout > 0.f).select(inout * factor, inout);
    }