TaoCtx.cpp

Go to the documentation of this file.

/**
 * @brief Context for Tao optimization routines
 * @author Anonymous author(s) committing under MIT license
 */
 
 
namespace MoFEM {
 
MoFEMErrorCode TaoCtx::copyLoops(const TaoCtx &tao_ctx) {
  MoFEMFunctionBeginHot;
  loopsHessian = tao_ctx.loopsHessian;
  loopsObjective = tao_ctx.loopsObjective;
  loopsGradient = tao_ctx.loopsGradient;
  preHessian = tao_ctx.preHessian;
  postHessian = tao_ctx.postHessian;
  preObjective = tao_ctx.preObjective;
  postObjective = tao_ctx.postObjective;
  preGradient = tao_ctx.preGradient;
  postGradient = tao_ctx.postGradient;
  MoFEMFunctionReturnHot(0);
}
 
MoFEMErrorCode TaoCtx::clearLoops() {
  MoFEMFunctionBeginHot;
  loopsHessian.clear();
  loopsObjective.clear();
  loopsGradient.clear();
  preHessian.clear();
  postHessian.clear();
  preObjective.clear();
  postObjective.clear();
  preGradient.clear();
  postGradient.clear();
  MoFEMFunctionReturnHot(0);
}
 
PetscErrorCode TaoSetObjective(Tao tao, Vec x, PetscReal *f, void *ctx) {
  TaoCtx *tao_ctx = (TaoCtx *)ctx;
  MoFEMFunctionBegin;
  PetscLogEventBegin(tao_ctx->MOFEM_EVENT_TaoObjective, 0, 0, 0, 0);
  CHKERR VecGhostUpdateBegin(x, INSERT_VALUES, SCATTER_FORWARD);
  CHKERR VecGhostUpdateEnd(x, INSERT_VALUES, SCATTER_FORWARD);
  CHKERR tao_ctx->mField.getInterface<VecManager>()->setLocalGhostVector(
      tao_ctx->problemName, COL, x, INSERT_VALUES, SCATTER_REVERSE);
 
  auto cache_ptr = boost::make_shared<CacheTuple>();
  CHKERR tao_ctx->mField.cache_problem_entities(tao_ctx->problemName,
                                                cache_ptr);
 
  auto set = [&](auto &fe) {
    fe.tao = tao;
    fe.tao_x = x;
    fe.tao_ctx = TaoMethod::CTX_TAO_OBJECTIVE;
    fe.data_ctx = PetscData::CtxSetX;
    fe.cacheWeakPtr = cache_ptr;
  };
 
  auto unset = [&](auto &fe) {
    fe.tao_ctx = TaoMethod::CTX_TAO_NONE;
    fe.data_ctx = PetscData::CtxSetNone;
  };
 
  for (auto &bit : tao_ctx->preObjective) {
    bit->vecAssembleSwitch = boost::move(tao_ctx->vecAssembleSwitch);
    set(*bit);
    CHKERR tao_ctx->mField.problem_basic_method_preProcess(tao_ctx->problemName,
                                                           *bit);
    unset(*bit);
  }
 
  for (auto &lit : tao_ctx->loopsObjective) {
    set(*lit.second);
    CHKERR tao_ctx->mField.loop_finite_elements(tao_ctx->problemName, lit.first,
                                                *lit.second, nullptr, MF_EXIST,
                                                cache_ptr);
    unset(*lit.second);
  }
 
  for (auto &bit : tao_ctx->postObjective) {
    bit->vecAssembleSwitch = boost::move(tao_ctx->vecAssembleSwitch);
    set(*bit);
    CHKERR tao_ctx->mField.problem_basic_method_postProcess(
        tao_ctx->problemName, *bit);
    unset(*bit);
    tao_ctx->vecAssembleSwitch = boost::move(bit->vecAssembleSwitch);
  }
 
  PetscLogEventEnd(tao_ctx->MOFEM_EVENT_TaoObjective, 0, 0, 0, 0);
  MoFEMFunctionReturn(0);
}
 
PetscErrorCode TaoSetGradient(Tao tao, Vec x, Vec f, void *ctx) {
  TaoCtx *tao_ctx = (TaoCtx *)ctx;
  MoFEMFunctionBegin;
  PetscLogEventBegin(tao_ctx->MOFEM_EVENT_TaoGradient, 0, 0, 0, 0);
  CHKERR VecGhostUpdateBegin(x, INSERT_VALUES, SCATTER_FORWARD);
  CHKERR VecGhostUpdateEnd(x, INSERT_VALUES, SCATTER_FORWARD);
  CHKERR tao_ctx->mField.getInterface<VecManager>()->setLocalGhostVector(
      tao_ctx->problemName, COL, x, INSERT_VALUES, SCATTER_REVERSE);
 
  auto zero_ghost_vec = [](Vec g) {
    MoFEMFunctionBegin;
    Vec l;
    CHKERR VecGhostGetLocalForm(g, &l);
    double *a;
    CHKERR VecGetArray(l, &a);
    int s;
    CHKERR VecGetLocalSize(l, &s);
    for (int i = 0; i != s; ++i)
      a[i] = 0;
    CHKERR VecRestoreArray(l, &a);
    CHKERR VecGhostRestoreLocalForm(g, &l);
    MoFEMFunctionReturn(0);
  };
  CHKERR zero_ghost_vec(f);
 
  tao_ctx->vecAssembleSwitch = boost::movelib::make_unique<bool>(true);
  auto cache_ptr = boost::make_shared<CacheTuple>();
  CHKERR tao_ctx->mField.cache_problem_entities(tao_ctx->problemName,
                                                cache_ptr);
 
  auto set = [&](auto &fe) {
    fe.tao = tao;
    fe.tao_x = x;
    fe.tao_f = f;
    fe.snes_ctx = SnesMethod::CTX_SNESSETFUNCTION;
    fe.ksp_ctx = KspMethod::CTX_SETFUNCTION;
    fe.tao_ctx = TaoMethod::CTX_TAO_GRADIENT;
    fe.data_ctx = PetscData::CtxSetF | PetscData::CtxSetX;
    fe.cacheWeakPtr = cache_ptr;
  };
 
  auto unset = [&](auto &fe) {
    fe.snes_ctx = SnesMethod::CTX_SNESNONE;
    fe.ksp_ctx = KspMethod::CTX_KSPNONE;
    fe.tao_ctx = TaoMethod::CTX_TAO_GRADIENT;
    fe.data_ctx = PetscData::CtxSetNone;
  };
 
  for (auto &bit : tao_ctx->preGradient) {
    bit->vecAssembleSwitch = boost::move(tao_ctx->vecAssembleSwitch);
    set(*bit);
    CHKERR tao_ctx->mField.problem_basic_method_preProcess(tao_ctx->problemName,
                                                           *bit);
    unset(*bit);
    tao_ctx->vecAssembleSwitch = boost::move(bit->vecAssembleSwitch);
  }
 
  for (auto &lit : tao_ctx->loopsGradient) {
    lit.second->vecAssembleSwitch = boost::move(tao_ctx->vecAssembleSwitch);
    set(*lit.second);
    CHKERR tao_ctx->mField.loop_finite_elements(tao_ctx->problemName, lit.first,
                                                *lit.second, nullptr,
                                                MF_EXIST, cache_ptr);
    unset(*lit.second);
    tao_ctx->vecAssembleSwitch = boost::move(lit.second->vecAssembleSwitch);
  }
 
  for (auto &bit : tao_ctx->postGradient) {
    bit->vecAssembleSwitch = boost::move(tao_ctx->vecAssembleSwitch);
    set(*bit);
    CHKERR tao_ctx->mField.problem_basic_method_postProcess(
        tao_ctx->problemName, *bit);
    unset(*bit);
    tao_ctx->vecAssembleSwitch = boost::move(bit->vecAssembleSwitch);
  }
 
  if (*(tao_ctx->vecAssembleSwitch)) {
    CHKERR VecGhostUpdateBegin(f, ADD_VALUES, SCATTER_REVERSE);
    CHKERR VecGhostUpdateEnd(f, ADD_VALUES, SCATTER_REVERSE);
    CHKERR VecAssemblyBegin(f);
    CHKERR VecAssemblyEnd(f);
  }
  PetscLogEventEnd(tao_ctx->MOFEM_EVENT_TaoGradient, 0, 0, 0, 0);
  MoFEMFunctionReturn(0);
}
 
PetscErrorCode TaoSetObjectiveAndGradient(Tao tao, Vec x, PetscReal *f, Vec g,
                                          void *ctx) {
  MoFEMFunctionBegin;
  CHKERR TaoSetObjective(tao, x, f, ctx);
  CHKERR TaoSetGradient(tao, x, g, ctx);
  MoFEMFunctionReturn(0);
}
 
PetscErrorCode TaoSetHessian(Tao tao, Vec x, Mat H, Mat Hpre, void *ctx) {
  TaoCtx *tao_ctx = (TaoCtx *)ctx;
  MoFEMFunctionBegin;
  PetscLogEventBegin(tao_ctx->MOFEM_EVENT_TaoHessian, 0, 0, 0, 0);
 
  tao_ctx->matAssembleSwitch = boost::movelib::make_unique<bool>(true);
  auto cache_ptr = boost::make_shared<CacheTuple>();
  CHKERR tao_ctx->mField.cache_problem_entities(tao_ctx->problemName,
                                                 cache_ptr);
 
  auto set = [&](auto &fe) {
    fe.tao = tao;
    fe.tao_A = H;
    fe.tao_B = Hpre;
    fe.ksp_ctx = KspMethod::CTX_OPERATORS;
    fe.snes_ctx = SnesMethod::CTX_SNESSETJACOBIAN;
    fe.tao_ctx = TaoMethod::CTX_TAO_HESSIAN;
    fe.data_ctx = PetscData::CtxSetA | PetscData::CtxSetB | PetscData::CtxSetX;
    fe.cacheWeakPtr = cache_ptr;
  };
 
  auto unset = [&](auto &fe) {
    fe.snes_ctx = SnesMethod::CTX_SNESNONE;
    fe.ksp_ctx = KspMethod::CTX_KSPNONE;
    fe.tao_ctx = TaoMethod::CTX_TAO_NONE;
    fe.data_ctx = PetscData::CtxSetNone;
  };
 
 
  CHKERR VecGhostUpdateBegin(x, INSERT_VALUES, SCATTER_FORWARD);
  CHKERR VecGhostUpdateEnd(x, INSERT_VALUES, SCATTER_FORWARD);
  CHKERR tao_ctx->mField.getInterface<VecManager>()->setLocalGhostVector(
      tao_ctx->problemName, COL, x, INSERT_VALUES, SCATTER_REVERSE);
 
  for (auto &bit : tao_ctx->preHessian) {
    bit->matAssembleSwitch = boost::move(tao_ctx->matAssembleSwitch);
    set(*bit);
    CHKERR tao_ctx->mField.problem_basic_method_preProcess(
        tao_ctx->problemName, *bit);
    unset(*bit);
    tao_ctx->matAssembleSwitch = boost::move(bit->matAssembleSwitch);
  }
 
  for (auto &lit : tao_ctx->loopsHessian) {
    lit.second->matAssembleSwitch = boost::move(tao_ctx->matAssembleSwitch);
    set(*lit.second);
    CHKERR tao_ctx->mField.loop_finite_elements(
        tao_ctx->problemName, lit.first, *(lit.second), nullptr, MF_EXIST,
        cache_ptr);
    unset(*lit.second);
    tao_ctx->matAssembleSwitch = boost::move(lit.second->matAssembleSwitch);
  }
 
  for (auto &bit : tao_ctx->postHessian) {
    bit->matAssembleSwitch = boost::move(tao_ctx->matAssembleSwitch);
    set(*bit);
    CHKERR tao_ctx->mField.problem_basic_method_postProcess(
        tao_ctx->problemName, *bit);
    unset(*bit);
    tao_ctx->matAssembleSwitch = boost::move(bit->matAssembleSwitch);
  }
 
  if (*tao_ctx->matAssembleSwitch) {
    CHKERR MatAssemblyBegin(Hpre, MAT_FINAL_ASSEMBLY);
    CHKERR MatAssemblyEnd(Hpre, MAT_FINAL_ASSEMBLY);
  }
  PetscLogEventEnd(tao_ctx->MOFEM_EVENT_TaoHessian, 0, 0, 0, 0);
  MoFEMFunctionReturn(0);  
}
 
} // namespace MoFEM