pcg.h

Go to the documentation of this file.

#ifndef _DG_PCG_
#define _DG_PCG_
 
#include <cmath>
 
#include "blas.h"
#include "functors.h"
#include "extrapolation.h"
#include "backend/typedefs.h"
 
#include "backend/timer.h"
 
namespace dg{
namespace detail
{
struct NORM
{
    template<class T, class Z>
    DG_DEVICE
    auto operator()( T w, Z z) {
        return w*norm(z);} // returns floating point
};
struct DOT
{
    template<class Z0, class T, class Z1>
    DG_DEVICE
    auto operator()( Z0 z0, T w, Z1 z1) { // returns floating point
        return w*(z0.real()*z1.real() + z0.imag()*z1.imag());
    }
};
} // namespace detail
 
// TODO check for better stopping criteria using condition number estimates?
enum ComplexMode{
    complex_symmetric, 
    complex_hermitian 
};
 
// TODO test complex behaviour
 
template<class ContainerType, ComplexMode complex_mode = dg::complex_hermitian>
class PCG
{
  public:
    using container_type = ContainerType;
    using value_type = get_value_type<ContainerType>; 
 
    PCG() = default;
    PCG( const ContainerType& copyable, unsigned max_iterations):
        m_r(copyable), m_p(m_r), m_ap(m_r), m_max_iter(max_iterations){}
    void set_max( unsigned new_max) {m_max_iter = new_max;}
    unsigned get_max() const {return m_max_iter;}
    const ContainerType& copyable()const{ return m_r;}
 
    void set_verbose( bool verbose){ m_verbose = verbose;}
    void set_throw_on_fail( bool throw_on_fail){
        m_throw_on_fail = throw_on_fail;
    }
 
    template<class ...Params>
    void construct( Params&& ...ps)
    {
        //construct and swap
        *this = PCG( std::forward<Params>( ps)...);
    }
    template< class MatrixType0, class ContainerType0, class ContainerType1, class MatrixType1, class ContainerType2 >
    unsigned solve( MatrixType0&& A, ContainerType0& x, const ContainerType1& b, MatrixType1&& P, const ContainerType2& W, double eps = 1e-12, double nrmb_correction = 1, int test_frequency = 1);
  private:
    // !! Calling this "norm" would shadow std::norm in NORM
    template<class ContainerType1, class ContainerType2>
    auto nrm( const ContainerType1& w, const ContainerType2& x)
    {
        using value_type_x = dg::get_value_type<ContainerType2>;
        constexpr bool is_complex = dg::is_scalar_v<value_type_x, dg::ComplexTag>;
        if constexpr (is_complex)
            return sqrt( blas1::vdot( detail::NORM(), w, x));
        else
            return sqrt( blas2::dot( w, x));
    }
    template<class ContainerType0, class ContainerType1, class ContainerType2>
    auto dot( const ContainerType0& x0, const ContainerType1& w, const ContainerType2& x1)
    {
        using value_type_x1 = dg::get_value_type<ContainerType2>;
        constexpr bool is_complex = dg::is_scalar_v<value_type_x1, dg::ComplexTag>;
        if constexpr (not is_complex or complex_mode == dg::complex_symmetric)
            return blas2::dot( x0, w, x1); // this returns value_type
        else // For Hermitian matrices all dot products in CG are real
            return blas1::vdot( detail::DOT(), x0, w, x1); // this returns floating point
    }
    ContainerType m_r, m_p, m_ap;
    unsigned m_max_iter;
    bool m_verbose = false, m_throw_on_fail = true, m_cocg = false;
};
 
template<class ContainerType>
using COCG = dg::PCG<ContainerType, dg::complex_symmetric>;
 
template< class ContainerType, ComplexMode mode>
template< class Matrix, class ContainerType0, class ContainerType1, class Preconditioner, class ContainerType2>
unsigned PCG< ContainerType, mode>::solve( Matrix&& A, ContainerType0& x, const ContainerType1& b, Preconditioner&& P, const ContainerType2& W, double eps, double nrmb_correction, int save_on_dots )
{
    // self-adjoint: apply PCG algorithm to (P 1/W) (W A) x = (P 1/W) (W b) : P' A' x = P' b'
    // This effectively just replaces all scalar products with the weighted one
    auto nrmb = nrm( W, b), nrmr(nrmb);
    double tol = eps*(nrmb + nrmb_correction);
#ifdef MPI_VERSION
    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
#endif //MPI
    if( m_verbose)
    {
        DG_RANK0 std::cout << "# Norm of W b "<<nrmb <<"\n";
        DG_RANK0 std::cout << "# Residual errors: \n";
    }
    if( nrmb == 0)
    {
        blas1::copy( value_type(0), x);
        return 0;
    }
    blas2::symv( std::forward<Matrix>(A), x, m_r);
    blas1::axpby( 1., b, -1., m_r);
    nrmr = nrm( W, m_r);
 
    if( nrmr < tol) //if x happens to be the solution
        return 0;
    blas2::symv( std::forward<Preconditioner>(P), m_r, m_p );
    auto dotzr_old = dot( m_p, W, m_r), alpha(dotzr_old), dotzr_new(dotzr_old);
    for( unsigned i=1; i<m_max_iter; i++)
    {
        blas2::symv( std::forward<Matrix>(A), m_p, m_ap);
        alpha =  dotzr_old/dot( m_p, W, m_ap);
        blas1::axpby( alpha, m_p, 1., x);
        blas1::axpby( -alpha, m_ap, 1., m_r);
        if( 0 == i%save_on_dots )
        {
            nrmr = nrm( W, m_r);
            if( m_verbose)
            {
                DG_RANK0 std::cout << "# Absolute r*W*r "<<nrmr <<"\t ";
                DG_RANK0 std::cout << "#  < Critical "<<tol <<"\t ";
                DG_RANK0 std::cout << "# (Relative "<<nrmr/nrmb << ")\n";
            }
            if( nrmr < tol)
                return i;
        }
        blas2::symv(std::forward<Preconditioner>(P),m_r,m_ap);
        dotzr_new = dot( m_ap, W, m_r);
        blas1::axpby(1., m_ap, dotzr_new/dotzr_old, m_p );
        dotzr_old=dotzr_new;
    }
    if( m_throw_on_fail)
    {
        throw dg::Fail( tol, Message(_ping_)
            <<"After "<<m_max_iter<<" PCG iterations with rtol "<<eps<<" and atol "<<eps*nrmb_correction );
    }
    return m_max_iter;
}
 
} //namespace dg
 
 
 
#endif //_DG_PCG_