tao_ctx_constraints.cpp

Go to the documentation of this file.

/** \file tao_ctx_constraints.cpp
 * \brief Test TAO constraint pipelines and application evaluation hooks
 */
 
#include <MoFEM.hpp>
 
using namespace MoFEM;
 
static char help[] = "Test TaoCtx constraint callbacks.\n\n";
 
int main(int argc, char *argv[]) {
  MoFEM::Core::Initialize(&argc, &argv, nullptr, help);
 
  try {
    moab::Core moab_core;
    MoFEM::Core mofem_core(moab_core);
    MoFEM::Interface &m_field = mofem_core;
 
    CHKERR DMRegister_MoFEM("DMMOFEM");
 
    auto simple = m_field.getInterface<Simple>();
    CHKERR simple->getOptions();
    CHKERR simple->loadFile();
    CHKERR simple->addDomainField("X", H1, AINSWORTH_LEGENDRE_BASE, 1);
    CHKERR simple->setFieldOrder("X", 1);
    CHKERR simple->setUp();
 
    auto dm = simple->getDM();
    auto x = createDMVector(dm);
    auto c = vectorDuplicate(x);
    auto J = createDMMatrix(dm);
    auto diagonal = vectorDuplicate(x);
 
    Tao tao;
    CHKERR TaoCreate(PETSC_COMM_WORLD, &tao);
    CHKERR TaoSetType(tao, TAOALMM);
    CHKERR TaoSetSolution(tao, x);
 
    TaoCtx tao_ctx(m_field, simple->getProblemName());
 
    auto check_vector_pipeline = [&](auto context, const double value,
                                     auto &pre_process, auto &post_process) {
      auto pre = boost::make_shared<BasicMethod>();
      auto post = boost::make_shared<BasicMethod>();
      auto *pre_ptr = pre.get();
      auto *post_ptr = post.get();
 
      pre->preProcessHook = [pre_ptr, context, value, c]() {
        MoFEMFunctionBegin;
        if (pre_ptr->tao_ctx != context || pre_ptr->tao_c != c)
          SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                  "Incorrect TAO constraint context");
        CHKERR VecSet(pre_ptr->tao_c, value);
        MoFEMFunctionReturn(0);
      };
      post->postProcessHook = [post_ptr, context, c]() {
        MoFEMFunctionBegin;
        if (post_ptr->tao_ctx != context || post_ptr->tao_c != c)
          SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                  "Incorrect TAO constraint context");
        MoFEMFunctionReturn(0);
      };
 
      pre_process.push_back(pre);
      post_process.push_back(post);
    };
 
    auto check_jacobian_pipeline = [&](auto context, const double value,
                                       auto &pre_process, auto &post_process) {
      auto pre = boost::make_shared<BasicMethod>();
      auto post = boost::make_shared<BasicMethod>();
      auto *pre_ptr = pre.get();
      auto *post_ptr = post.get();
 
      pre->preProcessHook = [pre_ptr, context, value, J]() {
        MoFEMFunctionBegin;
        if (pre_ptr->tao_ctx != context || pre_ptr->tao_J != J ||
            pre_ptr->tao_Jpre != J)
          SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                  "Incorrect TAO constraint Jacobian context");
        CHKERR MatShift(pre_ptr->tao_Jpre, value);
        MoFEMFunctionReturn(0);
      };
      post->postProcessHook = [post_ptr, context, J]() {
        MoFEMFunctionBegin;
        if (post_ptr->tao_ctx != context || post_ptr->tao_J != J ||
            post_ptr->tao_Jpre != J)
          SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                  "Incorrect TAO constraint Jacobian context");
        MoFEMFunctionReturn(0);
      };
 
      pre_process.push_back(pre);
      post_process.push_back(post);
    };
 
    check_vector_pipeline(TaoMethod::CTX_TAO_EQUALITY_CONSTRAINT, 2.0,
                          tao_ctx.getPreProcEqualityConstraint(),
                          tao_ctx.getPostProcEqualityConstraint());
    check_vector_pipeline(TaoMethod::CTX_TAO_INEQUALITY_CONSTRAINT, 3.0,
                          tao_ctx.getPreProcInequalityConstraint(),
                          tao_ctx.getPostProcInequalityConstraint());
    check_jacobian_pipeline(TaoMethod::CTX_TAO_EQUALITY_JACOBIAN, 4.0,
                            tao_ctx.getPreProcEqualityJacobian(),
                            tao_ctx.getPostProcEqualityJacobian());
    check_jacobian_pipeline(TaoMethod::CTX_TAO_INEQUALITY_JACOBIAN, 5.0,
                            tao_ctx.getPreProcInequalityJacobian(),
                            tao_ctx.getPostProcInequalityJacobian());
 
    CHKERR TaoSetEqualityConstraintsRoutine(
        tao, c, MoFEM::TaoSetEqualityConstraints, &tao_ctx);
    CHKERR TaoSetInequalityConstraintsRoutine(
        tao, c, MoFEM::TaoSetInequalityConstraints, &tao_ctx);
    CHKERR TaoSetJacobianEqualityRoutine(
        tao, J, J, MoFEM::TaoSetJacobianEquality, &tao_ctx);
    CHKERR TaoSetJacobianInequalityRoutine(
        tao, J, J, MoFEM::TaoSetJacobianInequality, &tao_ctx);
 
    auto check_vector = [&](const double expected) {
      MoFEMFunctionBegin;
      CHKERR VecShift(c, -expected);
      PetscReal norm;
      CHKERR VecNorm(c, NORM_INFINITY, &norm);
      if (norm > 10 * PETSC_MACHINE_EPSILON)
        SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                "Incorrect TAO constraint vector");
      MoFEMFunctionReturn(0);
    };
 
    auto check_jacobian = [&](const double expected) {
      MoFEMFunctionBegin;
      CHKERR MatGetDiagonal(J, diagonal);
      CHKERR VecShift(diagonal, -expected);
      PetscReal norm;
      CHKERR VecNorm(diagonal, NORM_INFINITY, &norm);
      if (norm > 10 * PETSC_MACHINE_EPSILON)
        SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                "Incorrect TAO constraint Jacobian");
      MoFEMFunctionReturn(0);
    };
 
    CHKERR TaoComputeEqualityConstraints(tao, x, c);
    CHKERR check_vector(2.0);
    CHKERR TaoComputeInequalityConstraints(tao, x, c);
    CHKERR check_vector(3.0);
    CHKERR TaoComputeJacobianEquality(tao, x, J, J);
    CHKERR check_jacobian(4.0);
    CHKERR TaoComputeJacobianInequality(tao, x, J, J);
    CHKERR check_jacobian(5.0);
 
    TaoCtx copied_ctx(m_field, simple->getProblemName());
    CHKERR copied_ctx.copyLoops(tao_ctx);
    if (copied_ctx.getPreProcEqualityConstraint().empty() ||
        copied_ctx.getPreProcInequalityConstraint().empty() ||
        copied_ctx.getPreProcEqualityJacobian().empty() ||
        copied_ctx.getPreProcInequalityJacobian().empty())
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "TAO constraint pipelines were not copied");
    CHKERR copied_ctx.clearLoops();
    if (!copied_ctx.getPreProcEqualityConstraint().empty() ||
        !copied_ctx.getPreProcInequalityConstraint().empty() ||
        !copied_ctx.getPreProcEqualityJacobian().empty() ||
        !copied_ctx.getPreProcInequalityJacobian().empty())
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "TAO constraint pipelines were not cleared");
 
    // A reduced objective owns its mesh state and can reuse its gradient.
    // Neither wrapper may scatter a trial or clear that cached result.
    CHKERR VecZeroEntries(x);
    CHKERR m_field.getInterface<VecManager>()->setLocalGhostVector(
        simple->getProblemName(), COL, x, INSERT_VALUES, SCATTER_REVERSE);
    CHKERR VecSet(x, 2.0);
    auto gradient = vectorDuplicate(x);
    auto expected_gradient = vectorDuplicate(x);
    CHKERR VecSet(gradient, 7.0);
    CHKERR VecCopy(gradient, expected_gradient);
    PetscInt objective_calls = 0, gradient_calls = 0;
    tao_ctx.getObjectiveHook() = [&](Tao, Vec trial, PetscReal *objective) {
      MoFEMFunctionBegin;
      if (trial != x)
        SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                "Incorrect TAO objective-hook controls");
      ++objective_calls;
      *objective = 3.0;
      MoFEMFunctionReturn(0);
    };
    tao_ctx.getGradientHook() = [&](Tao, Vec trial, Vec g) {
      MoFEMFunctionBegin;
      PetscBool equal;
      CHKERR VecEqual(g, expected_gradient, &equal);
      if (trial != x || !equal)
        SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
                "TAO gradient hook did not receive the cached gradient");
      ++gradient_calls;
      MoFEMFunctionReturn(0);
    };
    CHKERR ::TaoSetObjective(tao, MoFEM::TaoSetObjective, &tao_ctx);
    CHKERR ::TaoSetGradient(tao, gradient, MoFEM::TaoSetGradient, &tao_ctx);
    CHKERR ::TaoSetObjectiveAndGradient(
        tao, gradient, MoFEM::TaoSetObjectiveAndGradient, &tao_ctx);
    PetscReal objective;
    CHKERR TaoComputeObjective(tao, x, &objective);
    if (objective != 3.0 || objective_calls != 1 || gradient_calls != 0)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Objective-only wrapper did not dispatch exclusively to its hook");
    CHKERR TaoComputeGradient(tao, x, gradient);
    if (objective_calls != 1 || gradient_calls != 1)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Gradient-only wrapper did not dispatch exclusively to its hook");
    CHKERR TaoComputeObjectiveAndGradient(tao, x, &objective, gradient);
    PetscBool equal;
    CHKERR VecEqual(gradient, expected_gradient, &equal);
    if (objective != 3.0 || objective_calls != 2 || gradient_calls != 2 ||
        !equal)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "Combined wrapper did not fall back to the separate hooks");
    auto mesh_controls = vectorDuplicate(x);
    CHKERR m_field.getInterface<VecManager>()->setLocalGhostVector(
        simple->getProblemName(), COL, mesh_controls, INSERT_VALUES,
        SCATTER_FORWARD);
    PetscReal mesh_norm;
    CHKERR VecNorm(mesh_controls, NORM_INFINITY, &mesh_norm);
    if (mesh_norm != 0)
      SETERRQ(PETSC_COMM_SELF, MOFEM_DATA_INCONSISTENCY,
              "TAO evaluation hooks implicitly scattered trial controls");
 
    CHKERR TaoDestroy(&tao);
  }
  CATCH_ERRORS;
 
  CHKERR MoFEM::Core::Finalize();
  return 0;
}