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/container/deque.hpp>
#include <prc/core/functors/apply.hpp>
#include <prc/core/functors/eval.hpp>
#include <prc/core/tensor/declarations.hpp>
#include <prc/core/tensor/operator/hint.hpp>
#include <prc/core/tensor/operator/restrict.hpp>
#include <prc/core/tensor/operator/transform.hpp>
 
namespace pRC::Solver
{
    template<Size M = 32, Bool MGS = true, Size K = 0>
        requires(M > K)
    class GMRES
    {
    private:
        static constexpr Size defaultMaxIterations()
        {
            return 1000;
        }
 
    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 = RemoveReference<XA>, class XB,
            class RB = RemoveReference<XB>, class XX = decltype(zero<RB>()),
            class RX = RemoveReference<XX>,
            IsFloat VT = Common<Value<RA>, Value<RB>, Value<RX>>>
            requires IsInvocable<Apply<OT, OR, OH>, XA, XX> &&
            IsInvocable<Sub, XB, ResultOf<Apply<OT, OR, OH>, XA, XX>>
        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 + 1> g = zero();
                Tensor<T, M + 1, M> R = identity();
 
                v[0] = apply<OT, OR, OH>(A, x);
                v[0] = b - v[0];
 
                auto vNorm = norm(v[0])();
                g(0) = vNorm;
 
                Index m = 0;
                while(m < M)
                {
                    if(norm(g(m)) <= 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])();
                        if constexpr(MGS)
                        {
                            v[m + 1] -= R(i, m) * v[i];
                        }
                    }
 
                    if constexpr(!MGS)
                    {
                        for(Index i = 0; i < m + 1; ++i)
                        {
                            v[m + 1] -= R(i, m) * v[i];
                        }
                    }
 
                    vNorm = norm(v[m + 1])();
                    R(m + 1, m) = vNorm;
 
                    for(Index i = 0; i < m; ++i)
                    {
                        apply(G[i], chip<1>(R, m), i, i + 1);
                    }
 
                    G[m] = JacobiRotation<T>(chip<1>(R, m), m, m + 1);
                    apply(G[m], chip<1>(R, m), m, m + 1);
                    apply(G[m], g, m, m + 1);
 
                    ++m;
 
                    Logging::debug("GMRES", "Iteration:", iteration,
                        "Inner Iteration:", m, "Residual:", norm(g(m)),
                        "Target:", scaledTolerance);
                }
 
                auto const y = solve<BackwardSubstitution>(slice<M, M>(R, 0, 0),
                    slice<M>(g, 0));
 
                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");
                    return GMRES<2 * M, true, 2 * K>()(A, b, x, tolerance);
                }
                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(b, y, 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