PetscRunner.hpp

#pragma once

#include <petscpc.h>
#include <petscksp.h>

#include <dune/common/unused.hh>

#include <amdis/Initfile.hpp>
#include <amdis/linearalgebra/RunnerInterface.hpp>
#include <amdis/linearalgebra/SolverInfo.hpp>
#include <amdis/linearalgebra/petsc/Interface.hpp>

#ifndef KSPDGGMRES
#define KSPDGGMRES "dggmres"
#endif


namespace AMDiS
{

  template <class Matrix, class Vector>
  class PetscRunner
      : public RunnerInterface<Matrix,Vector>
  {
    using M = typename Matrix::BaseMatrix;
    using X = typename Vector::BaseVector;
    using Y = typename Vector::BaseVector;

  public:

    explicit PetscRunner(std::string const& prefix)
      : prefix_(prefix)
    {
      Parameters::get(prefix + "->info", info_);
    }

    ~PetscRunner()
    {
      exit();
    }

    void init(M const& A) override
    {
      exit();
#if HAVE_MPI
      KSPCreate(Environment::comm(), &ksp_);
#else
      KSPCreate(PETSC_COMM_SELF, &ksp_);
#endif

      PETSc::KSPSetOperators(ksp_, A, A);
      initKSP(ksp_, prefix_);
      initialized_ = true;
    }

    void exit() override
    {
      if (initialized_)
        KSPDestroy(&ksp_);
      initialized_ = false;
    }

    int solve([[maybe_unused]] M const& A, X& x, Y const& b, SolverInfo& solverInfo) override
    {
      assert(initialized_);

      KSPSolve(ksp_, b, x);

      if (info_ > 0)
        KSPReasonView(ksp_, PETSC_VIEWER_STDOUT_WORLD);

      KSPConvergedReason reason;
      KSPGetConvergedReason(ksp_, &reason);
      return reason > 0 ? 0 : int(reason);
    }

    KSP ksp() { return ksp_; }

  protected:
    // initialize the KSP solver from the initfile
    virtual void initKSP(KSP ksp, std::string prefix) const
    {
      // see https://www.mcs.anl.gov/petsc/petsc-current/docs/manualpages/KSP/KSPType.html
      auto kspType = Parameters::get<std::string>(prefix + "->ksp");
      std::string kspTypeStr = kspType.value_or("default");

      if (!kspType)
        Parameters::get(prefix, kspTypeStr);

      if (kspTypeStr == "direct")
        initDirectSolver(ksp, prefix);
      else if (kspTypeStr != "default") {
        KSPSetType(ksp, kspTypeStr.c_str());

        // initialize some KSP specific parameters
        initKSPParameters(ksp, kspTypeStr.c_str(), prefix);

        // set initial guess to nonzero only for non-preonly ksp type
        if (kspTypeStr != "preonly")
          KSPSetInitialGuessNonzero(ksp, PETSC_TRUE);
      }

      // set a KSPMonitor if info > 0
      int info = 0;
      Parameters::get(prefix + "->info", info);
      if (info == 1) {
        KSPMonitorSet(ksp, PETSc::KSPMonitorCyclic, nullptr, nullptr);
      } else if (info > 1) {
        KSPMonitorSet(ksp, PETSc::KSPMonitorNoisy, nullptr, nullptr);
      }

      // set solver tolerances
      auto maxIt = Parameters::get<PetscInt>(prefix + "->max it");
      auto rTol = Parameters::get<PetscReal>(prefix + "->rtol");
      auto aTol = Parameters::get<PetscReal>(prefix + "->atol");
      auto dTol = Parameters::get<PetscReal>(prefix + "->divtol");
      if (maxIt || rTol || aTol || dTol)
        KSPSetTolerances(ksp,
          rTol.value_or(PETSC_DEFAULT), aTol.value_or(PETSC_DEFAULT),
          dTol.value_or(PETSC_DEFAULT), maxIt.value_or(PETSC_DEFAULT));

      auto kspPrefixParam = Parameters::get<std::string>(prefix + "->prefix");
      if (kspPrefixParam) {
        std::string kspPrefix = kspPrefixParam.value() + "_";
        KSPSetOptionsPrefix(ksp, kspPrefix.c_str());
        KSPSetFromOptions(ksp);
      }

      PC pc;
      KSPGetPC(ksp, &pc);
      initPC(pc, prefix + "->pc");

      // show details of ksp object
      if (info >= 10) {
        KSPView(ksp, PETSC_VIEWER_STDOUT_WORLD);
      }
    }


    // initialize a direct solver from the initfile
    virtual void initDirectSolver(KSP ksp, std::string prefix) const
    {
      KSPSetInitialGuessNonzero(ksp, PETSC_TRUE);
      KSPSetType(ksp, KSPRICHARDSON);

      PC pc;
      KSPGetPC(ksp, &pc);
      PCSetType(pc, PCLU);
      auto matSolverType = Parameters::get<std::string>(prefix + "->mat solver");
      if (matSolverType)
        PETSc::PCFactorSetMatSolverType(pc, matSolverType.value().c_str());
      else {
        Mat Amat, Pmat;
        KSPGetOperators(ksp, &Amat, &Pmat);
        MatType type;
        MatGetType(Pmat, &type);
        if (std::strcmp(type, MATSEQAIJ) == 0)
          PETSc::PCFactorSetMatSolverType(pc, MATSOLVERUMFPACK);
        else if (std::strcmp(type, MATAIJ) == 0)
          PETSc::PCFactorSetMatSolverType(pc, MATSOLVERMUMPS);
      }
    }


    // initialize the preconditioner pc from the initfile
    virtual void initPC(PC pc, std::string prefix) const
    {
      // see https://www.mcs.anl.gov/petsc/petsc-current/docs/manualpages/PC/PCType.html
      auto pcType = Parameters::get<std::string>(prefix);
      std::string pcTypeStr = "default";
      if (pcType) {
        pcTypeStr = pcType.value();
        PCSetType(pc, pcTypeStr.c_str());
      }

      if (pcTypeStr == "lu") {
        // configure matsolverpackage
        // see https://www.mcs.anl.gov/petsc/petsc-current/docs/manualpages/Mat/MatSolverType.html
        auto matSolverType = Parameters::get<std::string>(prefix + "->mat solver");
        if (matSolverType)
          PETSc::PCFactorSetMatSolverType(pc, matSolverType.value().c_str());
      } else if (pcTypeStr == "bjacobi") {
        // configure sub solver
        PetscInt n_local, first;
        KSP* ksps;
        PCSetUp(pc);
        PCBJacobiGetSubKSP(pc, &n_local, &first, &ksps);
        for (PetscInt i = 0; i < n_local; ++i)
          initKSP(ksps[i], prefix + "->sub ksp");
      } else if (pcTypeStr == "ksp") {
        // configure sub ksp type
        KSP ksp;
        PCSetUp(pc);
        PCKSPGetKSP(pc, &ksp);
        initKSP(ksp, prefix + "->ksp");
      }

      auto pcPrefixParam = Parameters::get<std::string>(prefix + "->prefix");
      if (pcPrefixParam) {
        std::string pcPrefix = pcPrefixParam.value() + "_";
        PCSetOptionsPrefix(pc, pcPrefix.c_str());
        PCSetFromOptions(pc);
      }
    }


    // provide initfile parameters for some PETSc KSP parameters
    virtual void initKSPParameters(KSP ksp, char const* ksptype, std::string prefix) const
    {
      // parameters for the Richardson solver
      if (std::strcmp(ksptype, KSPRICHARDSON) == 0)
      {
        auto scale = Parameters::get<PetscReal>(prefix + "->scale");
        if (scale)
          KSPRichardsonSetScale(ksp, scale.value());
      }
      // parameters for the gmres solver
      else if (std::strcmp(ksptype, KSPGMRES) == 0 ||
               std::strcmp(ksptype, KSPFGMRES) == 0 ||
               std::strcmp(ksptype, KSPLGMRES) == 0 ||
               std::strcmp(ksptype, KSPDGGMRES) == 0 ||
               std::strcmp(ksptype, KSPPGMRES) == 0)
      {
        auto restart = Parameters::get<PetscInt>(prefix + "->restart");
        if (restart)
          KSPGMRESSetRestart(ksp, restart.value());

        auto gramschmidt = Parameters::get<std::string>(prefix + "->gramschmidt");
        if (gramschmidt) {
          if (gramschmidt.value() == "classical") {
            KSPGMRESSetOrthogonalization(ksp, KSPGMRESClassicalGramSchmidtOrthogonalization);
            auto refinementType = Parameters::get<std::string>(prefix + "->refinement type");
            if (refinementType) {
              if (refinementType.value() == "never")
                KSPGMRESSetCGSRefinementType(ksp, KSP_GMRES_CGS_REFINE_NEVER);
              else if (refinementType.value() == "ifneeded")
                KSPGMRESSetCGSRefinementType(ksp, KSP_GMRES_CGS_REFINE_IFNEEDED);
              else if (refinementType.value() == "always")
                KSPGMRESSetCGSRefinementType(ksp, KSP_GMRES_CGS_REFINE_ALWAYS);
            }
          } else if (gramschmidt.value() == "modified") {
            KSPGMRESSetOrthogonalization(ksp, KSPGMRESModifiedGramSchmidtOrthogonalization);
          }
        }
      }
      // parameters for the BiCGStab_ell solver
      else if (std::strcmp(ksptype, KSPBCGSL) == 0)
      {
        auto ell = Parameters::get<PetscInt>(prefix + "->ell");
        if (ell)
           KSPBCGSLSetEll(ksp, ell.value());
      }
      // parameters for the GCR solver
      else if (std::strcmp(ksptype, KSPGCR) == 0)
      {
        auto restart = Parameters::get<PetscInt>(prefix + "->restart");
        if (restart)
           KSPGCRSetRestart(ksp, restart.value());
      }
    }

  protected:
    std::string prefix_;
    int info_ = 0;

    KSP ksp_;
    bool initialized_ = false;
  };

} // end namespace AMDiS