ConjugateGradient.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_ConjugateGradient_h__
#define __solve_ConjugateGradient_h__

#define CG_DEBUG            FALSE
#define CG_TRACE            FALSE

#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)
#define CG_EXCEPTION        (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 ConjugateGradient
{
    public:
                ConjugateGradient(void):
                        m_threshold(1e-6f)                                  {}

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

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

    private:
        VECTOR          r;          //* residual
        VECTOR          d;          //* direction
        VECTOR          temp;       //* persistent temporary
        VECTOR          q;          //* A*d
        F64             m_threshold;
};

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

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

#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("ConjugateGradient symmetry error %d,%d %.6G %.6G\n",
                        ix,iy,A(ix,iy),A(iy,ix));
                if(CG_EXCEPTION)
                {
                    feX("ConjugateGradient::solve","unsymmetrical");
                }
            }
        }
    }
#endif

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

    const MATRIX* pA=&A;
    const VECTOR* pb=&b;
#if CG_CHECK_POSITIVE
    for(U32 k=0;k<N;k++)
    {
        if((*pA)(k,k)<=0.0)
        {
            feLog("ConjugateGradient::solve"
                    " non-positive diagonal %d,%d %.6G\n",
                    k,k,(*pA)(k,k));
        }
    }
#endif
#if CG_CHECK_PEAK
    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("ConjugateGradient::solve non-diagonal peak %d,%d %.6G\n",
                peak_i,peak_j,peak);
    }
#endif

    U32 i=0;
    r= *pb;
    d=r;
    F64 dnew=dot(r,r);
//  F64 d0=dnew;
    while(i<N)
    {
        transformVector(*pA,d,q);
#if CG_TRACE
        feLog("\n%d q=A*d                           >>> %.6G <<<\n",i,dnew);
        feLog("d<%s>\n",c_print(d));
        feLog("q<%s>\n",c_print(q));
#endif

        F64 alpha=dnew/dot(d,q);
#if CG_TRACE
        feLog("alpha=dnew/dot(d,q) %.6G=%.6G/%.6G\n",alpha,dnew,dot(d,q));
#endif

#if CG_TRACE
        feLog("x<%s>\n",c_print(x));
#endif
//      x=x+alpha*d;

//      temp=d;
//      temp*=alpha;
//      x+=temp;

        addScaled(x,alpha,d);
#if CG_TRACE
        feLog("x+=alpha*d <%s>\n",c_print(temp));
        feLog("x<%s>\n",c_print(x));
#endif

#if CG_TRACE
        feLog("r<%s>\n",c_print(r));
#endif
//      r=r-alpha*q;

//      temp=q;
//      temp*=alpha;
//      r-=temp;

        addScaled(r,-alpha,q);
#if CG_TRACE
        feLog("r-=alpha*q <%s>\n",c_print(temp));
        feLog("r<%s>\n",c_print(r));
#endif

        F64 dold=dnew;
        dnew=dot(r,r);
        F64 beta=dnew/dold;
#if CG_TRACE
        feLog("dnew=%.6G dold=%.6G beta=%.6G\n",dnew,dold,beta);
#endif

//      d=r+beta*d;

//      temp=d;
//      temp*=beta;
//      d=r;
//      d+=temp;

        scaleAndAdd(d,beta,r);

#if CG_TRACE
        feLog("d=r+beta*d <%s>\n",c_print(temp));
        feLog("d<%s>\n",c_print(d));
#endif

        if(magnitudeSquared(d)==0.0f)
        {
//          feX("ConjugateGradient::solve","direction lost its magnitude");

#if CG_DEBUG
            feLog("ConjugateGradient::solve direction lost its magnitude\n");
#endif
            break;
        }

        i++;

#if CG_TRACE
        feLog("ConjugateGradient::solve"
                " ran %d/%d alpha %.6G |r| %.6G |d| %.6G\n",
                i,N,alpha,magnitude(r),magnitude(d));
#endif

        if(magnitudeSquared(r)<m_threshold)
        {
#if CG_DEBUG
            feLog("ConjugateGradient::solve early solve %d/%d\n",i,N);
#endif
            break;
        }

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

#if CG_DEBUG
    feLog("\nx=<%s>\nA*x=<%s>\n",c_print(x),c_print(A*x));
    feLog("\nb=<%s>\n",c_print(b));
#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("ConjugateGradient::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",
                    c_print(A),c_print(x),
                    c_print(Ax),c_print(b));
        }
//      feX("ConjugateGradient::solve","failed to converge");
    }
#endif
}

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

#endif /* __solve_ConjugateGradient_h__ */