gmres.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef pRC_ALGORITHMS_SOLVER_GMRES_H
#define pRC_ALGORITHMS_SOLVER_GMRES_H
 
#include <prc/algorithms/jacobi_rotation.hpp>
#include <prc/algorithms/solve.hpp>
#include <prc/algorithms/solver/backward_substitution.hpp>
#include <prc/core/basic/functions/copy.hpp>
#include <prc/core/complex/type_traits.hpp>
#include <prc/core/container/array.hpp>
#include <prc/core/container/deque.hpp>
#include <prc/core/functors/eval.hpp>
#include <prc/core/tensor/operator/functions/apply.hpp>
#include <prc/core/tensor/type_traits.hpp>
#include <prc/core/value/functions/is_approx.hpp>
#include <prc/core/value/limits.hpp>
 
namespace pRC::Solver
{
    template<Size M = 32, Size K = 2, If<IsSatisfied<(M > K)>> = 0>
    class GMRES
    {
    private:
        static constexpr Size defaultMaxIterations()
        {
            return 10000;
        }
 
    public:
        constexpr GMRES(Size const maxIterations = defaultMaxIterations())
            : mMaxIterations(maxIterations)
        {
        }
 
        constexpr auto maxIterations() const
        {
            return mMaxIterations;
        }
 
        template<Operator::Transform OT = Operator::Transform::None,
            Operator::Restrict OR = Operator::Restrict::None,
            Operator::Hint OH = Operator::Hint::None, class XA,
            class RA = RemoveConstReference<XA>, class TA = typename RA::Type,
            class VA = typename TA::Value, class XB,
            class RB = RemoveConstReference<XB>, class TB = typename RB::Type,
            class VB = typename TB::Value, class XX = decltype(zero<RB>()),
            class RX = RemoveConstReference<XX>, class TX = typename RX::Type,
            class VX = typename TX::Value, class VT = Common<VA, VB, VX>,
            If<All<IsFloat<VA>, IsFloat<VB>, IsFloat<VX>, IsFloat<VT>>> = 0>
        inline constexpr auto operator()(XA &&A, XB &&b, XX &&x0 = zero<RX>(),
            VT const &tolerance = NumericLimits<VT>::tolerance()) const
        {
            decltype(auto) x =
                copy<!(!IsReference<XX>() && !IsConst<RX>())>(eval(x0));
 
            using V = RemoveConstReference<ResultOf<Eval,
                decltype(round(apply<OT, OR, OH>(forward<XA>(A), x)))>>;
            using T = typename V::Type;
 
            auto const scaledTolerance = tolerance * norm(b)();
 
            Deque<V, K> z;
            Bool converged = false;
            Index iteration = 0;
            do
            {
                Array<V, M + 1> v;
                Array<JacobiRotation<T>, M> G;
                Tensor<T, M> g = zero();
                Tensor<T, M, M> R = identity();
 
                v[0] = apply<OT, OR, OH>(A, x);
                v[0] = b - v[0];
 
                auto vNorm = eval(norm(v[0]));
                auto error = vNorm();
 
                Index m = 0;
                while(m < M)
                {
                    if(norm(error) <= scaledTolerance)
                    {
                        converged = true;
                        break;
                    }
 
                    ++iteration;
                    v[m] /= vNorm;
 
                    if constexpr(K > 0)
                    {
                        if(m >= M - z.size())
                        {
                            v[m + 1] = apply<OT, OR, OH>(A,
                                z.front(m - (M - z.size())));
                        }
                        else
                        {
                            v[m + 1] = apply<OT, OR, OH>(A, v[m]);
                        }
                    }
                    else
                    {
                        v[m + 1] = apply<OT, OR, OH>(A, v[m]);
                    }
 
                    for(Index i = 0; i < m + 1; ++i)
                    {
                        R(i, m) = innerProduct(v[m + 1], v[i])();
                        v[m + 1] -= R(i, m) * v[i];
                    }
 
                    for(Index i = 0; i < m; ++i)
                    {
                        apply(G[i], chip<1>(R, m), i, i + 1);
                    }
 
                    vNorm = norm(v[m + 1]);
 
                    G[m] = JacobiRotation<T>::MakeGivens(R(m, m), vNorm());
                    R(m, m) = G[m].c() * R(m, m) + conj(G[m].s()) * vNorm();
 
                    g(m) = G[m].c() * error;
                    error *= -G[m].s();
 
                    ++m;
 
                    Logging::debug("GMRES", "Iteration:", iteration,
                        "Inner Iteration:", m, "Residual:", norm(error),
                        "Target:", scaledTolerance);
                }
 
                auto const y = solve<BackwardSubstitution>(R, g);
 
                if constexpr(K > 0)
                {
                    auto dx = eval(v[0] * y(0));
                    for(Index i = 1; i < m && i < M - z.size(); ++i)
                    {
                        dx += v[i] * y(i);
                    }
                    for(Index i = 0; i < z.size(); ++i)
                    {
                        dx += z.front(i) * y((M - z.size()) + i);
                    }
                    x += dx;
 
                    auto const dxn = eval(norm(dx));
                    z.pushFront(dx / dxn);
                }
                else
                {
                    for(Index i = 0; i < m; ++i)
                    {
                        x += v[i] * y(i);
                    }
                }
            }
            while(iteration < maxIterations() && !converged);
 
            if(!converged)
            {
                if constexpr (M <= 512)
                {
                    Logging::info(
                        "GMRES failed to converge within allowed max iterations for M =",
                        M, "and K =", K, "- doubling M and K and continuing");
                    x = GMRES<2*M, 2*K>()(A, b, x, tolerance);
                    return x;
                }
                else
                {
                    Logging::error(
                        "GMRES failed to converge within allowed max iterations.");
                }
            }
 
            if constexpr(cDebugLevel >= DebugLevel::High)
            {
                auto const y = eval(apply<OT, OR, OH>(A, x));
                if(!isApprox(y, b, tolerance))
                {
                    Logging::error("GMRES failed.");
                }
            }
 
            Logging::debug("GMRES converged after", iteration, "iterations.");
 
            if constexpr(IsReference<decltype(x)>())
            {
                return forward<XX>(x0);
            }
            else
            {
                return x;
            }
        }
 
    private:
        Size const mMaxIterations;
    };
}
#endif // pRC_ALGORITHMS_SOLVER_GMRES_H