NNLS

/* Non-Negagive Least Squares Algorithm for Eigen.
 *
 * Copyright (C) 2021 Essex Edwards, <essex.edwards@gmail.com>
 * Copyright (C) 2013 Hannes Matuschek, hannes.matuschek at uni-potsdam.de
 *
 * This Source Code Form is subject to the terms of the Mozilla
 * Public License v. 2.0. If a copy of the MPL was not distributed
 * with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */

 
#ifndef EIGEN_NNLS_H
#define EIGEN_NNLS_H
 
#include "../../Eigen/Core"
#include "../../Eigen/QR"
 
#include <limits>
 
namespace Eigen {

template <class MatrixType_>
class NNLS {
 public:
  typedef MatrixType_ MatrixType;
 
  enum {
    RowsAtCompileTime = MatrixType::RowsAtCompileTime,
    ColsAtCompileTime = MatrixType::ColsAtCompileTime,
    Options = MatrixType::Options,
    MaxRowsAtCompileTime = MatrixType::MaxRowsAtCompileTime,
    MaxColsAtCompileTime = MatrixType::MaxColsAtCompileTime
  };
 
  typedef typename MatrixType::Scalar Scalar;
  typedef typename MatrixType::RealScalar RealScalar;
  typedef typename MatrixType::Index Index;

  typedef Matrix<Scalar, ColsAtCompileTime, 1> SolutionVectorType;
  typedef Matrix<Scalar, RowsAtCompileTime, 1> RhsVectorType;
  typedef Matrix<Index, ColsAtCompileTime, 1> IndicesType;

  NNLS();

  NNLS(const MatrixType &A, Index max_iter = -1, Scalar tol = NumTraits<Scalar>::dummy_precision());

  template <typename MatrixDerived>
  NNLS<MatrixType> &compute(const EigenBase<MatrixDerived> &A);

  const SolutionVectorType &solve(const RhsVectorType &b);

  const SolutionVectorType &x() const { return x_; }

  Scalar tolerance() const { return tolerance_; }

  NNLS<MatrixType> &setTolerance(const Scalar &tolerance) {
    tolerance_ = tolerance;
    return *this;
  }

  Index maxIterations() const { return max_iter_ < 0 ? 2 * A_.cols() : max_iter_; }

  NNLS<MatrixType> &setMaxIterations(Index maxIters) {
    max_iter_ = maxIters;
    return *this;
  }

  Index iterations() const { return iterations_; }

  ComputationInfo info() const { return info_; }
 
 private:
  void moveToInactiveSet_(Index idx);

  void moveToActiveSet_(Index idx);

  void solveInactiveSet_(const RhsVectorType &b);
 
 private:
  typedef Matrix<Scalar, ColsAtCompileTime, ColsAtCompileTime> MatrixAtAType;

  Index max_iter_;
  Index iterations_;
  ComputationInfo info_;
  Index numInactive_;
  Scalar tolerance_;
  MatrixType A_;
  MatrixAtAType AtA_;
  SolutionVectorType x_;
  SolutionVectorType gradient_;
  SolutionVectorType y_;
  SolutionVectorType Atb_;
  IndicesType index_sets_;
  MatrixType QR_;
  SolutionVectorType qrCoeffs_;
  SolutionVectorType tempSolutionVector_;
  RhsVectorType tempRhsVector_;
};
 
/* ********************************************************************************************
 * Implementation
 * ******************************************************************************************** */
 
template <typename MatrixType>
NNLS<MatrixType>::NNLS()
    : max_iter_(-1),
      iterations_(0),
      info_(ComputationInfo::InvalidInput),
      numInactive_(0),
      tolerance_(NumTraits<Scalar>::dummy_precision()) {}
 
template <typename MatrixType>
NNLS<MatrixType>::NNLS(const MatrixType &A, Index max_iter, Scalar tol) : max_iter_(max_iter), tolerance_(tol) {
  compute(A);
}
 
template <typename MatrixType>
template <typename MatrixDerived>
NNLS<MatrixType> &NNLS<MatrixType>::compute(const EigenBase<MatrixDerived> &A) {
  // Ensure Scalar type is real. The non-negativity constraint doesn't obviously extend to complex numbers.
  EIGEN_STATIC_ASSERT(!NumTraits<Scalar>::IsComplex, NUMERIC_TYPE_MUST_BE_REAL);
 
  // max_iter_: unchanged
  iterations_ = 0;
  info_ = ComputationInfo::Success;
  numInactive_ = 0;
  // tolerance: unchanged
  A_ = A.derived();
  AtA_.noalias() = A_.transpose() * A_;
  x_.resize(A_.cols());
  gradient_.resize(A_.cols());
  y_.resize(A_.cols());
  Atb_.resize(A_.cols());
  index_sets_.resize(A_.cols());
  QR_.resize(A_.rows(), A_.cols());
  qrCoeffs_.resize(A_.cols());
  tempSolutionVector_.resize(A_.cols());
  tempRhsVector_.resize(A_.rows());
 
  return *this;
}
 
template <typename MatrixType>
const typename NNLS<MatrixType>::SolutionVectorType &NNLS<MatrixType>::solve(const RhsVectorType &b) {
  // Initialize solver
  iterations_ = 0;
  info_ = ComputationInfo::NumericalIssue;
  x_.setZero();
 
  index_sets_ = IndicesType::LinSpaced(A_.cols(), 0, A_.cols() - 1);  // Identity permutation.
  numInactive_ = 0;
 
  // Precompute A^T*b
  Atb_.noalias() = A_.transpose() * b;
 
  const Index maxIterations = this->maxIterations();
 
  // OUTER LOOP
  while (true) {
    // Early exit if all variables are inactive, which breaks 'maxCoeff' below.
    if (A_.cols() == numInactive_) {
      info_ = ComputationInfo::Success;
      return x_;
    }
 
    // Find the maximum element of the gradient in the active set.
    // If it is small or negative, then we have converged.
    // Else, we move that variable to the inactive set.
    gradient_.noalias() = Atb_ - AtA_ * x_;
 
    const Index numActive = A_.cols() - numInactive_;
    Index argmaxGradient = -1;
    const Scalar maxGradient = gradient_(index_sets_.tail(numActive)).maxCoeff(&argmaxGradient);
    argmaxGradient += numInactive_;  // because tail() skipped the first numInactive_ elements
 
    if (maxGradient < tolerance_) {
      info_ = ComputationInfo::Success;
      return x_;
    }
 
    moveToInactiveSet_(argmaxGradient);
 
    // INNER LOOP
    while (true) {
      // Check if max. number of iterations is reached
      if (iterations_ >= maxIterations) {
        info_ = ComputationInfo::NoConvergence;
        return x_;
      }
 
      // Solve least-squares problem in inactive set only,
      // this step is rather trivial as moveToInactiveSet_ & moveToActiveSet_
      // updates the QR decomposition of inactive columns A^N.
      // solveInactiveSet_ puts the solution in y_
      solveInactiveSet_(b);
      ++iterations_;  // The solve is expensive, so that is what we count as an iteration.
 
      // Check feasibility...
      bool feasible = true;
      Scalar alpha = NumTraits<Scalar>::highest();
      Index infeasibleIdx = -1;  // Which variable became infeasible first.
      for (Index i = 0; i < numInactive_; i++) {
        Index idx = index_sets_[i];
        if (y_(idx) < 0) {
          // t should always be in [0,1].
          Scalar t = -x_(idx) / (y_(idx) - x_(idx));
          if (alpha > t) {
            alpha = t;
            infeasibleIdx = i;
            feasible = false;
          }
        }
      }
      eigen_assert(feasible || 0 <= infeasibleIdx);
 
      // If solution is feasible, exit to outer loop
      if (feasible) {
        x_ = y_;
        break;
      }
 
      // Infeasible solution -> interpolate to feasible one
      for (Index i = 0; i < numInactive_; i++) {
        Index idx = index_sets_[i];
        x_(idx) += alpha * (y_(idx) - x_(idx));
      }
 
      // Remove these indices from the inactive set and update QR decomposition
      moveToActiveSet_(infeasibleIdx);
    }
  }
}
 
template <typename MatrixType>
void NNLS<MatrixType>::moveToInactiveSet_(Index idx) {
  // Update permutation matrix:
  std::swap(index_sets_(idx), index_sets_(numInactive_));
  numInactive_++;
 
  // Perform rank-1 update of the QR decomposition stored in QR_ & qrCoeff_
  internal::householder_qr_inplace_update(QR_, qrCoeffs_, A_.col(index_sets_(numInactive_ - 1)), numInactive_ - 1,
                                          tempSolutionVector_.data());
}
 
template <typename MatrixType>
void NNLS<MatrixType>::moveToActiveSet_(Index idx) {
  // swap index with last inactive one & reduce number of inactive columns
  std::swap(index_sets_(idx), index_sets_(numInactive_ - 1));
  numInactive_--;
  // Update QR decomposition starting from the removed index up to the end [idx, ..., numInactive_]
  for (Index i = idx; i < numInactive_; i++) {
    Index col = index_sets_(i);
    internal::householder_qr_inplace_update(QR_, qrCoeffs_, A_.col(col), i, tempSolutionVector_.data());
  }
}
 
template <typename MatrixType>
void NNLS<MatrixType>::solveInactiveSet_(const RhsVectorType &b) {
  eigen_assert(numInactive_ > 0);
 
  tempRhsVector_ = b;
 
  // tmpRHS(0:numInactive_-1) := Q'*b
  // tmpRHS(numInactive_:end) := useless stuff we would rather not compute at all.
  tempRhsVector_.applyOnTheLeft(
      householderSequence(QR_.leftCols(numInactive_), qrCoeffs_.head(numInactive_)).transpose());
 
  // tempSol(0:numInactive_-1) := inv(R) * Q' * b
  //  = the least-squares solution for the inactive variables.
  tempSolutionVector_.head(numInactive_) =            //
      QR_.topLeftCorner(numInactive_, numInactive_)   //
          .template triangularView<Upper>()           //
          .solve(tempRhsVector_.head(numInactive_));  //
 
  // tempSol(numInactive_:end) := 0 = the value for the constrained variables.
  tempSolutionVector_.tail(y_.size() - numInactive_).setZero();
 
  // Back permute into original column order of A
  y_.noalias() = index_sets_.asPermutation() * tempSolutionVector_.head(y_.size());
}
 
}  // namespace Eigen
 
#endif  // EIGEN_NNLS_H