mhn_operator.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef cMHN_NONTT_MHN_OPERATOR_H
#define cMHN_NONTT_MHN_OPERATOR_H
 
#include <prc.hpp>
 
namespace cMHN::nonTT
{
    template<class T, pRC::Size D>
    class MHNOperator
    {
    public:
        using Type = T;
 
    public:
        constexpr MHNOperator(pRC::Tensor<T, D, D> const &smallTheta)
            : mBigTheta(exp(smallTheta))
        {
        }
 
        constexpr auto &bigTheta(pRC::Index const i, pRC::Index const j) const
        {
            return mBigTheta(i, j);
        }
 
        constexpr auto &bigTheta() const
        {
            return mBigTheta;
        }
 
        auto &bigTheta()
        {
            return mBigTheta;
        }
 
        constexpr auto smallTheta(pRC::Index const i, pRC::Index const j) const
        {
            return log(mBigTheta(i, j));
        }
 
        constexpr auto smallTheta() const
        {
            return eval(log(mBigTheta));
        }
 
    private:
        pRC::Tensor<T, D, D> mBigTheta;
    };
 
    // apply dQ/dtheta_ii to a vector x
    template<pRC::Operator::Transform OT = pRC::Operator::Transform::None,
        pRC::Operator::Restrict OR = pRC::Operator::Restrict::None,
        pRC::Operator::Hint OH = pRC::Operator::Hint::None, class T1,
        pRC::Size D, class T2, pRC::Size... Ns>
        requires(OR == pRC::Operator::Restrict::None) &&
        (OT == pRC::Operator::Transform::None)
    static inline constexpr auto applyDerivative(MHNOperator<T1, D> const &op,
        pRC::Tensor<T2, Ns...> const &x, pRC::Index const &i)
 
    {
        auto r = x;
        pRC::range<pRC::Context::CompileTime, D>(
            [&](auto const j)
            {
                if(i == j)
                {
                    pRC::chip<j>(r, 0) *= -op.bigTheta(i, i);
                    pRC::chip<j>(r, 1) = -pRC::chip<j>(r, 0);
                }
                else
                {
                    pRC::chip<j>(r, 1) *= op.bigTheta(i, j);
                }
            });
        return r;
    }
 
    template<pRC::Operator::Transform OT = pRC::Operator::Transform::None,
        pRC::Operator::Restrict OR = pRC::Operator::Restrict::None,
        pRC::Operator::Hint OH = pRC::Operator::Hint::None, class T1,
        pRC::Size D, class T2, pRC::Size... Ns>
        requires(OR == pRC::Operator::Restrict::None)
    static inline constexpr auto apply_diag(MHNOperator<T1, D> const &op,
        pRC::Tensor<T2, Ns...> const &x)
    {
        auto b = x;
        for(pRC::Index i = 0; i < D; ++i)
        {
            auto v = x;
            pRC::range<pRC::Context::CompileTime, D>(
                [&](auto const j)
                {
                    if(i == j)
                    {
                        pRC::chip<j>(v, 0) *= -op.bigTheta(i, i);
                        pRC::chip<j>(v, 1) = pRC::zero();
                    }
                    else
                    {
                        pRC::chip<j>(v, 1) *= op.bigTheta(i, j);
                    }
                });
            b -= v;
        }
        return b;
    }
 
    template<pRC::Operator::Transform OT = pRC::Operator::Transform::None,
        pRC::Operator::Restrict OR = pRC::Operator::Restrict::None,
        pRC::Operator::Hint OH = pRC::Operator::Hint::None, class T1,
        pRC::Size D, class T2, pRC::Size... Ns>
        requires(OR == pRC::Operator::Restrict::None)
    static inline constexpr auto apply_offdiag(MHNOperator<T1, D> const &op,
        pRC::Tensor<T2, Ns...> const &x)
    {
        pRC::Tensor r = x;
        r = pRC::zero();
        for(pRC::Index i = 0; i < D; ++i)
        {
            auto v = x;
            pRC::range<pRC::Context::CompileTime, D>(
                [&](auto const j)
                {
                    if(i == j)
                    {
                        if constexpr(OT == pRC::Operator::Transform::Transpose)
                        {
                            pRC::chip<j>(v, 0) =
                                op.bigTheta(i, i) * pRC::chip<j>(v, 1);
                            pRC::chip<j>(v, 1) = pRC::zero();
                        }
                        else
                        {
                            pRC::chip<j>(v, 1) =
                                op.bigTheta(i, i) * pRC::chip<j>(v, 0);
                            pRC::chip<j>(v, 0) = pRC::zero();
                        }
                    }
                    else
                    {
                        pRC::chip<j>(v, 1) *= op.bigTheta(i, j);
                    }
                });
            r += v;
        }
        return r;
    }
 
    // apply (1-Q) or the transposed to a vector x
    template<pRC::Operator::Transform OT = pRC::Operator::Transform::None,
        pRC::Operator::Restrict OR = pRC::Operator::Restrict::None,
        pRC::Operator::Hint OH = pRC::Operator::Hint::None, class T1,
        pRC::Size D, class T2, pRC::Size... Ns>
        requires(OR == pRC::Operator::Restrict::None)
    static inline constexpr auto apply(MHNOperator<T1, D> const &op,
        pRC::Tensor<T2, Ns...> const &x)
 
    {
        auto b = x;
        for(pRC::Index i = 0; i < D; ++i)
        {
            auto v = x;
            pRC::range<pRC::Context::CompileTime, D>(
                [&](auto const j)
                {
                    if(i == j)
                    {
                        if constexpr(OT == pRC::Operator::Transform::Transpose)
                        {
                            pRC::chip<j>(v, 0) =
                                -op.bigTheta(i, i) * pRC::chip<j>(v, 0) +
                                op.bigTheta(i, i) * pRC::chip<j>(v, 1);
                            pRC::chip<j>(v, 1) = pRC::zero();
                        }
                        else
                        {
                            pRC::chip<j>(v, 0) *= -op.bigTheta(i, i);
                            pRC::chip<j>(v, 1) = -pRC::chip<j>(v, 0);
                        }
                    }
                    else
                    {
                        pRC::chip<j>(v, 1) *= op.bigTheta(i, j);
                    }
                });
            b -= v;
        }
        return b;
    }
 
}
 
#endif // cMHN_NONTT_MHN_OPERATOR_H