TaoCtx.hpp

Go to the documentation of this file.

/** \file TaoCtx.hpp
 * \brief Context for PETSc TAO, i.e. optimization solver
 * \author Anonymous author(s) committing under MIT license
 */
 
#ifndef __TAOTX_HPP__
#define __TAOTX_HPP__
 
namespace MoFEM {
 
/** \brief Interface for TAO solvers
 * \ingroup mofem_petsc_solvers
 */
struct TaoCtx {
 
  MoFEM::Interface &mField; ///< database Interface
  moab::Interface &moab;    ///< moab Interface
  std::string problemName;  ///< problem name
  bool vErify;              ///< If true verify vector
 
  using PairNameFEMethodPtr = MoFEM::PairNameFEMethodPtr;
  using FEMethodsSequence = MoFEM::FEMethodsSequence;
  using BasicMethodsSequence = MoFEM::BasicMethodsSequence;
 
  FEMethodsSequence
      loopsHessian; ///< Sequence of finite element methods for assembling the
                    ///< Hessian (second derivative) matrix in TAO optimization
  FEMethodsSequence
      loopsObjective; ///< Sequence of finite element methods for evaluating the
                      ///< objective function in TAO optimization
  FEMethodsSequence
      loopsGradient; ///< Sequence of finite element methods for assembling the
                     ///< gradient (first derivative) vector in TAO optimization
 
  BasicMethodsSequence
      preHessian; ///< Sequence of methods run before Hessian is assembled
  BasicMethodsSequence
      postHessian; ///< Sequence of methods run after Hessian is assembled
 
  BasicMethodsSequence
      preObjective; ///< Sequence of methods run before objective is evaluated
  BasicMethodsSequence
      postObjective; ///< Sequence of methods run after objective is evaluated
 
  BasicMethodsSequence
      preGradient; ///< Sequence of methods run before gradient is assembled
  BasicMethodsSequence
      postGradient; ///< Sequence of methods run after gradient is assembled
 
  TaoCtx(Interface &m_field, const std::string &problem_name)
      : mField(m_field), moab(m_field.get_moab()), problemName(problem_name),
        vErify(false) {
    PetscLogEventRegister("LoopTAOHessian", 0, &MOFEM_EVENT_TaoHessian);
    PetscLogEventRegister("LoopTAOObjective", 0, &MOFEM_EVENT_TaoObjective);
    PetscLogEventRegister("LoopTAOGradient", 0, &MOFEM_EVENT_TaoGradient);
    if (!LogManager::checkIfChannelExist("TAO_WORLD")) {
      auto core_log = logging::core::get();
      core_log->add_sink(
          LogManager::createSink(LogManager::getStrmWorld(), "TAO_WORLD"));
      LogManager::setLog("TAO_WORLD");
      MOFEM_LOG_TAG("TAO_WORLD", "TAO");
    }
  }
 
  virtual ~TaoCtx() = default;
 
  // TAO Hessian (second derivative) matrix assembly loop and pre/post sequences
  FEMethodsSequence &getHessianLoops() { return loopsHessian; }
  BasicMethodsSequence &getPreProcHessian() { return preHessian; }
  BasicMethodsSequence &getPostProcHessian() { return postHessian; }
 
  // TAO Objective function evaluation loop and pre/post sequences
  FEMethodsSequence &getObjectiveLoops() { return loopsObjective; }
  BasicMethodsSequence &getPreProcObjective() { return preObjective; }
  BasicMethodsSequence &getPostProcObjective() { return postObjective; }
 
  // TAO Gradient (first derivative) vector assembly loop and pre/post sequences
  FEMethodsSequence &getGradientLoops() { return loopsGradient; }
  BasicMethodsSequence &getPreProcGradient() { return preGradient; }
  BasicMethodsSequence &getPostProcGradient() { return postGradient; }
 
  /**
   * \brief Copy sequences from another TaoCtx
   * @param  tao_ctx TaoCtx from which sequences are copied
   * @return         error code
   */
  MoFEMErrorCode copyLoops(const TaoCtx &tao_ctx);
 
  /**
   * @brief Clear loops
   *
   * @return MoFEMErrorCode
   */
  MoFEMErrorCode clearLoops();
 
  friend PetscErrorCode TaoSetObjective(Tao tao, Vec x, PetscReal *f,
                                        void *ctx);
  friend PetscErrorCode TaoSetGradient(Tao tao, Vec x, Vec g, void *ctx);
  friend PetscErrorCode TaoSetObjectiveAndGradient(Tao tao, Vec x, PetscReal *f,
                                                   Vec g, void *ctx);
  friend PetscErrorCode TaoSetHessian(Tao tao, Vec x, Mat H, Mat Hpre,
                                      void *ctx);
 
private:
  boost::movelib::unique_ptr<bool> vecAssembleSwitch;
  boost::movelib::unique_ptr<bool> matAssembleSwitch;
  PetscLogEvent MOFEM_EVENT_SnesRhs; ///< Log events to assemble residual
  PetscLogEvent MOFEM_EVENT_SnesMat; ///< Log events to assemble tangent matrix
  PetscLogEvent MOFEM_EVENT_TaoHessian;   ///< Log event for TAO Hessian
  PetscLogEvent MOFEM_EVENT_TaoGradient;  ///< Log event for TAO Gradient
  PetscLogEvent MOFEM_EVENT_TaoObjective; ///< Log event for TAO Objective
};
 
/**
 * @brief Sets the objective function value for a TAO optimization context.
 *
 * @param tao The TAO optimization solver context.
 * @param x The input vector at which to evaluate the objective.
 * @param f Pointer to store the computed objective value.
 * @param ctx User-defined context for the objective function.
 *
 * @return PetscErrorCode Error code (0 on success).
 */
PetscErrorCode TaoSetObjective(Tao tao, Vec x, PetscReal *f, void *ctx);
 
/**
 * @brief Sets the gradient vector for a TAO optimization context.
 *
 * @param tao The TAO optimization solver context.
 * @param x The input vector at which to evaluate the gradient.
 * @param g Pointer to the gradient vector.
 * @param ctx User-defined context for the gradient function.
 *
 * @return PetscErrorCode Error code (0 on success).
 */
PetscErrorCode TaoSetGradient(Tao tao, Vec x, Vec g, void *ctx);
 
/**
 * @brief Sets the objective function value and gradient for a TAO optimization
 * solver.
 *
 * @param tao The TAO solver context.
 * @param x Input vector at which to evaluate the objective and gradient.
 * @param f Pointer to store the computed objective function value.
 * @param g Vector to store the computed gradient.
 * @param ctx User-defined context for objective and gradient evaluation.
 *
 * @return PetscErrorCode Error code (0 on success).
 *
 * @note This function is typically called within a user-defined routine to
 * provide both the objective value and its gradient for optimization
 * algorithms.
 *
 * @see
 * https://petsc.org/release/docs/manualpages/Tao/TaoSetObjectiveAndGradient.html
 */
PetscErrorCode TaoSetObjectiveAndGradient(Tao tao, Vec x, PetscReal *f, Vec g,
                                          void *ctx);
 
/**
 * @brief Sets the Hessian matrix for a TAO optimization context.
 *
 * @param tao The TAO optimization solver context.
 * @param x The input vector at which to evaluate the Hessian.
 * @param H Pointer to the Hessian matrix.
 * @param Hpre Preconditioner matrix (can be NULL).
 * @param ctx User-defined context for the Hessian function.
 *
 * @return PetscErrorCode Error code (0 on success).
 */
PetscErrorCode TaoSetHessian(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx);
 
} // namespace MoFEM
 
#endif // __TAOTX_HPP__