calculate_score_and_gradient.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef cMHN_COMMON_CALCULATE_SCORE_AND_GRADIENT_H
#define cMHN_COMMON_CALCULATE_SCORE_AND_GRADIENT_H
 
#include <map>
#include <tuple>
 
#include <cmhn/common/calculate_pTheta.hpp>
#include <cmhn/common/regulator.hpp>
#include <cmhn/common/score.hpp>
#include <cmhn/nontt/mhn_operator.hpp>
#include <cmhn/tt/mhn_operator.hpp>
 
#include <prc.hpp>
 
namespace cMHN
{
    template<class T, pRC::Size D, class S>
    std::tuple<T, pRC::Tensor<T, D, D>> calculateScoreAndGradient(
        nonTT::MHNOperator<T, D> const &op, std::map<S, T> const &pD,
        cMHN::Score<T> const &Score, cMHN::Regulator<T, D> const &Regulator,
        T const &toleranceSolverQ = 1e-8)
    {
        T score = pRC::zero();
        pRC::Tensor<T, D, D> g = pRC::zero();
 
        auto const pTheta = calculatePTheta(op);
 
        for(auto const &[k, v] : pD)
        {
            score += Score.pointwiseScore(v, pTheta(k));
        }
 
        // calculate g = dS/dtheta
        auto rhs = eval(expand(pRC::makeConstantSequence<pRC::Size, D, 2>(),
            [](auto const... Ns)
            {
                return pRC::zero<pRC::Tensor<T, Ns...>>();
            }));
 
        for(auto const &[k, v] : pD)
        {
            rhs(k) = Score.pointwiseDSDP(v, pTheta(k));
        }
 
        auto const q = pRC::solve<pRC::Solver::GMRES<>,
            pRC::Operator::Transform::Transpose>(op, rhs,
            eval(pRC::zero<decltype(rhs)>()), toleranceSolverQ);
 
        // this follows appendix C of Rudi's Thesis
        for(pRC::Index i = 0; i < D; ++i)
        {
            auto r = hadamardProduct(q, nonTT::applyDerivative(op, pTheta, i));
 
            pRC::range<pRC::Context::CompileTime, D>(
                [&](auto const j)
                {
                    if(i == j)
                    {
                        g(i, j) = pRC::reduce<pRC::Add>(r)();
                    }
                    else
                    {
                        g(i, j) = pRC::reduce<pRC::Add>(pRC::chip<j>(r, 1))();
                    }
                });
        }
 
        score += Regulator.score(log(op.bigTheta()));
        g += Regulator.grad(log(op.bigTheta()));
 
        return std::make_tuple(score, g);
    }
 
    template<pRC::Size RP, pRC::Size RQ, class T, pRC::Size D, class S, class X>
    std::tuple<T, pRC::Tensor<T, D, D>> calculateScoreAndGradient(
        TT::MHNOperator<T, D> const &op, std::map<S, T> const &pD,
        cMHN::Score<T> const &Score, cMHN::Regulator<T, D> const &Regulator,
        X &pInit, T const &toleranceSolverP = 1e-4,
        T const &toleranceSolverQ = 1e-4)
    {
        using ModeSizes = decltype(TT::getModeSizes<D>());
 
        T score = pRC::zero();
        pRC::Tensor<T, D, D> g = pRC::zero();
 
        auto const pTheta = calculatePTheta<RP>(op, pInit, toleranceSolverP);
 
        T scoreT = pRC::zero();
        pRC::Tensor<T, D, D> gT = pRC::zero();
#if defined(_OPENMP)
    #pragma omp declare reduction(+ : T, \
            pRC::Tensor<T, D, D> : omp_out = omp_in + omp_out) \
        initializer(omp_priv(pRC::Zero()))
    #pragma omp parallel for schedule(dynamic, 10) reduction(+ : scoreT, gT)
#endif
        for(pRC::Index s = 0; s < pD.size(); ++s)
        {
            auto it = pD.cbegin();
            std::advance(it, s);
            auto const [k, v] = *it;
            auto const pThetaE = pTheta(k);
 
            scoreT += Score.pointwiseScore(v, pThetaE);
 
            auto const rhs = pRC::TensorTrain::Tensor<T, ModeSizes>::Single(
                pRC::identity<T>(), k);
 
            // solve (1-Q)T * q = rhs
            auto const q = TT::als<RQ, pRC::Operator::Transform::Transpose>(op,
                rhs, toleranceSolverQ);
 
            pRC::Tensor<T, D, D> tmp = pRC::zero();
            for(pRC::Index i = 0; i < D; ++i)
            {
                for(pRC::Index j = 0; j < D; ++j)
                {
                    tmp(i, j) = Score.pointwiseDSDP(v, pThetaE) *
                        scalarProduct(q, op.derivative(i, j) * pTheta)();
                }
            }
 
            gT += tmp;
        }
 
        score += scoreT + Regulator.score(log(op.bigTheta()));
        g += gT + Regulator.grad(log(op.bigTheta()));
 
        // update pInit
        pInit = pTheta;
 
        return std::make_tuple(score, g);
    }
 
    template<pRC::Size RP, pRC::Size RQ, class T, pRC::Size D, class S>
    std::tuple<T, pRC::Tensor<T, D, D>> calculateScoreAndGradient(
        TT::MHNOperator<T, D> const &op, std::map<S, T> const &pD,
        cMHN::Score<T> const &Score, cMHN::Regulator<T, D> const &Regulator,
        T const &toleranceSolverP = 1e-4, T const &toleranceSolverQ = 1e-4)
    {
        using ModeSizes = decltype(TT::getModeSizes<D>());
        using Ranks = decltype(TT::getRanks<D, RP>());
 
        // use a random pInit
        pRC::SeedSequence seq(8, 16);
        pRC::DefaultRandomEngine rng(seq);
        pRC::GaussianDistribution<pRC::Float<>> dist;
        auto pInit = round<Ranks>(
            pRC::random<pRC::TensorTrain::Tensor<T, ModeSizes, Ranks>>(rng,
                dist));
        pInit /= scalarProduct(pInit,
            pRC::unit<pRC::TensorTrain::Tensor<T, ModeSizes>>())();
 
        return calculateScoreAndGradient<RP, RQ>(op, pD, Score, Regulator,
            pInit, toleranceSolverP, toleranceSolverQ);
    }
}
 
#endif