AMReX_PCG.H

Go to the documentation of this file.

#ifndef AMREX_PCG_H_
#define AMREX_PCG_H_
#include <AMReX_Config.H>
 
#include <AMReX_BLProfiler.H>
#include <AMReX_Math.H>
#include <AMReX_Print.H>
 
#include <cmath>
#include <iomanip>
#include <limits>
 
namespace amrex {
 
template <typename V, typename M>
class PCG
{
public:
 
    using RT = typename M::RT; // double or float
 
    PCG () = default;
 
    void define (M& linop) { m_linop = &linop; }
 
    void solve (V& a_sol, V const& a_rhs, RT a_tol_rel, RT a_tol_abs, int a_its=-1);
 
    void setInitialGuessNonzero (bool flag) { m_nonzero_guess = flag; }
 
    void setVerbose (int v) { m_verbose = v; }
 
    void setMaxIters (int niters) { m_maxiter = niters; }
 
    [[nodiscard]] int getNumIters () const { return m_its; }
 
    [[nodiscard]] int getStatus () const { return m_status; }
 
    [[nodiscard]] RT getResidualNorm () const { return m_res; }
 
private:
 
    int m_verbose = 0;
    bool m_nonzero_guess = false;
    int m_maxiter = 2000;
    int m_its = 0;
    int m_status = -1;
    RT m_res = std::numeric_limits<RT>::max();
    RT m_rtol = RT(0);
    RT m_atol = RT(0);
 
    M* m_linop = nullptr;
};
 
template <typename V, typename M>
void PCG<V,M>::solve (V& a_sol, V const& a_rhs, RT a_tol_rel, RT a_tol_abs, int a_its)
{
    BL_PROFILE("PCG::solve()");
 
    AMREX_ALWAYS_ASSERT(m_linop != nullptr);
 
    m_rtol = a_tol_rel;
    m_atol = a_tol_abs;
    if (a_its < 0) { a_its = m_maxiter; }
 
    m_its = 0;
    m_status = -1;
 
    V r = m_linop->makeVecRHS();
    V z = m_linop->makeVecLHS();
    V p = m_linop->makeVecLHS();
    V q = m_linop->makeVecRHS();
 
    if (m_nonzero_guess) {
        // r = b - A x. apply() needs an x made by makeVecLHS.
        m_linop->assign(p, a_sol);
        m_linop->apply(q, p);
        m_linop->linComb(r, RT(1), a_rhs, RT(-1), q);
    } else {
        m_linop->setToZero(a_sol);
        m_linop->assign(r, a_rhs);
    }
 
    RT rnorm = m_linop->norm2(r);
    RT const rnorm0 = rnorm;
    m_res = rnorm;
 
    if (m_verbose > 0) {
        amrex::Print() << "PCG: Initial residual (error0) = " << rnorm0 << '\n';
    }
 
    if (rnorm0 == RT(0) || rnorm0 < m_atol) {
        m_status = 0;
        return;
    }
 
    auto converged = [&] (RT rn) {
        return rn == RT(0) || rn < m_rtol*rnorm0 || rn < m_atol;
    };
 
    m_linop->precond(z, r);
    m_linop->assign(p, z);
    RT rz = m_linop->dotProduct(r, z);
    // A and the preconditioner may both be negative definite: every r.z and
    // p.Ap must then keep the sign of the first r.z.
    RT const sgn = (rz < RT(0)) ? RT(-1) : RT(1);
    if (rz == RT(0) || !amrex::isfinite(rz)) { m_status = 3; }
 
    for (int iter = 1; iter <= a_its && m_status == -1; ++iter)
    {
        m_linop->apply(q, p);
        RT const pq = m_linop->dotProduct(p, q);
        if (sgn*pq <= RT(0) || !amrex::isfinite(pq)) { m_status = 2; break; }
        RT const alpha = rz/pq;
 
        m_linop->increment(a_sol, p, alpha);
        m_linop->increment(r, q, -alpha);
 
        rnorm = m_linop->norm2(r);
        m_its = iter;
        m_res = rnorm;
 
        if (m_verbose > 1) {
            amrex::Print() << "PCG: Iteration " << std::setw(11) << iter
                           << " rel. err. " << rnorm/rnorm0 << '\n';
        }
 
        if (converged(rnorm)) { m_status = 0; break; }
        if (!amrex::isfinite(rnorm)) { m_status = 3; break; }
        if (iter == a_its) { break; }
 
        m_linop->precond(z, r);
        RT const rz_new = m_linop->dotProduct(r, z);
        if (sgn*rz_new <= RT(0) || !amrex::isfinite(rz_new)) { m_status = 3; break; }
        RT const beta = rz_new/rz;
        rz = rz_new;
 
        m_linop->scale(p, beta);
        m_linop->increment(p, z, RT(1));         // p = z + beta*p
    }
 
    if (m_status == -1) { m_status = 1; }
 
    if (m_verbose > 0) {
        amrex::Print() << "PCG: Final: Iteration " << std::setw(4) << m_its
                       << " rel. err. " << rnorm/rnorm0 << '\n';
    }
 
    if (m_status != 0 && m_verbose > 0) {
        amrex::Print() << "PCG: Failed to converge, status = " << m_status << '\n';
    }
}
 
}
 
#endif