Homec4science
Diffusion tamaas (archive/legacy_mpi) Browse

tamaas/src/solvers/petsc/tao.cpp
b0c44a420f99archive/legacy_mpi

Last Change
Raw File

tao.cpp
View Options

/*
* SPDX-License-Indentifier: AGPL-3.0-or-later
*
* Copyright (©) 2016-2024 EPFL (École Polytechnique Fédérale de Lausanne),
* Laboratory (LSMS - Laboratoire de Simulation en Mécanique des Solides)
* Copyright (©) 2020-2024 Lucas Frérot
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published
* by the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <https://www.gnu.org/licenses/>.
*
*/
/* -------------------------------------------------------------------------- */
#include "tao.hh"
#include "model_type.hh"
#include <iomanip>
/* -------------------------------------------------------------------------- */
namespace tamaas {
namespace petsc {
#define chk(error) PetscCallAbort(mpi::comm::world(), error)
#define PRECONDITION 0
PetscErrorCode monitor(Tao tao, void* ctx) {
Int iter = 0;
Real f = 0;
chk(TaoGetSolutionStatus(tao, &iter, &f, nullptr, nullptr, nullptr, nullptr));
auto& solver = *static_cast<OptimizationSolver*>(ctx);
solver.printState(iter, f, solver.computeError());
return 0;
}
PetscErrorCode convergence(Tao tao, void* ctx) {
auto& solver = *static_cast<OptimizationSolver*>(ctx);
Int iter = 0;
chk(TaoGetIterationNumber(tao, &iter));
auto error = solver.computeError();
auto tolerance = solver.getTolerance();
if (error < tolerance)
chk(TaoSetConvergedReason(tao, TAO_CONVERGED_USER));
else if (iter > static_cast<Int>(solver.getMaxIterations()))
chk(TaoSetConvergedReason(tao, TAO_DIVERGED_MAXITS));
else
chk(TaoSetConvergedReason(tao, TAO_CONTINUE_ITERATING));
return 0;
}
PetscErrorCode form_gradient(Tao, Vec x, PetscReal* f, Vec g, void* ctx) {
PetscInt local_n = 0;
chk(VecGetLocalSize(x, &local_n));
Real *xarray = nullptr, *garray = nullptr;
chk(VecGetArray(x, &xarray));
chk(VecGetArray(g, &garray));
UInt N = static_cast<UInt>(local_n);
auto& grad_ctx = *static_cast<OptimizationSolver::gradient_context*>(ctx);
auto xgrid = viewGrid<true>(std::get<1>(grad_ctx), std::get<0>(grad_ctx), 1,
span<Real>{xarray, N});
auto ggrid = viewGrid<true>(std::get<1>(grad_ctx), std::get<0>(grad_ctx), 1,
span<Real>{garray, N});
auto* functional = std::get<2>(grad_ctx);
functional->computeGradF(*xgrid, *ggrid);
*f = functional->computeF(*xgrid, *ggrid);
return 0;
}
PetscErrorCode form_hessian(Tao, Vec /*x*/, Mat /*H*/, Mat /*Hpre*/,
void* /*ctx*/) {
// empty on purpose, the matrix shell does all the work of forming the
// hessian
return 0;
}
PetscErrorCode hessian_shell(Mat mat, Vec x, Vec y) {
IntegralOperator* integral_op = nullptr;
chk(MatShellGetContext(mat, &integral_op));
auto& traction = integral_op->getModel().getTraction();
auto local_n = traction.dataSize();
auto nc = traction.getNbComponents();
Real *xarray = nullptr, *yarray = nullptr;
chk(VecGetArray(x, &xarray));
chk(VecGetArray(y, &yarray));
auto type = integral_op->getModel().getType();
auto sizes = integral_op->getModel().getBoundaryDiscretization();
auto xgrid = viewGrid<true>(type, sizes, 1, span<Real>{xarray, nc * local_n});
auto ygrid = viewGrid<true>(type, sizes, 1, span<Real>{yarray, nc * local_n});
// Apply hessian
integral_op->apply(*xgrid, *ygrid);
// Restore petsc vectors
chk(VecRestoreArray(x, &xarray));
chk(VecRestoreArray(y, &yarray));
return 0;
}
PetscErrorCode apply_preconditionner(PC pc, Vec x, Vec y) {
Tao* tao = nullptr;
Mat P;
Real pmin = 0;
chk(PCShellGetContext(pc, &tao));
chk(TaoGetHessian(*tao, nullptr, &P, nullptr, nullptr));
chk(MatMult(P, x, y));
chk(VecMin(y, nullptr, &pmin));
chk(VecShift(y, -pmin));
return 0;
}
OptimizationSolver::OptimizationSolver(Model& model,
const GridBase<Real>& surface,
Real tolerance, std::string petsc_args)
: PolonskyKeerRey(model, surface, tolerance, PolonskyKeerRey::pressure,
PolonskyKeerRey::gap) {
TAMAAS_ASSERT(model.getType() == model_type::basic_1d or
model.getType() == model_type::basic_2d,
"TAO solver only compatible with basic model types");
// Create the options database and make petsc use it
chk(PetscOptionsCreate(&tao_options));
chk(PetscOptionsInsertString(tao_options, petsc_args.c_str()));
chk(PetscOptionsPush(tao_options));
// Wrap Petsc vectors on Tamaas' grids
auto local_n = primal->dataSize(), global_n = primal->globalDataSize();
chk(VecCreateMPIWithArray(mpi::comm::world(), 1, local_n, global_n,
primal->getInternalData(), &_primal_vec));
chk(VecCreateMPIWithArray(mpi::comm::world(), 1, local_n, global_n,
dual->getInternalData(), &_dual_vec));
// Creating bounds vectors
chk(VecCreateMPI(mpi::comm::world(), local_n, global_n, &_lower_b));
chk(VecCreateMPI(mpi::comm::world(), local_n, global_n, &_upper_b));
// Create matrix for hessian
chk(MatCreateShell(mpi::comm::world(), local_n, local_n, global_n, global_n,
this->integral_op.get(), &H));
chk(MatShellSetOperation(H, MATOP_MULT, (void (*)(void))hessian_shell));
#if PRECONDITION
// Create matrix for preconditioner
model.getBEEngine().registerDirichlet();
chk(MatCreateShell(mpi::comm::world(), local_n, local_n, global_n, global_n,
model.getIntegralOperator("westergaard_dirichlet").get(),
&P));
chk(MatShellSetOperation(P, MATOP_MULT, (void (*)(void))hessian_shell));
#endif
// Create gradient context
grad_ctx = std::make_tuple(model.getBoundaryDiscretization(), model.getType(),
&this->functional);
// Create TAO solver context
chk(TaoCreate(mpi::comm::world(), &tao));
chk(TaoSetType(tao, TAOGPCG)); // projected conjugate gradient
chk(TaoSetSolution(tao, _primal_vec));
chk(TaoSetObjectiveAndGradient(tao, _dual_vec, form_gradient, &grad_ctx));
chk(TaoSetHessian(tao, H, H, form_hessian, nullptr));
chk(TaoSetFromOptions(tao));
// Set monitor and convergence
chk(TaoSetMonitor(tao, monitor, this, nullptr));
chk(TaoSetConvergenceTest(tao, convergence, this));
// chk(TaoSetTolerances(tao, this->tolerance, 0, 0));
// Set defaults for linear solver and preconditioner
KSP ksp;
PC pc;
chk(TaoGetKSP(tao, &ksp));
chk(KSPGetPC(ksp, &pc));
#if PRECONDITION
chk(TaoSetHessian(tao, H, P, form_hessian, nullptr));
chk(PCSetType(pc, PCSHELL));
chk(PCShellSetContext(pc, &tao));
chk(PCShellSetApply(pc, apply_preconditionner));
#else
chk(PCSetType(pc, PCNONE));
#endif
// Real atol = 1e-10;
UInt maxiter = 50;
chk(KSPSetTolerances(ksp, this->tolerance, PETSC_DEFAULT, PETSC_DEFAULT,
maxiter));
// Setting bounds
chk(VecZeroEntries(_lower_b));
chk(VecSet(_upper_b, PETSC_INFINITY));
chk(TaoSetVariableBounds(tao, _lower_b, _upper_b));
// Clearing solver-specific options
chk(PetscOptionsPop());
}
OptimizationSolver::~OptimizationSolver() {
chk(PetscOptionsDestroy(&tao_options));
chk(TaoDestroy(&tao));
chk(MatDestroy(&H));
#if PRECONDITION
chk(MatDestroy(&P));
#endif
for (auto* v : {&_dual_vec, &_primal_vec, &_lower_b, &_upper_b})
chk(VecDestroy(v));
}
/* -------------------------------------------------------------------------- */
Real OptimizationSolver::solve(std::vector<Real> /*target_gap*/) {
// Printing column headers
Logger().get(LogLevel::info) << std::setw(5) << "Iter"
<< " " << std::setw(15) << "Cost_f"
<< " " << std::setw(15) << "Error" << '\n'
<< std::fixed;
// Petsc solve
chk(TaoSolve(tao));
// Final update of dual variable
functional.computeGradF(*primal, *dual);
auto error = computeError();
enforceAdmissibleState();
Int iter = 0;
Real f = 0;
chk(TaoGetSolutionStatus(tao, &iter, &f, nullptr, nullptr, nullptr, nullptr));
logIteration(iter, f, error);
return error;
}
} // namespace petsc
} // namespace tamaas
c4science · Help