KrylovMethods.hpp

/*
 * This file is part of bogus, a C++ sparse block matrix library.
 *
 * Copyright 2013 Gilles Daviet <gdaviet@gmail.com>
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
*/

#ifndef BOGUS_KRYLOV_METHODS_HPP
#define BOGUS_KRYLOV_METHODS_HPP

#include "../BlockSolvers.fwd.hpp"

#include "../Utils/LinearSolverBase.hpp"
#include "../Utils/Signal.hpp"
#include "../Utils/NumTraits.hpp"
#include "../Utils/CppTools.hpp"

#define BOGUS_KRYLOV_METHODS \
    BOGUS_PROCESS_KRYLOV_METHOD(CG      )\
    BOGUS_PROCESS_KRYLOV_METHOD(BiCG    )\
    BOGUS_PROCESS_KRYLOV_METHOD(BiCGSTAB)\
    BOGUS_PROCESS_KRYLOV_METHOD(CGS     )\
    BOGUS_PROCESS_KRYLOV_METHOD(GMRES   )\
    BOGUS_PROCESS_KRYLOV_METHOD(TFQMR   )


namespace bogus {

namespace krylov {

template< template < typename, typename, typename > class Method,
          typename Matrix, typename Preconditioner, typename Traits >
struct KrylovSolverBase
        : public LinearSolverBase< Method< Matrix, Preconditioner, Traits > >
{
    typedef Method< Matrix, Preconditioner, Traits > Derived ;
    typedef LinearSolverBase< Derived > Base ;
    typedef typename Traits::Scalar Scalar ;
    typedef Signal< unsigned, Scalar > SignalType ;

    const Matrix *m_A ;
    const Preconditioner *m_P;
    const SignalType *m_callback ;

    Scalar m_tol ;
    unsigned m_maxIters;

    KrylovSolverBase( const Matrix &A,
                      unsigned maxIters, Scalar tol,
                      const Preconditioner *P,
                      const SignalType *callback
            )
        : m_A( &A ), m_P( P ), m_callback( callback ),
          m_tol( tol ), m_maxIters( maxIters ),
          m_parallelizeRhs( false ), m_enableResCaching( false )
    {}

    KrylovSolverBase( )
        : m_A( BOGUS_NULL_PTR(const Matrix) ),
          m_P( BOGUS_NULL_PTR(const Preconditioner) ),
          m_callback( BOGUS_NULL_PTR(const SignalType) ),
          m_tol( 0 ), m_maxIters( 0 ),
          m_parallelizeRhs( false ), m_enableResCaching( false )
    {}

    template < typename RhsT >
    typename LinearSolverTraits< Derived >::template Result< RhsT >::Type
    solve( const RhsT& rhs ) const
    {
        typedef typename LinearSolverTraits< Derived >::template Result< RhsT >::Type ReturnType ;
        static ReturnType s_cachedRes ;

        if( !m_enableResCaching ) {
            ReturnType x( m_A->rows(), rhs.cols() ) ;
            x.setZero() ;
            Base::solve( rhs, x ) ;
            return x ;
        }

        if( s_cachedRes.rows() != m_A->rows() || s_cachedRes.cols() != rhs.cols() ) {
            s_cachedRes.resize( m_A->rows(), rhs.cols() ) ;
            s_cachedRes.setZero() ;
        }

        Base::solve( rhs, s_cachedRes ) ;
        return s_cachedRes ;
    }

    template < typename ResT, typename RhsT >
    Scalar solve( const RhsT& rhs, ResT& x ) const
    {
        Scalar res = 0;
#ifndef BOGUS_DONT_PARALLELIZE
#pragma omp parallel for reduction( +:res ) if ( m_parallelizeRhs )
#endif
        for( std::ptrdiff_t c = 0 ; c < (std::ptrdiff_t) rhs.cols() ; ++c )
        {
            res += Base::derived().vectorSolve( rhs.col( c ), x.col( c ) ) ;
        }

        return res ;
    }

    Derived& parallelizeRhs( bool parallelize = true )
    {
        m_parallelizeRhs = parallelize ;
        return Base::derived() ;
    }


    Derived& enableResCaching( bool doCache = true )
    {
        m_enableResCaching = doCache ;
        return Base::derived() ;
    }

protected:
    bool m_parallelizeRhs ;
    bool m_enableResCaching ;

} ;

namespace solvers {

// Useful macros for common declarations

#define BOGUS_MAKE_KRYLOV_SOLVER_TYPEDEFS( MethodName ) \
    typedef KrylovSolverBase< solvers::MethodName, Matrix, Preconditioner, Traits > Base ; \
    typedef typename Traits::Scalar Scalar ;            \
                                                        \
    using Base::m_A ;                                   \
    using Base::m_P ;                                   \
    using Base::m_maxIters ;                            \
    using Base::m_tol ;                                 \
    using Base::m_callback ;                            \

#define BOGUS_MAKE_KRYLOV_SOLVER_HEADER( MethodName )   \
    BOGUS_MAKE_KRYLOV_SOLVER_TYPEDEFS( MethodName )     \
    MethodName( const Matrix &A,                        \
    unsigned maxIters,                                  \
    Scalar tol = NumTraits< Scalar >::epsilon(),        \
    const Preconditioner *P = BOGUS_NULL_PTR(const Preconditioner), \
    const typename Base::SignalType *callback = BOGUS_NULL_PTR(const typename Base::SignalType ) )\
    : Base( A, maxIters, tol, P, callback )             \
        {}                                              \
                                                        \
    MethodName():Base() {}                              \
                                                        \
    template < typename RhsT, typename ResT >           \
    Scalar vectorSolve( const RhsT &b, ResT x ) const ;     \


// Conjugate Gradient


template < typename Matrix,
           typename Preconditioner = TrivialPreconditioner< Matrix >,
           typename Traits = ProblemTraits< typename MatrixTraits<Matrix>::Scalar > >
struct CG : public KrylovSolverBase< CG, Matrix, Preconditioner, Traits >
{


    BOGUS_MAKE_KRYLOV_SOLVER_HEADER( CG )

} ;


template < typename Matrix,
           typename Preconditioner = TrivialPreconditioner< Matrix >,
           typename Traits = ProblemTraits< typename MatrixTraits<Matrix>::Scalar > >
struct BiCG : public KrylovSolverBase< BiCG, Matrix, Preconditioner, Traits>
{

    BOGUS_MAKE_KRYLOV_SOLVER_HEADER( BiCG )

} ;


template < typename Matrix,
           typename Preconditioner = TrivialPreconditioner< Matrix >,
           typename Traits = ProblemTraits< typename MatrixTraits<Matrix>::Scalar > >
struct BiCGSTAB : public KrylovSolverBase< BiCGSTAB, Matrix, Preconditioner, Traits>
{

    BOGUS_MAKE_KRYLOV_SOLVER_HEADER( BiCGSTAB )

} ;


template < typename Matrix,
           typename Preconditioner = TrivialPreconditioner< Matrix >,
           typename Traits = ProblemTraits< typename MatrixTraits<Matrix>::Scalar > >
struct CGS : public KrylovSolverBase< CGS, Matrix, Preconditioner, Traits>
{

    BOGUS_MAKE_KRYLOV_SOLVER_HEADER( CGS )
} ;



template < typename Matrix,
           typename Preconditioner = TrivialPreconditioner< Matrix >,
           typename Traits = ProblemTraits< typename MatrixTraits<Matrix>::Scalar > >
struct GMRES : public KrylovSolverBase< GMRES, Matrix, Preconditioner, Traits>
{
    BOGUS_MAKE_KRYLOV_SOLVER_TYPEDEFS( GMRES )

    GMRES() : Base(), m_restart( 0 )
    {}

    GMRES( const Matrix &A,
           unsigned maxIters,
           Scalar tol = NumTraits< Scalar >::epsilon(),
           const Preconditioner *P = BOGUS_NULL_PTR( const Preconditioner),
           const typename Base::SignalType *callback = BOGUS_NULL_PTR(const typename Base::SignalType),
           unsigned restart = 0 )
        : Base( A, maxIters, tol, P, callback ),
          m_restart( restart )
    {}

    GMRES &setRestart( unsigned restart )
    {
        m_restart = restart ;
        return *this ;
    }

    template < typename RhsT, typename ResT >
    Scalar vectorSolve( const RhsT &b, ResT x ) const ;

protected:
    unsigned m_restart ;
} ;


template < typename Matrix,
           typename Preconditioner = TrivialPreconditioner< Matrix >,
           typename Traits = ProblemTraits< typename MatrixTraits<Matrix>::Scalar > >
struct TFQMR : public KrylovSolverBase< TFQMR, Matrix, Preconditioner, Traits>
{
    BOGUS_MAKE_KRYLOV_SOLVER_HEADER( TFQMR )
};


} //namespace solvers

} //namespace krylov


// Sovlers traits ( in bogus namespace )

#define BOGUS_PROCESS_KRYLOV_METHOD( MethodName ) \
    template< typename Matrix, typename Preconditioner, class Traits > \
    struct LinearSolverTraits< krylov::solvers::MethodName< Matrix, Preconditioner, Traits > > \
    { \
      typedef Matrix MatrixType ; \
      template < typename RhsT > struct Result { \
          typedef typename Traits::template MutableClone< RhsT >::Type Type ; \
      } ; \
    } ; \
    template< typename Matrix, typename Preconditioner, class Traits, \
              typename RhsBlockT, bool TransposeLhs, bool TransposeRhs > \
    struct BlockBlockProductTraits <  krylov::solvers::MethodName< Matrix, Preconditioner, Traits >, RhsBlockT, TransposeLhs, TransposeRhs > \
    { \
        typedef typename BlockBlockProductTraits < Matrix, RhsBlockT, TransposeLhs, TransposeRhs >::ReturnType \
        ReturnType ; \
    } ; \

BOGUS_KRYLOV_METHODS
#undef BOGUS_PROCESS_KRYLOV_METHOD

} //namespace bogus

#endif