ConjugateGradient2.h

Go to the documentation of this file.

/*  Copyright (C) 2003-2021 Free Electron Organization
    Any use of this software requires a license.  If a valid license
    was not distributed with this file, visit freeelectron.org. */

/** @file */

#ifndef __solve_ConjugateGradient2_h__
#define __solve_ConjugateGradient2_h__

#define CG_NORMALIZE        TRUE
#define CG_DEBUG            TRUE
#define CG_TRACE            TRUE

#define CG_CHECK_SYMMETRY   (FE_CODEGEN<=FE_DEBUG)
#define CG_CHECK_POSITIVE   (FE_CODEGEN<=FE_DEBUG)
#define CG_CHECK_PEAK       (FE_CODEGEN<=FE_DEBUG)
#define CG_VERIFY           (FE_CODEGEN<=FE_DEBUG)

namespace fe
{
namespace ext
{

/**************************************************************************//**
    @brief solve Ax=b for x

    @ingroup solve

    Uses Conjugate-Gradient.  The matrix must be positive-definite
    and symmetric.

    The arguments are templated, so any argument types should work,
    given that they have the appropriate methods and operators.

    TODO try seeding with previous x instead of clearing to zero.
*//***************************************************************************/
template <typename MATRIX, typename VECTOR>
class ConjugateGradient2
{
    public:
                /// @brief Choice to alter matrix during solve
        enum    Preconditioning
                {
                    e_none,
                    e_diagonal_A
                };

                ConjugateGradient2(void):
                        m_threshold(1e-6f),
                        m_preconditioning(e_none)                           {}

        void    solve(VECTOR& x, const MATRIX& A, const VECTOR& b);

        void    setThreshold(F64 threshold) { m_threshold=threshold; }

        void    setPreconditioning(Preconditioning preconditioning)
                {   m_preconditioning=preconditioning; }

    private:
        VECTOR          r,r_1,r_2;          //* residual    (at k, k-1, k-2)
        VECTOR          p,p_1;              //* direction
        VECTOR          temp;               //* persistent temporary
        VECTOR          Ap;
        VECTOR          y;
        VECTOR          m_AT_b;
        VECTOR          m_preconditionedVector;
        MATRIX          m_tempMatrix;
        MATRIX          m_AT;
        MATRIX          m_AT_A;
        MATRIX          m_preconditionedMatrix;
        F64             m_threshold;
        Preconditioning m_preconditioning;
};

template <typename MATRIX, typename VECTOR>
inline void ConjugateGradient2<MATRIX,VECTOR>::solve(VECTOR& x,
        const MATRIX& A, const VECTOR& b)
{
    U32 N=size(b);
    if(size(y)!=N)
    {
        y=b;        // adopt size
    }
    if(size(Ap)!=N)
    {
        Ap=b;       // adopt size
    }
    set(y);

    F64 dot_r_1;

#if CG_CHECK_SYMMETRY
    for(U32 iy=0;iy<N;iy++)
    {
        for(U32 ix=iy;ix<N;ix++)
        {
            if(fabs(A(ix,iy)-A(iy,ix))>1e-9f)
            {
                feLog("ConjugateGradient2 symmetry error %d,%d %.6G %.6G\n",
                        ix,iy,A(ix,iy),A(iy,ix));
            }
        }
    }
#endif

#if CG_DEBUG
    feLog("\nA\n%s\nb=<%s>\n",print(A).c_str(),print(b).c_str());
#endif

    if(magnitudeSquared(b)<m_threshold)
    {
#if CG_DEBUG
        feLog("ConjugateGradient2::solve has trivial solution\n");
#endif
        x=y;
        return;
    }

    const MATRIX* pA=&A;
    const VECTOR* pb=&b;
    if(CG_NORMALIZE)
    {
        m_AT.setTranspose(A);
        m_AT_A.setProduct(m_AT,A);
        pA=&m_AT_A;

        m_AT_b=m_AT*b;
        pb=&m_AT_b;

#if CG_DEBUG
        feLog("AT\n%s\n",print(m_AT).c_str());
        feLog("AT*A\n%s\n",print(m_AT_A).c_str());
        feLog("AT*b=<%s>\n",print(m_AT_b).c_str());
#endif
        // HACK sloppy cut&paste
        if(m_preconditioning==e_diagonal_A)
        {
            premultiplyInverseDiagonal(m_tempMatrix,m_AT_A,m_AT_A);
            postmultiplyInverseDiagonal(m_preconditionedMatrix,
                    m_tempMatrix,m_AT_A);
            pA=&m_preconditionedMatrix;

            premultiplyInverseDiagonal(m_preconditionedVector,m_AT_A,m_AT_b);
            pb=&m_preconditionedVector;

#if CG_DEBUG
            feLog("mid\n%s\n",print(m_tempMatrix).c_str());
            feLog("A'\n%s\nb`=<%s>\n",print(*pA).c_str(),print(*pb).c_str());
#endif
        }
    }
    else if(m_preconditioning==e_diagonal_A)
    {
        premultiplyInverseDiagonal(m_tempMatrix,A,A);
        postmultiplyInverseDiagonal(m_preconditionedMatrix,
                m_tempMatrix,A);
        pA=&m_preconditionedMatrix;

        premultiplyInverseDiagonal(m_preconditionedVector,A,b);
        pb=&m_preconditionedVector;

#if CG_DEBUG
        feLog("mid\n%s\n",print(m_tempMatrix).c_str());
        feLog("A'\n%s\nb`=<%s>\n",print(*pA).c_str(),print(*pb).c_str());
#endif
    }
#if CG_CHECK_POSITIVE
    for(U32 k=0;k<N;k++)
    {
        if((*pA)(k,k)<=0.0)
        {
            feLog("ConjugateGradient2::solve"
                    " non-positive diagonal %d,%d %.6G\n",
                    k,k,(*pA)(k,k));
        }
    }
#endif
#if CG_CHECK_POSITIVE
    F64 peak=0.0;
    U32 peak_i;
    U32 peak_j;
    for(U32 i=0;i<N;i++)
    {
        for(U32 j=0;j<N;j++)
        {
            F64 value=fabs((*pA)(i,j));
            if(peak<value)
            {
                peak=value;
                peak_i=i;
                peak_j=j;
            }
        }
    }
    if(peak_i!=peak_j)
    {
        feLog("ConjugateGradient2::solve non-diagonal peak %d,%d %.6G\n",
                peak_i,peak_j,peak);
    }
#endif

    for(U32 k=1;k<=N;k++)
    {
        if(k==1)
        {
            p= *pb;
            r_1=p;

            dot_r_1=dot(r_1,r_1);
        }
        else
        {
            r_2=r_1;
            r_1=r;
            p_1=p;

            dot_r_1=dot(r_1,r_1);

            F64 beta=dot_r_1/dot(r_2,r_2);

//          p=r_1+beta*p_1;
            temp=p_1;
            temp*=beta;
            p=r_1;
            p+=temp;
        }

        transformVector(*pA,p,Ap);
        F64 alpha=dot_r_1/dot(p,Ap);

#if CG_TRACE==FALSE
        if(magnitudeSquared(p)==0.0f)
#endif
        {
            feLog("\n%d alpha=%.6G r_1*r_1 %.6G y<%s>\n",
                    k,alpha,dot_r_1,print(y).c_str());
            feLog("r_1<%s>\nr_2<%s>\n",
                    print(r_1).c_str(),print(r_2).c_str());
            feLog("p<%s>\np_1<%s>\nA*p<%s>\n",
                    print(p).c_str(),print(p_1).c_str(),print(*pA*p).c_str());
        }

        if(magnitudeSquared(p)==0.0f)
        {
            feX("ConjugateGradient2::solve","direction lost its magnitude");
        }

//      y+=alpha*p;
        temp=p;
        temp*=alpha;
        y+=temp;

//      r=r_1-alpha*(Ap);
        temp=Ap;
        temp*=alpha;
        r=r_1;
        r-=temp;

#if CG_TRACE
        feLog("r<%s>\n",print(r).c_str());
#endif

        feLog("ConjugateGradient2::solve ran %d/%d alpha %.6G r %.6G p %.6G\n",
                k,N,alpha,magnitude(r),magnitude(p));

//      if(magnitudeSquared(r)<m_threshold)
        if(magnitudeSquared(r)<1e-4)
        {
#if CG_DEBUG
#endif
            feLog("ConjugateGradient2::solve early solve %d/%d\n",k,N);
            break;
        }

        if(k==N)
        {
            feLog("ConjugateGradient2::solve ran %d/%d\n",k,N);
        }
    }

    if(m_preconditioning==e_diagonal_A)
    {
        premultiplyInverseDiagonal(x,A,y);
    }
    else
    {
        // TODO use x all along to save copy
        x=y;
    }

#if CG_DEBUG
    feLog("\ny=<%s>\nA'*y=<%s>\n",print(y).c_str(),print(*pA*y).c_str());
    feLog("\nx=<%s>\nA*x=<%s>\n",print(x).c_str(),print(A*x).c_str());
#endif

#if CG_VERIFY
    BWORD invalid=FALSE;
    for(U32 k=0;k<N;k++)
    {
        if(FE_INVALID_SCALAR(x[k]))
        {
            invalid=TRUE;
        }
    }
    VECTOR Ax=A*x;
    F64 distance=magnitude(Ax-b);
    if(invalid || distance>1.0f)
    {
        feLog("ConjugateGradient2::solve failed to converge (dist=%.6G)\n",
                distance);
        if(size(x)<100)
        {
            feLog("  collecting state ...\n");
            feLog("A=\n%s\nx=<%s>\nA*x=<%s>\nb=<%s>\n",
                    print(A).c_str(),print(x).c_str(),
                    print(Ax).c_str(),print(b).c_str());
        }
//      feX("ConjugateGradient2::solve","failed to converge");
    }
#endif
}

} /* namespace ext */
} /* namespace fe */

#endif /* __solve_ConjugateGradient2_h__ */