constraint-set-product.hpp

#pragma once
 
#include "proxsuite-nlp/constraint-set.hpp"
 
namespace proxsuite {
namespace nlp {
template <typename Derived>
auto blockMatrixGetRow(const Eigen::MatrixBase<Derived> &matrix,
                       const std::vector<Eigen::Index> &rowBlockSizes,
                       std::size_t rowIdx) {
  Eigen::Index startIdx = 0;
  for (std::size_t kk = 0; kk < rowIdx; kk++) {
    startIdx += rowBlockSizes[kk];
  }
  return matrix.const_cast_derived().middleRows(startIdx,
                                                rowBlockSizes[rowIdx]);
}
 
template <typename Derived>
auto blockVectorGetRow(const Eigen::MatrixBase<Derived> &matrix,
                       const std::vector<Eigen::Index> &blockSizes,
                       std::size_t blockIdx) {
  EIGEN_STATIC_ASSERT_VECTOR_ONLY(Derived);
  Eigen::Index startIdx = 0;
  for (std::size_t kk = 0; kk < blockIdx; kk++) {
    startIdx += blockSizes[kk];
  }
  return matrix.const_cast_derived().segment(startIdx, blockSizes[blockIdx]);
}

template <typename Scalar>
struct ConstraintSetProductTpl : ConstraintSetTpl<Scalar> {
  PROXSUITE_NLP_DYNAMIC_TYPEDEFS(Scalar);
  using Base = ConstraintSetTpl<Scalar>;
  using ActiveType = typename Base::ActiveType;
 
  ConstraintSetProductTpl(const std::vector<xyz::polymorphic<Base>> components,
                          const std::vector<Eigen::Index> &blockSizes)
      : m_components(components), m_blockSizes(blockSizes) {
    if (components.size() != blockSizes.size()) {
      PROXSUITE_NLP_RUNTIME_ERROR("Number of components and corresponding "
                                  "block sizes should be the same.");
    }
  }
 
  ConstraintSetProductTpl(const ConstraintSetProductTpl &) = default;
  ConstraintSetProductTpl &operator=(const ConstraintSetProductTpl &) = default;
  ConstraintSetProductTpl(ConstraintSetProductTpl &&) = default;
  ConstraintSetProductTpl &operator=(ConstraintSetProductTpl &&) = default;
 
  Scalar evaluate(const ConstVectorRef &zproj) const override {
    Scalar res = 0.;
    for (std::size_t i = 0; i < m_components.size(); i++) {
      auto zblock = blockVectorGetRow(zproj, m_blockSizes, i);
      res += m_components[i]->evaluate(zblock);
    }
    return res;
  }
 
  void projection(const ConstVectorRef &z, VectorRef zout) const override {
    for (std::size_t i = 0; i < m_components.size(); i++) {
      ConstVectorRef inblock = blockVectorGetRow(z, m_blockSizes, i);
      VectorRef outblock = blockVectorGetRow(zout, m_blockSizes, i);
      m_components[i]->projection(inblock, outblock);
    }
  }
 
  void normalConeProjection(const ConstVectorRef &z,
                            VectorRef zout) const override {
    for (std::size_t i = 0; i < m_components.size(); i++) {
      ConstVectorRef inblock = blockVectorGetRow(z, m_blockSizes, i);
      VectorRef outblock = blockVectorGetRow(zout, m_blockSizes, i);
      m_components[i]->normalConeProjection(inblock, outblock);
    }
  }
 
  void applyProjectionJacobian(const ConstVectorRef &z,
                               MatrixRef Jout) const override {
    for (std::size_t i = 0; i < m_components.size(); i++) {
      ConstVectorRef inblock = blockVectorGetRow(z, m_blockSizes, i);
      MatrixRef outblock = blockMatrixGetRow(Jout, m_blockSizes, i);
      m_components[i]->applyProjectionJacobian(inblock, outblock);
    }
  }
 
  void applyNormalConeProjectionJacobian(const ConstVectorRef &z,
                                         MatrixRef Jout) const override {
    for (std::size_t i = 0; i < m_components.size(); i++) {
      ConstVectorRef inblock = blockVectorGetRow(z, m_blockSizes, i);
      MatrixRef outblock = blockMatrixGetRow(Jout, m_blockSizes, i);
      m_components[i]->applyNormalConeProjectionJacobian(inblock, outblock);
    }
  }
 
  void computeActiveSet(const ConstVectorRef &z,
                        Eigen::Ref<ActiveType> out) const override {
    for (std::size_t i = 0; i < m_components.size(); i++) {
      ConstVectorRef inblock = blockVectorGetRow(z, m_blockSizes, i);
      decltype(out) outblock = blockVectorGetRow(out, m_blockSizes, i);
      m_components[i]->computeActiveSet(inblock, outblock);
    }
  }
 
  const std::vector<xyz::polymorphic<Base>> &components() const {
    return m_components;
  }
  const std::vector<Eigen::Index> &blockSizes() const { return m_blockSizes; }
 
private:
  std::vector<xyz::polymorphic<Base>> m_components;
  std::vector<Eigen::Index> m_blockSizes;
};
 
#ifdef PROXSUITE_NLP_ENABLE_TEMPLATE_INSTANTIATION
extern template struct PROXSUITE_NLP_EXPLICIT_INSTANTIATION_DECLARATION_DLLAPI
    ConstraintSetProductTpl<context::Scalar>;
#endif
 
} // namespace nlp
} // namespace proxsuite