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;
  using ObjectiveHook =
      boost::function<MoFEMErrorCode(Tao tao, Vec x, PetscReal *objective)>;
  using GradientHook =
      boost::function<MoFEMErrorCode(Tao tao, Vec x, Vec gradient)>;
  using ObjectiveAndGradientHook = boost::function<MoFEMErrorCode(
      Tao tao, Vec x, PetscReal *objective, Vec gradient)>;
 
  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
  FEMethodsSequence loopsEqualityConstraint;   ///< Sequence of finite element
                                               ///< methods for evaluating
                                               ///< equality constraints
  FEMethodsSequence loopsInequalityConstraint; ///< Sequence of finite element
                                               ///< methods for evaluating
                                               ///< inequality constraints
  FEMethodsSequence loopsEqualityJacobian;     ///< Sequence of finite element
                                               ///< methods for assembling the
                                               ///< equality constraint Jacobian
  FEMethodsSequence loopsInequalityJacobian;   ///< Sequence of finite element
                                               ///< methods for assembling the
                                               ///< inequality constraint Jacobian
 
  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
 
  BasicMethodsSequence preEqualityConstraint;   ///< Sequence of methods run
                                                ///< before equality constraints
                                                ///< are evaluated
  BasicMethodsSequence postEqualityConstraint;  ///< Sequence of methods run
                                                ///< after equality constraints
                                                ///< are evaluated
  BasicMethodsSequence preInequalityConstraint; ///< Sequence of methods run
                                                ///< before inequality
                                                ///< constraints are evaluated
  BasicMethodsSequence postInequalityConstraint; ///< Sequence of methods run
                                                 ///< after inequality
                                                 ///< constraints are evaluated
  BasicMethodsSequence preEqualityJacobian;  ///< Sequence of methods run before
                                             ///< the equality Jacobian is
                                             ///< assembled
  BasicMethodsSequence postEqualityJacobian; ///< Sequence of methods run after
                                             ///< the equality Jacobian is
                                             ///< assembled
  BasicMethodsSequence preInequalityJacobian;  ///< Sequence of methods run
                                               ///< before the inequality
                                               ///< Jacobian is assembled
  BasicMethodsSequence postInequalityJacobian; ///< Sequence of methods run
                                               ///< after the inequality
                                               ///< Jacobian 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);
    PetscLogEventRegister("LoopTAOEqualityConstraint", 0,
                          &MOFEM_EVENT_TaoEqualityConstraint);
    PetscLogEventRegister("LoopTAOInequalityConstraint", 0,
                          &MOFEM_EVENT_TaoInequalityConstraint);
    PetscLogEventRegister("LoopTAOEqualityJacobian", 0,
                          &MOFEM_EVENT_TaoEqualityJacobian);
    PetscLogEventRegister("LoopTAOInequalityJacobian", 0,
                          &MOFEM_EVENT_TaoInequalityJacobian);
    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; }
 
  // TAO equality constraint evaluation loop and pre/post sequences
  FEMethodsSequence &getEqualityConstraintLoops() {
    return loopsEqualityConstraint;
  }
  BasicMethodsSequence &getPreProcEqualityConstraint() {
    return preEqualityConstraint;
  }
  BasicMethodsSequence &getPostProcEqualityConstraint() {
    return postEqualityConstraint;
  }
 
  // TAO inequality constraint evaluation loop and pre/post sequences
  FEMethodsSequence &getInequalityConstraintLoops() {
    return loopsInequalityConstraint;
  }
  BasicMethodsSequence &getPreProcInequalityConstraint() {
    return preInequalityConstraint;
  }
  BasicMethodsSequence &getPostProcInequalityConstraint() {
    return postInequalityConstraint;
  }
 
  // TAO equality constraint Jacobian loop and pre/post sequences
  FEMethodsSequence &getEqualityJacobianLoops() {
    return loopsEqualityJacobian;
  }
  BasicMethodsSequence &getPreProcEqualityJacobian() {
    return preEqualityJacobian;
  }
  BasicMethodsSequence &getPostProcEqualityJacobian() {
    return postEqualityJacobian;
  }
 
  // TAO inequality constraint Jacobian loop and pre/post sequences
  FEMethodsSequence &getInequalityJacobianLoops() {
    return loopsInequalityJacobian;
  }
  BasicMethodsSequence &getPreProcInequalityJacobian() {
    return preInequalityJacobian;
  }
  BasicMethodsSequence &getPostProcInequalityJacobian() {
    return postInequalityJacobian;
  }
 
  /**
   * @brief Get an optional objective evaluator.
   *
   * When set, MoFEM::TaoSetObjective delegates the complete evaluation to this
   * hook before scattering controls to the mesh or running element loops.
   * The hook owns state synchronization and objective evaluation.
   */
  ObjectiveHook &getObjectiveHook() { return objectiveHook; }
 
  /**
   * @brief Get an optional gradient evaluator.
   *
   * When set, MoFEM::TaoSetGradient delegates the complete evaluation to this
   * hook before scattering controls to the mesh or running element loops.
   * The hook owns state synchronization and gradient assembly.
   */
  GradientHook &getGradientHook() { return gradientHook; }
 
  /**
   * @brief Get an optional combined objective-and-gradient evaluator.
   *
   * When set, MoFEM::TaoSetObjectiveAndGradient delegates the complete
   * evaluation to this hook. This supports reduced objectives which require
   * an application-owned state solve before returning the value and gradient.
   * Otherwise, the separate objective and gradient callbacks are used.
   */
  ObjectiveAndGradientHook &getObjectiveAndGradientHook() {
    return objectiveAndGradientHook;
  }
 
  /**
   * \brief Copy sequences from another TaoCtx
   *
   * Evaluator hooks are unchanged.
   * @param  tao_ctx TaoCtx from which sequences are copied
   * @return         error code
   */
  MoFEMErrorCode copyLoops(const TaoCtx &tao_ctx);
 
  /**
   * @brief Clear loops
   *
   * Evaluator hooks are unchanged.
   *
   * @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);
  friend PetscErrorCode TaoSetEqualityConstraints(Tao tao, Vec x, Vec c,
                                                  void *ctx);
  friend PetscErrorCode TaoSetInequalityConstraints(Tao tao, Vec x, Vec c,
                                                    void *ctx);
  friend PetscErrorCode TaoSetJacobianEquality(Tao tao, Vec x, Mat J, Mat Jpre,
                                               void *ctx);
  friend PetscErrorCode TaoSetJacobianInequality(Tao tao, Vec x, Mat J,
                                                 Mat Jpre, void *ctx);
 
private:
  PetscErrorCode evaluateConstraint(Tao tao, Vec x, Vec c,
                                    FEMethodsSequence &loops,
                                    BasicMethodsSequence &pre_process,
                                    BasicMethodsSequence &post_process,
                                    TaoMethod::TAOContext evaluation_context,
                                    PetscLogEvent log_event);
  PetscErrorCode evaluateConstraintJacobian(
      Tao tao, Vec x, Mat J, Mat Jpre, FEMethodsSequence &loops,
      BasicMethodsSequence &pre_process, BasicMethodsSequence &post_process,
      TaoMethod::TAOContext evaluation_context, PetscLogEvent log_event);
 
  ObjectiveHook objectiveHook;
  GradientHook gradientHook;
  ObjectiveAndGradientHook objectiveAndGradientHook;
  boost::movelib::unique_ptr<bool> vecAssembleSwitch;
  boost::movelib::unique_ptr<bool> matAssembleSwitch;
  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
  PetscLogEvent MOFEM_EVENT_TaoEqualityConstraint;
  PetscLogEvent MOFEM_EVENT_TaoInequalityConstraint;
  PetscLogEvent MOFEM_EVENT_TaoEqualityJacobian;
  PetscLogEvent MOFEM_EVENT_TaoInequalityJacobian;
};
 
/**
 * @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);
 
/**
 * @brief Evaluate equality constraints for a TAO optimization context.
 *
 * Register this callback with TaoSetEqualityConstraintsRoutine().
 */
PetscErrorCode TaoSetEqualityConstraints(Tao tao, Vec x, Vec c, void *ctx);
 
/**
 * @brief Evaluate inequality constraints for a TAO optimization context.
 *
 * Register this callback with TaoSetInequalityConstraintsRoutine(). PETSc
 * defines feasible inequality constraints by c(x) >= 0.
 */
PetscErrorCode TaoSetInequalityConstraints(Tao tao, Vec x, Vec c, void *ctx);
 
/**
 * @brief Assemble the equality constraint Jacobian.
 *
 * Register this callback with TaoSetJacobianEqualityRoutine().
 */
PetscErrorCode TaoSetJacobianEquality(Tao tao, Vec x, Mat J, Mat Jpre,
                                      void *ctx);
 
/**
 * @brief Assemble the inequality constraint Jacobian.
 *
 * Register this callback with TaoSetJacobianInequalityRoutine().
 */
PetscErrorCode TaoSetJacobianInequality(Tao tao, Vec x, Mat J, Mat Jpre,
                                        void *ctx);
 
} // namespace MoFEM
 
#endif // __TAOTX_HPP__