nlcg.hpp

// FEAT3: Finite Element Analysis Toolbox, Version 3
// Copyright (C) 2010 by Stefan Turek & the FEAT group
// FEAT3 is released under the GNU General Public License version 3,
// see the file 'copyright.txt' in the top level directory for details.
 
#pragma once
#include <kernel/base_header.hpp>
#include <kernel/solver/base.hpp>
#include <kernel/solver/iterative.hpp>
#include <kernel/solver/linesearch.hpp>
#include <kernel/solver/nloptls.hpp>
#include <kernel/solver/nlopt_precond.hpp>
 
#include <deque>
namespace FEAT
{
  namespace Solver
  {
    enum class NLCGDirectionUpdate
    {
      undefined = 0,
      DaiYuan,
      DYHSHybrid,
      FletcherReeves,
      HagerZhang, // warning: Experimental, inefficient
      HestenesStiefel,
      PolakRibiere
    };
 
 
    inline std::ostream& operator<<(std::ostream& os, NLCGDirectionUpdate update)
    {
      switch(update)
      {
        case NLCGDirectionUpdate::undefined:
          return os << "undefined";
        case NLCGDirectionUpdate::DaiYuan:
          return os << "DaiYuan";
        case NLCGDirectionUpdate::DYHSHybrid:
          return os << "DYHSHybrid";
        case NLCGDirectionUpdate::FletcherReeves:
          return os << "FletcherReeves";
        case NLCGDirectionUpdate::HagerZhang:
          return os << "HagerZhang";
        case NLCGDirectionUpdate::HestenesStiefel:
          return os << "HestenesStiefel";
        case NLCGDirectionUpdate::PolakRibiere:
          return os << "PolakRibiere";
        default:
          return os << "-unknown-";
      }
    }
 
    inline void operator<<(NLCGDirectionUpdate& update, const String& update_name)
    {
        if(update_name == "undefined")
          update = NLCGDirectionUpdate::undefined;
        else if(update_name == "DaiYuan")
          update = NLCGDirectionUpdate::DaiYuan;
        else if(update_name == "DYHSHybrid")
          update = NLCGDirectionUpdate::DYHSHybrid;
        else if(update_name == "FletcherReeves")
          update = NLCGDirectionUpdate::FletcherReeves;
        else if(update_name == "HagerZhang")
          update = NLCGDirectionUpdate::HagerZhang;
        else if(update_name == "HestenesStiefel")
          update = NLCGDirectionUpdate::HestenesStiefel;
        else if(update_name == "PolakRibiere")
          update = NLCGDirectionUpdate::PolakRibiere;
        else
          XABORTM("Unknown NLCGDirectionUpdate identifier string " +update_name);
    }
 
 
    template<typename Functional_, typename Filter_>
    class NLCG : public NLOptLS<Functional_, Filter_>
    {
      public:
        typedef Functional_ FunctionalType;
        typedef Filter_ FilterType;
        typedef Solver::Linesearch<Functional_, Filter_> LinesearchType;
 
        typedef typename Functional_::GradientType GradientType;
        typedef typename Functional_::VectorTypeR VectorType;
        typedef typename Functional_::DataType DataType;
 
        typedef NLOptLS<FunctionalType, FilterType> BaseClass;
        typedef NLOptPrecond<typename Functional_::VectorTypeL, Filter_> PrecondType;
        static constexpr NLCGDirectionUpdate direction_update_default = NLCGDirectionUpdate::DYHSHybrid;
 
      protected:
        std::shared_ptr<LinesearchType> _linesearch;
        std::shared_ptr<PrecondType> _precond;
 
        NLCGDirectionUpdate _direction_update;
 
        VectorType _vec_r;
        VectorType _vec_p;
        VectorType _vec_z;
        VectorType _vec_y;
 
        Index _max_num_restarts;
        Index _num_restarts;
        Index _restart_freq;
 
      public:
        std::deque<VectorType>* iterates;
 
      public:
        explicit NLCG(Functional_& functional, Filter_& filter, std::shared_ptr<LinesearchType> linesearch,
        const NLCGDirectionUpdate du = direction_update_default,
        bool keep_iterates = false, std::shared_ptr<PrecondType> precond = nullptr) :
          BaseClass("NLCG", functional, filter, precond),
          _linesearch(linesearch),
          _precond(precond),
          _direction_update(du),
          _max_num_restarts(10),
          _num_restarts(0),
          _restart_freq(0),
          iterates(nullptr)
          {
            XASSERT(_linesearch != nullptr);
 
            this->_min_stag_iter = 0;
            _restart_freq = this->_functional.columns() + Index(3);
 
            this->set_ls_iter_digits(Math::ilog10(_linesearch->get_max_iter()));
 
            if(keep_iterates)
            {
              iterates = new std::deque<VectorType>;
            }
          }
 
        explicit NLCG(const String& section_name, const PropertyMap* section, Functional_& functional, Filter_& filter,
        std::shared_ptr<LinesearchType> linesearch, std::shared_ptr<PrecondType> precond = nullptr) :
          BaseClass("NLCG", section_name, section, functional, filter, precond),
          _linesearch(linesearch),
          _precond(precond),
          _direction_update(direction_update_default),
          _max_num_restarts(10),
          _num_restarts(0),
          _restart_freq(0),
          iterates(nullptr)
          {
            XASSERT(_linesearch != nullptr);
 
            this->_min_stag_iter = 0;
            _restart_freq = this->_functional.columns() + Index(3);
 
            this->set_ls_iter_digits(Math::ilog10(_linesearch->get_max_iter()));
 
            // Get direction update
            auto direction_update_p = section->query("direction_update");
            if(direction_update_p.second)
            {
              NLCGDirectionUpdate my_update;
              my_update << direction_update_p.first;
              set_direction_update(my_update);
            }
 
            // Check if we have to keep the iterates
            auto keep_iterates_p = section->query("keep_iterates");
            if(keep_iterates_p.second && std::stoul(keep_iterates_p.first) == 1)
            {
              iterates = new std::deque<VectorType>;
            }
 
            auto max_num_restarts_p = section->query("max_num_restarts");
            if(max_num_restarts_p.second)
            {
              set_max_num_restarts(Index(std::stoul(max_num_restarts_p.first)));
            }
 
          }
 
        virtual ~NLCG()
        {
          if(iterates != nullptr)
            delete iterates;
        }
 
        virtual void init_symbolic() override
        {
          BaseClass::init_symbolic();
          // create three temporary vectors
          _vec_r = this->_functional.create_vector_r();
          _vec_p = this->_functional.create_vector_r();
          _vec_y = this->_functional.create_vector_r();
          _vec_z = this->_functional.create_vector_r();
 
          _vec_z.format();
          _linesearch->init_symbolic();
        }
 
        virtual void done_symbolic() override
        {
          if(iterates != nullptr)
            iterates->clear();
 
          this->_vec_p.clear();
          this->_vec_r.clear();
          this->_vec_y.clear();
          this->_vec_z.clear();
          _restart_freq = Index(0);
          _linesearch->done_symbolic();
          BaseClass::done_symbolic();
        }
 
        virtual String name() const override
        {
          return "NLCG";
        }
 
        virtual Status apply(VectorType& vec_cor, const VectorType& vec_def) override
        {
          // clear solution vector
          vec_cor.format();
 
          // Evaluate the functional at the new state
          this->_functional.prepare(vec_cor, this->_filter);
          this->_functional.eval_fval_grad(this->_fval, this->_vec_r);
 
          // Copy back given defect
          this->_vec_r.copy(vec_def);
          //this->_system_filter.filter_def(this->_vec_r);
 
          // Prepare the preconditioner (if any)
          if(this->_precond != nullptr)
          {
            this->_precond->prepare(vec_cor, this->_filter);
          }
 
          // apply
          this->_status = _apply_intern(vec_cor);
          this->plot_summary();
          return this->_status;
        }
 
        virtual Status correct(VectorType& vec_sol, const VectorType& DOXY(vec_rhs)) override
        {
          //std::cout << std::scientific;
          //std::cout << std::setprecision(16);
          // Evaluate the functional at the new state
          this->_functional.prepare(vec_sol, this->_filter);
          // Compute defect
          this->_functional.eval_fval_grad(this->_fval, this->_vec_r);
          this->_vec_r.scale(this->_vec_r,DataType(-1));
          this->_filter.filter_def(this->_vec_r);
 
          // Prepare the preconditioner (if any)
          if(this->_precond != nullptr)
          {
            this->_precond->prepare(vec_sol, this->_filter);
          }
 
          // apply
          this->_status = _apply_intern(vec_sol);
          this->plot_summary();
          return this->_status;
        }
 
        void set_keep_iterates(bool keep_iterates)
        {
          if(iterates != nullptr)
          {
            delete iterates;
          }
 
          if(keep_iterates)
          {
            iterates = new std::deque<VectorType>;
          }
 
        }
 
        void set_max_num_restarts(Index max_num_restarts)
        {
          XASSERT(_max_num_restarts > Index(1));
 
          _max_num_restarts = max_num_restarts;
        }
 
        void set_restart_freq(Index restart_freq)
        {
          XASSERT(_restart_freq > Index(1));
 
          _restart_freq = restart_freq;
        }
 
        void set_direction_update(NLCGDirectionUpdate update_)
        {
          _direction_update = update_;
        }
 
      protected:
        virtual Status _apply_intern(VectorType& vec_sol)
        {
          IterationStats pre_iter(*this);
          Statistics::add_solver_expression(std::make_shared<ExpressionStartSolve>(this->name()));
 
          // p[k+1] <- r[k+1] + _beta * p[k+1]
          DataType beta;
          // eta[k] = <r[k], p[k]>
          DataType eta;
          // Reset member variables in the LineSearch
          _linesearch->reset();
 
          if(iterates != nullptr)
          {
            iterates->push_back(std::move(vec_sol.clone()));
          }
 
          // Set initial defect. The defect vector was calculated in the calling function
          Status status = this->_set_initial_defect(this->_vec_r, vec_sol);
          if(status != Status::progress)
          {
            Statistics::add_solver_expression(std::make_shared<ExpressionEndSolve>(this->name(), status, this->get_num_iter()));
            return status;
          }
 
          // apply preconditioner to defect vector
          if(!this->_apply_precond(this->_vec_z, this->_vec_r, this->_filter))
          {
            Statistics::add_solver_expression(std::make_shared<ExpressionEndSolve>(this->name(), Status::aborted, this->get_num_iter()));
            return Status::aborted;
          }
          //this->_vec_z.copy(this->_vec_r);
 
          // The first direction has to be the (preconditioned) steepest descent direction
          this->_vec_p.copy(this->_vec_z);
 
 
          //DataType max_elem(this->_vec_p.max_element());
          //DataType scale1 = DataType(1)/max_elem;
          //std::cout << "max_elem   " << max_elem << " " << scale1 << "\n";
          //DataType scale2(DataType(1)/new_norm);
          //std::cout << " new norm  " << new_norm << " " << scale2 << "\n";
 
          // Compute initial eta = < p, r>
          eta = this->_vec_p.dot(this->_vec_r);
 
          // If the preconditioner was not pd in the first step, reset the search direction to steepest descent
          if(eta <= DataType(0))
          {
            this->_vec_p.copy(this->_vec_r);
          }
 
          // compute initial gamma = < z[0], r[0] >
          DataType gamma(this->_vec_z.dot(this->_vec_r));
          DataType gamma_prev(0);
 
          if(this->_def_init <= this->_tol_rel)
          {
            Statistics::add_solver_expression(std::make_shared<ExpressionEndSolve>(this->name(), Status::success, this->get_num_iter()));
            return Status::success;
          }
 
          Index its_since_restart(0);
          _num_restarts = Index(0);
          Index first_restart(_restart_freq+Index(1));
 
          pre_iter.destroy();
 
          // start iterating
          while(status == Status::progress)
          {
            IterationStats stat(*this);
 
            // Clear statistics if we have a preconditioner. Otherwise this tends to overflow the main memory if the
            // preconditioner does too many iterations
            // \todo: fix this
            //FEAT::Statistics::reset_solver_statistics();
 
            this->_fval_prev = this->_fval;
 
            this->_vec_y.copy(this->_vec_r);
 
            // Copy information to the linesearch
            _linesearch->set_initial_fval(this->_fval);
            _linesearch->set_grad_from_defect(this->_vec_r);
            // If we are using a preconditioner, use the additional information about the preconditioned step length in the line search
            if(this->_precond != nullptr)
            {
              _linesearch->set_dir_scaling(true);
            }
 
            // Call the linesearch to update vec_sol
            Status linesearch_status = _linesearch->correct(vec_sol, this->_vec_p);
 
            // Copy back information from the linesearch
            this->_fval = _linesearch->get_final_fval();
            this->_ls_its = _linesearch->get_num_iter();
            this->_steplength = _linesearch->get_rel_update();
            _linesearch->get_defect_from_grad(this->_vec_r);
            this->_vec_y.scale(this->_vec_y, -DataType(1));
            this->_vec_y.axpy(this->_vec_r); 
 
            // Log iterates if necessary
            if(iterates != nullptr)
            {
              iterates->push_back(vec_sol.clone());
            }
 
            // Compute defect norm. This also performs the convergence/divergence checks.
            status = this->_set_new_defect(this->_vec_r, vec_sol);
 
            if(status != Status::progress)
            {
              stat.destroy();
              Statistics::add_solver_expression(std::make_shared<ExpressionEndSolve>(this->name(), status, this->get_num_iter()));
              return status;
            }
 
            // Update preconditioner if necessary
            if(this->_precond != nullptr)
            {
              this->_precond->prepare(vec_sol, this->_filter);
            }
 
            // apply preconditioner
            if(!this->_apply_precond(_vec_z, _vec_r, this->_filter))
            {
              stat.destroy();
              Statistics::add_solver_expression(std::make_shared<ExpressionEndSolve>(this->name(), Status::aborted, this->get_num_iter()));
              return Status::aborted;
            }
 
            // Save old gamma and compute new
            gamma_prev = gamma;
            gamma = this->_vec_r.dot(this->_vec_z);
 
            //String info("");
            bool restart(false);
            // We need to perform a steepest descent step if the Wolfe conditions do not hold (linesearch status
            // != success).
            if(linesearch_status != Status::success)
            {
              //info += " Wolfe Conditions do not hold";
              restart = true;
              _num_restarts++;
            }
            else
            {
              _num_restarts = Index(0);
            }
 
            if(_restart_freq > Index(0) && this->_num_iter >= first_restart && its_since_restart%_restart_freq == 0)
            {
              //info += " Scheduled restart";
              restart = true;
              its_since_restart = Index(0);
            }
            // This needs to be done after all the checks
            ++its_since_restart;
 
            // Set beta = 0 or compute depending on the restart flag
            restart ? beta = DataType(0) : beta = compute_beta(gamma, gamma_prev);
 
            //std::cout << "Beta " << beta << info << "\n";
 
            // We need to check beta again here as some direction updates might set it to zero
            // Discard the old search direction and perform (preconditioned) steepest descent
            if(beta == DataType(0))
            {
               this->_vec_p.copy(this->_vec_z);
            }
            else
            {
              this->_vec_p.scale(this->_vec_p, beta);
              this->_vec_p.axpy(this->_vec_z); 
            }
 
 
            //max_elem = this->_vec_p.max_element();
            //scale1 = DataType(1)/max_elem;
            //std::cout << "max_elem   " << max_elem << " " << scale1 << "\n";
            //scale2 = DataType(1)/new_norm;
            //std::cout << " new norm  " << new_norm << " " << scale2 << "\n";
 
            eta = this->_vec_p.dot(this->_vec_r);
 
            // Safeguard as this should not happen
            // TODO: Correct the output of the preconditioner if it turns out to not have been positive definite
            if(eta <= DataType(0))
            {
              this->_vec_p.copy(this->_vec_r);
              //return Status::aborted;
            }
 
          }
 
          // We should never come to this point
          Statistics::add_solver_expression(std::make_shared<ExpressionEndSolve>(this->name(), Status::undefined, this->get_num_iter()));
          return Status::undefined;
        }
 
        DataType compute_beta(const DataType& gamma, const DataType& gamma_prev) const
        {
          switch(_direction_update)
          {
            case NLCGDirectionUpdate::DaiYuan:
              return dai_yuan(gamma);
            case NLCGDirectionUpdate::DYHSHybrid:
              return dy_hs_hybrid(gamma);
            case NLCGDirectionUpdate::FletcherReeves:
              return fletcher_reeves(gamma, gamma_prev);
            case NLCGDirectionUpdate::HagerZhang:
              // Warning: This update is not very efficient in its current implementation.
              return hager_zhang(gamma);
            case NLCGDirectionUpdate::HestenesStiefel:
              return hestenes_stiefel();
            case NLCGDirectionUpdate::PolakRibiere:
              return polak_ribiere(gamma_prev);
            default:
              XABORTM("Unhandled direction update: "+stringify(_direction_update));
          }
        }
 
        virtual Status _set_new_defect(const VectorType& vec_r, const VectorType& vec_sol) override
        {
 
          Status st = BaseClass::_set_new_defect(vec_r, vec_sol);
 
          if(st != Status::progress)
          {
            return st;
          }
 
          // If there were too many subsequent restarts, the solver is stagnated
          if(_max_num_restarts > Index(0) && _num_restarts > _max_num_restarts)
          {
            return Status::stagnated;
          }
 
          // continue iterating
          return Status::progress;
        }
 
        DataType dai_yuan(const DataType gamma) const
        {
          // <r[k+1] - r[k], p[k]>
          DataType eta_dif = this->_vec_p.dot(this->_vec_y);
 
          return gamma/(-eta_dif);
        }
 
        DataType dy_hs_hybrid(const DataType gamma) const
        {
 
          // <r[k+1] - r[k], z[k+1]>
          DataType gamma_dif = this->_vec_z.dot(this->_vec_y);
          // <r[k+1] - r[k], p[k]>
          DataType eta_dif = this->_vec_p.dot(this->_vec_y);
 
          DataType beta(Math::min(gamma/(-eta_dif), gamma_dif/(-eta_dif)));
          beta = Math::max(DataType(0), beta);
 
          return beta;
 
        }
 
        DataType fletcher_reeves(const DataType gamma, const DataType gamma_prev) const
        {
          return gamma/gamma_prev;
        }
 
        DataType hager_zhang(const DataType gamma) const
        {
 
          DataType thresh = DataType(0.01);
 
          // <r[k+1] - r[k], p[k]>
          DataType eta_dif = this->_vec_p.dot(this->_vec_y);
          // <r[k+1] - r[k], z[k+1]>
          DataType gamma_dif = this->_vec_y.dot(this->_vec_z);
          // <p[k], z[k+1]>
          DataType omega_mid = this->_vec_p.dot(this->_vec_z);
          // || r[k+1] - r[k]
          DataType norm_y = this->_vec_y.norm2();
 
          DataType beta = -( gamma_dif/(eta_dif) + DataType(2)*(norm_y)*omega_mid/Math::sqr(eta_dif));
          beta = Math::max(beta,-DataType(1)/(this->_vec_p.norm2()*Math::min(thresh, Math::sqrt(gamma))));
 
          return beta;
        }
 
        DataType hestenes_stiefel() const
        {
          // <r[k+1] - r[k], p[k]>
          DataType eta_dif = this->_vec_p.dot(this->_vec_y);
          // <r[k+1] - r[k], z[k+1]>
          DataType gamma_dif = this->_vec_z.dot(this->_vec_y);
 
          return gamma_dif/( -eta_dif);
        }
 
        DataType polak_ribiere(const DataType gamma_prev) const
        {
          // <r[k+1] - r[k], z[k+1]>
          DataType gamma_dif = this->_vec_z.dot(this->_vec_y);
 
          return  Math::max(gamma_dif/gamma_prev, DataType(0));
        }
 
    }; // class NLCG
 
    template<typename Functional_, typename Filter_, typename Linesearch_>
    inline std::shared_ptr<NLCG<Functional_, Filter_>> new_nlcg(
      Functional_& functional, Filter_& filter, Linesearch_& linesearch,
      NLCGDirectionUpdate direction_update = NLCG<Functional_, Filter_>::direction_update_default,
      bool keep_iterates = false,
      std::shared_ptr<NLOptPrecond<typename Functional_::VectorTypeL, Filter_>> precond = nullptr)
      {
        return std::make_shared<NLCG<Functional_, Filter_>>(functional, filter, linesearch, direction_update,
        keep_iterates, precond);
      }
 
    template<typename Functional_, typename Filter_, typename Linesearch_>
    inline std::shared_ptr<NLCG<Functional_, Filter_>> new_nlcg(
      const String& section_name, const PropertyMap* section,
      Functional_& functional, Filter_& filter, Linesearch_& linesearch,
      std::shared_ptr<NLOptPrecond<typename Functional_::VectorTypeL, Filter_>> precond = nullptr)
      {
        return std::make_shared<NLCG<Functional_, Filter_>>(section_name, section, functional, filter, linesearch,
        precond);
      }
  } //namespace Solver
} // namespace FEAT