optim.h

Go to the documentation of this file.

//
// Project NuriKit - Copyright 2024 SNU Compbio Lab.
// SPDX-License-Identifier: Apache-2.0
//
 
#ifndef NURI_ALGO_OPTIM_H_
#define NURI_ALGO_OPTIM_H_

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <functional>
#include <utility>
 
#include <absl/log/absl_check.h>
#include <Eigen/Dense>
 
#include "nuri/eigen_config.h"
#include "nuri/core/container/container_ext.h"
#include "nuri/utils.h"
 
namespace nuri {
enum class OptimResultCode {
  kSuccess,
  kMaxIterReached,
  kInvalidInput,
  kAbnormalTerm,
};
 
struct LbfgsResult {
  OptimResultCode code;
  int niter;
  double fx;
  ArrayXd gx;
};
 
namespace internal {
  constexpr double kEpsMach = 2.220446049250313e-16;
  constexpr double kSqrtEpsMach =
      1.4901161193847655978721599999999997530663236602513281125266e-8;
 
  enum class DcsrchStatus : std::uint8_t {
    kFound,
    kConverged,
    kContinue,
  };
 
  class Dcsrch {
  public:
    constexpr static double kStepMin = 0, kStepMax = 1e+10;
 
    Dcsrch(double f0, double g0, double step0, double stepmin = kStepMin,
           double stepmax = kStepMax, double ftol = 1e-3, double gtol = 0.9,
           double xtol = 0.1) noexcept;
 
    DcsrchStatus operator()(double f, double g);
 
    double step() const { return step_; }
 
    double finit() const { return finit_; }
 
    double ginit() const { return ginit_; }
 
  private:
    constexpr static double kXTrapL = 1.1, kXTrapU = 4;
 
    double finit_, ginit_;
    double stepmin_, stepmax_;
    double gtest_, gtol_, xtol_;
 
    double step_;
 
    // The variables stx, fx, gx contain the values of the step, function,
    // and derivative at the best step.
    // The variables sty, fy, gy contain the value of the step, function,
    // and derivative at sty.
    // The variables stp, f, g contain the values of the step, function,
    // and derivative at stp.
    double stx_ = 0, fx_, gx_;
    double sty_ = 0, fy_, gy_;
    double stmin_ = 0, stmax_;
    double width_, width1_;
 
    bool brackt_ = false, bisect_ = false;
  };
 
  class LBfgsBase {
  public:
    LBfgsBase(MutRef<ArrayXd> x, int m);
 
    /* State modifiers, only for implementations */
 
    int n() const { return static_cast<int>(x_.size()); }
 
    int m() const { return static_cast<int>(ws_.cols()); }
 
    auto &x() { return x_; }
 
    const auto &x() const { return x_; }
 
    auto wnt() { return wnt_.topLeftCorner(2 * col_, 2 * col_); }
    auto &wnt_raw() { return wnt_; }
 
    auto ws() { return ws_.leftCols(col_); }
 
    auto wy() { return wy_.leftCols(col_); }
 
    auto ss() { return ss_.topLeftCorner(col_, col_); }
 
    auto sy() { return sy_.topLeftCorner(col_, col_); }
 
    auto wtt() { return wtt_.topLeftCorner(col_, col_); }
 
    auto p() { return p_.head(2 * col_); }
 
    auto v() { return v_.head(2 * col_); }
 
    auto c() { return c_.head(2 * col_); }
 
    auto wbp() { return wbp_.head(2 * col_); }
 
    auto smul() { return smul_.head(nonnegative(col_ - 1)); }
 
    auto &z() { return z_; }
    auto &xcp() { return z_; }
 
    auto &r() { return r_; }
 
    auto &t() { return t_; }
    auto &xp() { return t_; }
 
    auto &d() { return d_; }
 
    double d_dot(const ArrayXd &gx) const {
      return gx.matrix().dot(d_.matrix());
    }
 
    double theta() const { return theta_; }
 
    int col() const { return col_; }
 
    int prev_col() const { return prev_col_; }
 
    bool updated() const { return updated_; }
 
    /* Test utils; normally not used */
 
    void update_col(int col) { col_ = col; }
 
    void update_theta(double theta) { theta_ = theta; }
 
  protected:
    bool prepare_next_iter(double gd, double ginit, double step, double dtd);
 
    void reset_memory() {
      theta_ = 1.0;
      col_ = prev_col_ = 0;
      updated_ = false;
    }
 
  private:
    /*  System, size n */
    MutRef<ArrayXd> x_;
 
    /* (2m, 2m) */
    MatrixXd wnt_;
    /* (n, m) */
    MatrixXd ws_, wy_;
    /* (m, m) */
    MatrixXd ss_, sy_, wtt_;
    /* 2*m */
    VectorXd p_, v_, c_, wbp_;
    /* n */
    ArrayXd z_, r_, t_, d_;
    /* m-1 */
    ArrayXd smul_;
 
    double theta_;
    int col_, prev_col_;
    bool updated_;
  };
 
  extern bool lbfgs_bmv(MutVecBlock<VectorXd> p, MutVecBlock<ArrayXd> smul,
                        ConstRef<VectorXd> v, ConstRef<MatrixXd> sy,
                        ConstRef<MatrixXd> wtt);
}  // namespace internal

template <class Impl>
class LBfgs: public internal::LBfgsBase {
public:
  LBfgs(MutRef<ArrayXd> x, Impl &&impl, const int m = 10)
      : internal::LBfgsBase(x, m), impl_(std::move(impl)) { }

  template <class FuncGrad>
  LbfgsResult minimize(FuncGrad fg, double factr = 1e+7, double pgtol = 1e-5,
                       int maxiter = 15000, int maxls = 20);

  Impl &impl() { return impl_; }
 
private:
  Impl impl_;
};
 
template <class Impl>
template <class FuncGrad>
LbfgsResult LBfgs<Impl>::minimize(FuncGrad fg, const double factr,
                                  const double pgtol, const int maxiter,
                                  const int maxls) {
  const double tol = factr * internal::kEpsMach;
 
  ArrayXd gx(n());
  double fx = fg(gx, x());
 
  internal::AllowEigenMallocScoped<false> ems;
 
  double sbgnrm = impl_.projgr(x(), gx);
  if (sbgnrm <= pgtol)
    return { OptimResultCode::kSuccess, 0, fx, std::move(gx) };
 
  impl_.reset();
  reset_memory();
 
  int iter = 0;
  for (; iter < maxiter; ++iter) {
    if (!impl_.prepare_lnsrch(*this, gx, sbgnrm, iter)) {
      reset_memory();
      continue;
    }
 
    d() = z() - x();
 
    double gd = d_dot(gx);
    if (gd >= 0) {
      if (col() == 0)
        return { OptimResultCode::kAbnormalTerm, iter + 1, fx, std::move(gx) };
 
      reset_memory();
      continue;
    }
 
    t() = x();
    r() = gx;
 
    auto lnsrl = impl_.lnsrch(x(), t(), z(), d(), fx, gd, iter);
    bool converged = false;
    for (int i = 0; !converged && i < maxls; ++i) {
      NURI_EIGEN_ALLOW_MALLOC(true);
      fx = fg(gx, x());
      NURI_EIGEN_ALLOW_MALLOC(false);
 
      gd = d_dot(gx);
      converged = lnsrl.search(fx, gd);
    }
 
    if (!converged) {
      x() = t();
      fx = lnsrl.finit();
      gx = r();
 
      if (col() == 0)
        return { OptimResultCode::kAbnormalTerm, iter + 1, fx, std::move(gx) };
 
      reset_memory();
      continue;
    }
 
    sbgnrm = impl_.projgr(x(), gx);
    if (sbgnrm <= pgtol)
      return { OptimResultCode::kSuccess, iter + 1, fx, std::move(gx) };
 
    double ddum = std::max({ std::abs(lnsrl.finit()), std::abs(fx), 1.0 });
    if (lnsrl.finit() - fx <= tol * ddum)
      return { OptimResultCode::kSuccess, iter + 1, fx, std::move(gx) };
 
    r() = gx - r();
    if (!prepare_next_iter(gd, lnsrl.ginit(), lnsrl.step(), lnsrl.dtd()))
      reset_memory();
  }
 
  return { OptimResultCode::kMaxIterReached, maxiter, fx, std::move(gx) };
}
 
namespace internal {
  class LbfgsbBounds {
  public:
    LbfgsbBounds(const ArrayXi &nbd, const Array2Xd &bounds)
        : nbd_(&nbd), bds_(&bounds) {
      check_sizes();
    }
 
    LbfgsbBounds(ArrayXi &&nbd, Array2Xd &&bounds) = delete;
 
    bool has_bound(int i) const { return nbd()[i] != 0; }
 
    bool has_lb(int i) const { return (nbd()[i] & 0x1) != 0; }
 
    bool has_ub(int i) const { return (nbd()[i] & 0x2) != 0; }
 
    bool has_both(int i) const { return nbd()[i] == (0x1 | 0x2); }
 
    int raw_nbd(int i) const { return nbd()[i]; }
 
    double lb(int i) const {
      ABSL_DCHECK(has_lb(i));
      return bds()(0, i);
    }
 
    double ub(int i) const {
      ABSL_DCHECK(has_ub(i));
      return bds()(1, i);
    }
 
  private:
    const ArrayXi &nbd() const { return *nbd_; }
 
    const Array2Xd &bds() const { return *bds_; }
 
    void check_sizes() const {
      ABSL_DCHECK(nbd().size() == bds().cols());
      ABSL_DCHECK(((nbd() >= 0) && (nbd() <= 3)).all());
      ABSL_DCHECK(
          ((nbd() != 3).transpose() || bds().row(0) <= bds().row(1)).all());
    }
 
    const ArrayXi *nbd_;
    const Array2Xd *bds_;
  };
 
  class LbfgsbLnsrch {
  public:
    LbfgsbLnsrch(MutRef<ArrayXd> &x, const ArrayXd &t, const ArrayXd &z,
                 const ArrayXd &d, const LbfgsbBounds &bounds, double f0,
                 double g0, int iter, bool constrained, bool boxed,
                 double ftol = 1e-3, double gtol = 0.9,
                 double xtol = 0.1) noexcept;
 
    bool search(double f, double g);
 
    double dtd() const { return dtd_; }
 
    double step() const { return dcsrch_.step(); }
 
    double xstep() const { return step() * dnorm_; }
 
    double finit() const { return dcsrch_.finit(); }
 
    double ginit() const { return dcsrch_.ginit(); }
 
  private:
    MutRef<ArrayXd> &x() { return *x_; }
 
    const ArrayXd &t() const { return *t_; }
 
    const ArrayXd &z() const { return *z_; }
 
    const ArrayXd &d() const { return *d_; }
 
    const LbfgsbBounds &bounds() const { return *bounds_; }
 
    void step_x();
 
    MutRef<ArrayXd> *x_;
    const ArrayXd *t_, *z_, *d_;
    const LbfgsbBounds *bounds_;
    double dtd_, dnorm_;
 
    Dcsrch dcsrch_;
  };
 
  struct CauchyBrkpt {
    int ibp;
    double tj;
  };
 
  inline bool operator>(CauchyBrkpt lhs, CauchyBrkpt rhs) {
    return lhs.tj > rhs.tj;
  }
 
  class LBfgsBImpl {
  public:
    LBfgsBImpl(MutRef<ArrayXd> x, LbfgsbBounds bounds, int m);
 
    LbfgsbLnsrch lnsrch(MutRef<ArrayXd> &x, const ArrayXd &t, const ArrayXd &z,
                        const ArrayXd &d, double f0, double g0, int iter,
                        double ftol = 1e-3, double gtol = 0.9,
                        double xtol = 0.1) const noexcept;
 
    void reset();
 
    bool prepare_lnsrch(LBfgsBase &lbfgsb, const ArrayXd &gx, double sbgnrm,
                        int iter);
 
    double projgr(ConstRef<ArrayXd> x, const ArrayXd &gx);
 
    // For internal use only
 
    MatrixXd &wn1() { return wn1_; }
 
    int n() const { return static_cast<int>(iwhere_.size()); }
 
    const LbfgsbBounds &bounds() const { return bounds_; }
 
    int nfree() const { return nfree_; }
 
    bool constrained() const { return constrained_; }
 
    bool boxed() const { return boxed_; }
 
    auto &iwhere() { return iwhere_; }
 
    auto free() const { return free_bound_.head(nfree_); }
 
    auto bound() const { return free_bound_.tail(n() - nfree_); }
 
    auto enter() const { return enter_leave_.head(nenter_); }
 
    auto leave() const { return enter_leave_.tail(nleave_); }
 
    auto &brks() { return brks_; }
 
    ArrayXi &free_bound() { return free_bound_; }
 
    void update_nfree(int nfree) { nfree_ = nfree; }
 
  private:
    bool freev(int iter);
 
    /* (2m, 2m) */
    MatrixXd wn1_;
 
    /* Bound constraints */
    LbfgsbBounds bounds_;
 
    /* n */
    ArrayXi iwhere_, free_bound_, enter_leave_;
 
    ClearablePQ<CauchyBrkpt, std::greater<>> brks_;
 
    int nfree_, nenter_, nleave_;
    bool constrained_, boxed_;
  };
 
  extern bool lbfgsb_errclb(const ArrayXd &x, const ArrayXi &nbd,
                            const Array2Xd &bounds, int m, double factr);
 
  extern bool lbfgsb_cauchy(LBfgsBase &lbfgsb, LBfgsBImpl &impl,
                            const ArrayXd &gx, double sbgnrm);
 
  extern bool lbfgsb_subsm(LBfgsBase &lbfgsb, LBfgsBImpl &impl,
                           const ArrayXd &gg);
}  // namespace internal

template <class FuncGrad>
LbfgsResult l_bfgs_b(FuncGrad &&fg, MutRef<ArrayXd> x, const ArrayXi &nbd,
                     const Array2Xd &bounds, const int m = 10,
                     const double factr = 1e+7, const double pgtol = 1e-5,
                     const int maxiter = 15000, const int maxls = 20) {
  bool args_ok = internal::lbfgsb_errclb(x, nbd, bounds, m, factr);
  if (!args_ok)
    return { OptimResultCode::kInvalidInput, 0, 0, {} };
 
  LBfgs<internal::LBfgsBImpl> lbfgsb(
      x, internal::LBfgsBImpl(x, { nbd, bounds }, m), m);
  return lbfgsb.minimize(std::forward<FuncGrad>(fg), factr, pgtol, maxiter,
                         maxls);
}
 
namespace internal {
  class LbfgsLnsrch {
  public:
    LbfgsLnsrch(MutRef<ArrayXd> &x, const ArrayXd &t, const ArrayXd &z,
                const ArrayXd &d, double f0, double g0, int iter,
                double ftol = 1e-3, double gtol = 0.9,
                double xtol = 0.1) noexcept;
 
    bool search(double f, double g);
 
    double dtd() const { return dtd_; }
 
    double step() const { return dcsrch_.step(); }
 
    double xstep() const { return step() * dnorm_; }
 
    double finit() const { return dcsrch_.finit(); }
 
    double ginit() const { return dcsrch_.ginit(); }
 
  private:
    MutRef<ArrayXd> &x() { return *x_; }
 
    const ArrayXd &t() const { return *t_; }
 
    const ArrayXd &z() const { return *z_; }
 
    const ArrayXd &d() const { return *d_; }
 
    void step_x();
 
    MutRef<ArrayXd> *x_;
    const ArrayXd *t_, *z_, *d_;
    double dtd_, dnorm_;
 
    Dcsrch dcsrch_;
  };
 
  class LBfgsImpl {
  public:
    LBfgsImpl(): LBfgsImpl(10) { }
 
    explicit LBfgsImpl(int m);
 
    static LbfgsLnsrch lnsrch(MutRef<ArrayXd> &x, const ArrayXd &t,
                              const ArrayXd &z, const ArrayXd &d, double f0,
                              double g0, int iter, double ftol = 1e-3,
                              double gtol = 0.9, double xtol = 0.1) noexcept;
 
    static void reset() { }
 
    bool prepare_lnsrch(LBfgsBase &lbfgs, const ArrayXd &gx, double sbgnrm,
                        int iter);
 
    static double projgr(ConstRef<ArrayXd> x, const ArrayXd &gx);
 
    MatrixXd &wn1() { return wn1_; }
 
  private:
    /* (2m, m) */
    MatrixXd wn1_;
  };
 
  extern bool lbfgs_errclb(const ArrayXd &x, int m, double factr);
}  // namespace internal

template <class FuncGrad>
LbfgsResult l_bfgs(FuncGrad &&fg, MutRef<ArrayXd> x, const int m = 10,
                   const double factr = 1e+7, const double pgtol = 1e-5,
                   const int maxiter = 15000, const int maxls = 20) {
  bool args_ok = internal::lbfgs_errclb(x, m, factr);
  if (!args_ok)
    return { OptimResultCode::kInvalidInput, 0, 0, {} };
 
  LBfgs<internal::LBfgsImpl> lbfgsb(x, internal::LBfgsImpl(m), m);
  return lbfgsb.minimize(std::forward<FuncGrad>(fg), factr, pgtol, maxiter,
                         maxls);
}
 
struct BfgsResult {
  OptimResultCode code;
  int niter;
  double fx;
  ArrayXd gx;
};

class Bfgs {
public:
  Bfgs(MutRef<ArrayXd> x);

  template <class FuncGrad>
  BfgsResult minimize(FuncGrad fg, const double pgtol = 1e-5,
                      const double xrtol = 0, int maxiter = -1,
                      const int maxls = 100, const double ftol = 1e-4,
                      const double gtol = 0.9, const double xtol = 1e-14) {
    using internal::Dcsrch;
    using internal::DcsrchStatus;
 
    if (maxiter < 0)
      maxiter = 200 * static_cast<int>(x().size());
 
    Hk_.setIdentity();
 
    ArrayXd gfk(x().size());
 
    const double f0 = fg(gfk, x());
    double gnorm = gfk.abs().maxCoeff();
    if (gnorm <= pgtol)
      return { OptimResultCode::kSuccess, 0, f0, std::move(gfk) };
 
    int k = 0;
    double fk = f0, fkm1 = fk + gfk.matrix().norm() * 0.5;
    for (; k < maxiter; ++k) {
      Dcsrch dcsrch = prepare_lnsrch(gfk, fk, fkm1, ftol, gtol, xtol);
      fkm1 = fk;
 
      bool success = false;
      for (int iter = 0; iter < maxls; ++iter) {
        xk() = x() + dcsrch.step() * pk().array();
        fk = fg(gfkp1(), xk());
 
        auto status = dcsrch(fk, gfkp1().matrix().dot(pk()));
        if (status == DcsrchStatus::kContinue)
          continue;
 
        success = true;
        break;
      }
      if (!success)
        return { OptimResultCode::kAbnormalTerm, k + 1, fk, std::move(gfk) };
 
      bool converged = prepare_next_iter(gfk, dcsrch.step(), pgtol, xrtol);
      if (converged)
        return { OptimResultCode::kSuccess, k + 1, fk, std::move(gfk) };
    }
 
    return { OptimResultCode::kMaxIterReached, maxiter, fk, std::move(gfk) };
  }
 
private:
  internal::Dcsrch prepare_lnsrch(const ArrayXd &gfk, double fk, double fkm1,
                                  double ftol, double gtol, double xtol);
 
  bool prepare_next_iter(ArrayXd &gfk, double step, double pgtol, double xrtol);
 
  MutRef<ArrayXd> &x() { return x_; }
 
  // NOLINTNEXTLINE(readability-identifier-naming)
  Eigen::SelfAdjointView<MatrixXd, Eigen::Upper> Hk() {
    return Hk_.selfadjointView<Eigen::Upper>();
  }
 
  VectorXd &pk() { return pk_; }
 
  ArrayXd &xk() { return xk_; }
  Eigen::MatrixWrapper<ArrayXd> sk() { return xk_.matrix(); }
 
  VectorXd &yk() { return yk_; }
 
  ArrayXd &gfkp1() { return gfkp1_; }
  // NOLINTNEXTLINE(readability-identifier-naming)
  Eigen::MatrixWrapper<ArrayXd> Hk_yk() { return gfkp1_.matrix(); }
 
  MutRef<ArrayXd> x_;
  ArrayXd xk_, gfkp1_;
 
  MatrixXd Hk_;
  VectorXd pk_, yk_;
};

template <class FuncGrad>
inline BfgsResult bfgs(FuncGrad &&fg, MutRef<ArrayXd> x,
                       const double pgtol = 1e-5, const double xrtol = 0,
                       int maxiter = -1, const int maxls = 100,
                       const double ftol = 1e-4, const double gtol = 0.9,
                       const double xtol = 1e-14) {
  Bfgs bfgs(x);
  return bfgs.minimize(std::forward<FuncGrad>(fg), pgtol, xrtol, maxiter, maxls,
                       ftol, gtol, xtol);
}
 
struct NMResult {
  OptimResultCode code;
  int argmin;
};

class NelderMead {
public:
  explicit NelderMead(MutRef<ArrayXXd> data);
 
  auto n() const { return data_.cols() - 1; }
 
  constexpr static int argmin() { return 0; }
  int argmax() const { return static_cast<int>(n()); }
 
  auto min() const { return data_.col(argmin()); }
  double minf() const { return min()[n()]; }
 
  auto max() { return data_.col(argmax()); }
  auto max() const { return data_.col(argmax()); }
  double maxf() const { return max()[n()]; }
 
  double max2f() const { return data_(n(), n() - 1); }

  template <class Func>
  NMResult minimize(Func f, int maxiter = -1, const double ftol = 1e-6,
                    const double alpha = 1, const double gamma = 2,
                    const double rho = 0.5, const double sigma = 0.5) {
    // NOLINTNEXTLINE(readability-identifier-naming)
    const auto N = n();
    if (maxiter <= 0)
      maxiter = static_cast<int>(N) * 200;
 
    auto eval_update = [&](auto &&simplexf) -> double {
      return simplexf[N] = f(simplexf.head(N));
    };
 
    for (int i = 0; i < data_.cols(); ++i)
      eval_update(data_.col(i));
 
    for (int iter = 0; iter < maxiter; ++iter) {
      partiton_min1_max2();
      if (maxf() - minf() < ftol)
        return { OptimResultCode::kSuccess, argmin() };
 
      centroid();
      eval_update(c_);
 
      reflection(alpha);
      double fr = eval_update(r_);
 
      if (fr < minf()) {
        expansion(ets_, gamma);
        const double fe = eval_update(ets_);
 
        max() = (fr < fe) ? r_ : ets_;
        continue;
      }
 
      if (fr < max2f()) {
        max() = r_;
        continue;
      }
 
      if (maxf() < fr) {
        r_ = max();
        fr = maxf();
      }
      contraction(ets_, rho);
      const double ft = eval_update(ets_);
 
      if (ft < fr) {
        max() = ets_;
        continue;
      }
 
      shrink(sigma);
      for (int i = 1; i < data_.cols(); ++i)
        eval_update(data_.col(i));
    }
 
    return { OptimResultCode::kMaxIterReached, argmin() };
  }
 
private:
  void partiton_min1_max2();
 
  void centroid();
 
  void reflection(double alpha);
 
  void expansion(ArrayXd &e, double gamma) const;
 
  void contraction(ArrayXd &t, double rho) const;
 
  void shrink(double sigma);
 
  /* (N + 1, N + 1) */
  MutRef<ArrayXXd> data_;
  /* (N + 1) */
  ArrayXd c_, r_, ets_;
};
 
namespace internal {
  extern bool nm_check_input(ConstRef<ArrayXXd> data, double alpha,
                             double gamma, double rho, double sigma);
}

template <class Func>
inline NMResult nelder_mead(Func &&f, MutRef<ArrayXXd> data, int maxiter = -1,
                            const double ftol = 1e-6, const double alpha = 1,
                            const double gamma = 2, const double rho = 0.5,
                            const double sigma = 0.5) {
  if (!internal::nm_check_input(data, alpha, gamma, rho, sigma))
    return { OptimResultCode::kInvalidInput, -1 };
 
  NelderMead nm(data);
  return nm.minimize(std::forward<Func>(f), maxiter, ftol, alpha, gamma, rho,
                     sigma);
}

extern ArrayXXd nm_prepare_simplex(ConstRef<ArrayXd> x0, double eps = 1e-6);
}  // namespace nuri
 
#endif /* NURI_ALGO_OPTIM_H_ */