mhn_operator.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef cMHN_TT_MHN_OPERATOR_H
#define cMHN_TT_MHN_OPERATOR_H
 
#include <prc.hpp>
 
namespace cMHN::TT
{
    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));
        }
 
        // calculates one TT term in eq. (2.2) of Peter's Thesis
        // -> Q=sum_i term(i)
        constexpr auto term(pRC::Index const i) const
        {
            TermType ret;
            pRC::range<pRC::Context::CompileTime, D>(
                [this, &ret, i](auto const j)
                {
                    ret.template core<j>() = this->core(i, j);
                });
 
            return ret;
        }
 
        // calculates a core of the TT (1-Q)
        template<pRC::Index i>
        constexpr auto core() const
        {
            return expand(pRC::makeSeries<pRC::Index, D>(),
                [&](auto const... seq)
                {
                    return (pRC::identity<TermType>() - (term(seq) + ...))
                        .template core<i>();
                });
        }
 
        // calculates dQ/d(theta_ij) - note that the derivative is wrt small
        // thetas
        constexpr auto derivative(pRC::Index const i, pRC::Index const j) const
        {
            TermType ret;
            pRC::range<pRC::Context::CompileTime, D>(
                [this, &ret, i, j](auto const m)
                {
                    ret.template core<m>() = this->derivative(i, j, m);
                });
 
            return ret;
        }
 
    private:
        template<pRC::Index... Ms, pRC::Index... Ns, pRC::Index... Rs>
        static constexpr auto termType(pRC::Sequence<pRC::Index, Ms...> const,
            pRC::Sequence<pRC::Index, Ns...> const,
            pRC::Sequence<pRC::Index, Rs...> const)
        {
            using M = pRC::Sizes<Ms...>;
            using N = pRC::Sizes<Ns...>;
            using R = pRC::Sizes<Rs...>;
 
            return pRC::TensorTrain::Operator<T, M, N, R>();
        }
 
        using TermType =
            decltype(termType(pRC::makeConstantSequence<pRC::Index, D, 2>(),
                pRC::makeConstantSequence<pRC::Index, D, 2>(),
                pRC::makeConstantSequence<pRC::Index, D - 1, 1>()));
 
    private:
        // cores for calculation of Q
        constexpr auto core(pRC::Size const i, pRC::Size const j) const
        {
            pRC::Tensor c = pRC::zero<pRC::Tensor<T, 1, 2, 2, 1>>();
            if(i == j)
            {
                c(0, 0, 0, 0) = -mBigTheta(i, i);
                c(0, 1, 0, 0) = mBigTheta(i, i);
            }
            else
            {
                c(0, 0, 0, 0) = pRC::identity<T>();
                c(0, 1, 1, 0) = mBigTheta(i, j);
            }
            return c;
        }
 
        constexpr auto derivative(pRC::Size const i, pRC::Size const j,
            pRC::Size const m) const
        {
            pRC::Tensor c = pRC::zero<pRC::Tensor<T, 1, 2, 2, 1>>();
            if(j == m)
            {
                if(i == j)
                {
                    c(0, 0, 0, 0) = mBigTheta(i, m);
                    c(0, 1, 0, 0) = -mBigTheta(i, m);
                }
                else
                {
                    c(0, 1, 1, 0) = mBigTheta(i, m);
                }
            }
            else
            {
                if(i == m)
                {
                    c(0, 0, 0, 0) = mBigTheta(i, m);
                    c(0, 1, 0, 0) = -mBigTheta(i, m);
                }
                else
                {
                    c(0, 0, 0, 0) = pRC::identity<T>();
                    c(0, 1, 1, 0) = mBigTheta(i, m);
                }
            }
 
            return c;
        }
 
    private:
        pRC::Tensor<T, D, D> mBigTheta;
    };
}
 
#endif // cMHN_TT_MHN_OPERATOR_H