AMReX_GMRES.H

Go to the documentation of this file.

#ifndef AMREX_GMRES_H_
#define AMREX_GMRES_H_
#include <AMReX_Config.H>
 
#include <AMReX_BLProfiler.H>
#include <AMReX_Print.H>
#include <AMReX_TableData.H>
#include <AMReX_Vector.H>
#include <cmath>
#include <limits>
#include <memory>
 
namespace amrex {


template <typename V, typename M>
class GMRES
{
public:
 
    using RT = typename M::RT; // double or float

    GMRES ();

    void define (M& 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 setFlexible (bool flag) { m_flexible = flag; }

    void setRestartLength (int rl);

    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; }

    [[nodiscard]] RT getInitialResidualNorm () const { return m_rnorm0; }
 
private:
    void clear ();
    void allocate_scratch ();
    void cycle (V& a_xx, int& a_status, int& a_itcount, RT& a_rnorm0, int a_maxiter);
    void build_solution (V& a_xx, int it);
    void compute_residual (V& a_rr, V const& a_xx, V const& a_bb);
 
    [[nodiscard]] bool converged (RT r0, RT r) const;
 
    void gram_schmidt_orthogonalization (int it);
    void update_hessenberg (int it, bool happyend, RT& res);
 
    int m_verbose = 0;
    int m_maxiter = 2000;
    int m_its = 0;
    int m_status = -1;
    int m_restrtlen = 30;
    bool m_nonzero_guess = false;
    bool m_flexible = false;
    RT m_res = std::numeric_limits<RT>::max();
    RT m_rnorm0 = std::numeric_limits<RT>::max();
    RT m_rtol = RT(0);
    RT m_atol = RT(0);
    Vector<RT> m_hh_1d;
    Vector<RT> m_hes_1d;
    Table2D<RT> m_hh;
    Table2D<RT> m_hes;
    Vector<RT> m_grs;
    Vector<RT> m_cc;
    Vector<RT> m_ss;
    std::unique_ptr<V> m_v_tmp_rhs;
    std::unique_ptr<V> m_v_tmp_lhs;
    Vector<V> m_vv;
    Vector<V> m_zz; // preconditioned vectors, FGMRES only
    M* m_linop = nullptr;
};
 
template <typename V, typename M>
GMRES<V,M>::GMRES ()
{
    allocate_scratch();
}
 
template <typename V, typename M>
void GMRES<V,M>::allocate_scratch ()
{
    int rs = m_restrtlen;
 
    m_hh_1d.resize(std::size_t(rs + 2) * (rs + 1));
    m_hh = Table2D<RT>(m_hh_1d.data(), {0,0}, {rs+1,rs}); // (0:rs+1,0:rs)
 
    m_hes_1d.resize(std::size_t(rs + 2) * (rs + 1));
    m_hes = Table2D<RT>(m_hes_1d.data(), {0,0}, {rs+1,rs}); // (0:rs+1,0:rs)
 
    m_grs.resize(rs + 2);
    m_cc.resize(rs + 1);
    m_ss.resize(rs + 1);
}
 
template <typename V, typename M>
void GMRES<V,M>::setRestartLength (int rl)
{
    if (m_restrtlen != rl) {
        m_restrtlen = rl;
        allocate_scratch();
        m_vv.clear();
        m_zz.clear();
    }
}
 
template <typename V, typename M>
void GMRES<V,M>::define (M& linop)
{
    clear();
    m_linop = &linop;
}
 
template <typename V, typename M>
void GMRES<V,M>::clear ()
{
    m_its = 0;
    m_status = -1;
    m_res = std::numeric_limits<RT>::max();
    m_v_tmp_rhs.reset();
    m_v_tmp_lhs.reset();
    m_vv.clear();
    m_zz.clear();
    m_linop = nullptr;
}
 
template <typename V, typename M>
bool GMRES<V,M>::converged (RT r0, RT r) const
{
    return (r < r0*m_rtol) || (r < m_atol);
}
 
template <typename V, typename M>
void GMRES<V,M>::solve (V& a_sol, V const& a_rhs, RT a_tol_rel, RT a_tol_abs, int a_its)
{
    BL_PROFILE("GMRES::solve()");
 
    AMREX_ALWAYS_ASSERT(m_linop != nullptr);
 
    auto t0 = amrex::second();
 
    if (m_v_tmp_rhs == nullptr) {
        m_v_tmp_rhs = std::make_unique<V>(m_linop->makeVecRHS());
    }
    if (m_v_tmp_lhs == nullptr) {
        m_v_tmp_lhs = std::make_unique<V>(m_linop->makeVecLHS());
    }
    if (m_vv.empty()) {
        m_vv.reserve(m_restrtlen+1);
        for (int i = 0; i < 2; ++i) { // to save space, start with just 2
            m_vv.emplace_back(m_linop->makeVecRHS());
        }
    }
 
    m_rtol = a_tol_rel;
    m_atol = a_tol_abs;
 
    if (a_its < 0) { a_its = m_maxiter; }
 
    auto rnorm0 = RT(0);
 
    if (m_nonzero_guess) {
        compute_residual(m_vv[0], a_sol, a_rhs);
    } else {
        m_linop->assign(m_vv[0], a_rhs);
        m_linop->setToZero(a_sol);
    }
 
    m_its = 0;
    m_status = -1;
    cycle(a_sol, m_status, m_its, rnorm0, a_its);
 
    while (m_status == -1 && m_its < a_its) {
        compute_residual(m_vv[0], a_sol, a_rhs);
        cycle(a_sol, m_status, m_its, rnorm0, a_its);
    }
 
    if (m_status == -1 && m_its >= a_its) { m_status = 1; }
    m_rnorm0 = rnorm0;
 
    m_v_tmp_rhs.reset();
    m_v_tmp_lhs.reset();
    m_vv.clear();
    m_zz.clear();
 
    auto t1 = amrex::second();
    if (m_verbose > 0) {
        amrex::Print() << "GMRES: Solve Time = " << t1-t0 << '\n';
    }
}
 
template <typename V, typename M>
void GMRES<V,M>::cycle (V& a_xx, int& a_status, int& a_itcount, RT& a_rnorm0, int a_maxiter)
{
    BL_PROFILE("GMRES::cycle()");
 
    m_res = m_linop->norm2(m_vv[0]);
    m_grs[0] = m_res;
 
    if (m_res == RT(0.0)) {
        a_status = 0;
        return;
    }
 
    m_linop->scale(m_vv[0], RT(1.0)/m_res);
 
    if (a_itcount == 0) { a_rnorm0 = m_res; }
 
    a_status = converged(a_rnorm0,m_res) ? 0 : -1;
 
    int it = 0;
    while (it < m_restrtlen && a_itcount < a_maxiter)
    {
        if (m_verbose > 1) {
            amrex::Print() << "GMRES: iter = " << a_itcount
                           << ", residual = " << m_res << ", " << m_res/a_rnorm0
                           << " (rel.)\n";
        }
 
        if (a_status == 0) { break; }
 
        while (m_vv.size() < it+2) {
            m_vv.emplace_back(m_linop->makeVecRHS());
        }
 
        if (m_flexible) {
            if (m_zz.empty()) { m_zz.reserve(m_restrtlen); }
            while (m_zz.size() < it+1) {
                m_zz.emplace_back(m_linop->makeVecLHS());
            }
        }
 
        auto const& vv_it  = m_vv[it  ];
        auto      & vv_it1 = m_vv[it+1];
        auto      & zz_it  = m_flexible ? m_zz[it] : *m_v_tmp_lhs;
 
        m_linop->precond(zz_it, vv_it);
        m_linop->apply(vv_it1, zz_it);
 
        gram_schmidt_orthogonalization(it);
 
        auto tt = m_linop->norm2(vv_it1);
 
        auto const sml = RT((sizeof(RT) == 8) ? 1.e-99 : 1.e-30);
        bool happyend = (tt < sml);
        if (!happyend) {
            m_linop->scale(vv_it1, RT(1.0)/tt);
        }
 
        m_hh (it+1,it) = tt;
        m_hes(it+1,it) = tt;
 
        update_hessenberg(it, happyend, m_res);
 
        ++it;
        ++a_itcount;
        a_status = converged(a_rnorm0, m_res) ? 0 : -1;
        if (happyend) { break; }
    }
 
    if ((m_verbose > 1) && (a_status != 0 || a_itcount >= a_maxiter)) {
        amrex::Print() << "GMRES: iter = " << a_itcount
                       << ", residual = " << m_res << ", " << m_res/a_rnorm0
                       << " (rel.)\n";
    }
 
    build_solution(a_xx, it-1);
}
 
template <typename V, typename M>
void GMRES<V,M>::gram_schmidt_orthogonalization (int const it)
{
    // Two unmodified Gram-Schmidt Orthogonalization
 
    BL_PROFILE("GMRES::GramSchmidt");
 
    auto& vv_1 = m_vv[it+1];
 
    Vector<RT> lhh(it+1);
 
    for (int j = 0; j <= it; ++j) {
        m_hh (j,it) = RT(0.0);
        m_hes(j,it) = RT(0.0);
    }
 
    for (int ncnt = 0; ncnt < 2 ; ++ncnt)
    {
        for (int j = 0; j <= it; ++j) {
            lhh[j] = m_linop->dotProduct(vv_1, m_vv[j]);
        }
 
        for (int j = 0; j <= it; ++j) {
            m_linop->increment(vv_1, m_vv[j], -lhh[j]);
            m_hh (j,it) += lhh[j];
            m_hes(j,it) -= lhh[j];
        }
    }
}
 
template <typename V, typename M>
void GMRES<V,M>::update_hessenberg (int it, bool happyend, RT& res)
{
    BL_PROFILE("GMRES::update_hessenberg()");
 
    for (int j = 1; j <= it; ++j) {
        auto tt = m_hh(j-1,it);
        m_hh(j-1,it) = m_cc[j-1] * tt         + m_ss[j-1] * m_hh(j,it);
        m_hh(j  ,it) = m_cc[j-1] * m_hh(j,it) - m_ss[j-1] * tt;
    }
 
    if (!happyend)
    {
        auto tt = std::sqrt(m_hh(it,it)*m_hh(it,it) + m_hh(it+1,it)*m_hh(it+1,it));
        m_cc[it] = m_hh(it  ,it) / tt;
        m_ss[it] = m_hh(it+1,it) / tt;
        m_grs[it+1] = - (m_ss[it] * m_grs[it]);
        m_grs[it  ] =    m_cc[it] * m_grs[it];
        m_hh(it,it) = m_cc[it] * m_hh(it,it) + m_ss[it] * m_hh(it+1,it);
        res = std::abs(m_grs[it+1]);
    }
    else
    {
        res = RT(0.0);
    }
}
 
template <typename V, typename M>
void GMRES<V,M>::build_solution (V& a_xx, int const it)
{
    BL_PROFILE("GMRES:build_solution()");
 
    if (it < 0) { return; }
 
    if (m_hh(it,it) != RT(0.0)) {
        m_grs[it] /= m_hh(it,it);
    } else {
        m_grs[it] = RT(0.0);
    }
 
    for (int ii = 1; ii <= it; ++ii) {
        int k = it - ii;
        auto tt = m_grs[k];
        for (int j = k+1; j <= it; ++j) {
            tt -= m_hh(k,j) * m_grs[j];
        }
        m_grs[k] = tt / m_hh(k,k);
    }
 
    if (m_flexible) {
        // x += Z y, with Z the stored preconditioned vectors
        for (int ii = 0; ii < it+1; ++ii) {
            m_linop->increment(a_xx, m_zz[ii], m_grs[ii]);
        }
    } else {
        // x += P^{-1} V y, valid only if P^{-1} is the same in all iterations
        m_linop->setToZero(*m_v_tmp_rhs);
        for (int ii = 0; ii < it+1; ++ii) {
            m_linop->increment(*m_v_tmp_rhs, m_vv[ii], m_grs[ii]);
        }
 
        m_linop->precond(*m_v_tmp_lhs, *m_v_tmp_rhs);
        m_linop->increment(a_xx, *m_v_tmp_lhs, RT(1.0));
    }
}
 
template <typename V, typename M>
void GMRES<V,M>::compute_residual (V& a_rr, V const& a_xx, V const& a_bb)
{
    BL_PROFILE("GMRES::compute_residual()");
    m_linop->assign(*m_v_tmp_lhs, a_xx);
    m_linop->apply(*m_v_tmp_rhs, *m_v_tmp_lhs);
    m_linop->linComb(a_rr, RT(1.0), a_bb, RT(-1.0), *m_v_tmp_rhs);
}
 
}
#endif