learn_theta.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: BSD-2-Clause
 
#ifndef cMHN_TT_LEARN_THETA_H
#define cMHN_TT_LEARN_THETA_H
 
#include <iostream>
#include <map>
#include <string>
#include <tuple>
 
#include <cmhn/common/calculate_pTheta.hpp>
#include <cmhn/tt/als.hpp>
#include <cmhn/tt/mamen.hpp>
#include <cmhn/tt/mhn_operator.hpp>
#include <cmhn/tt/utility.hpp>
#include <cmhn/utility/read_data.hpp>
#include <cmhn/utility/read_header.hpp>
#include <cmhn/utility/write_theta.hpp>
 
#include <prc.hpp>
 
namespace cMHN::TT
{
    template<pRC::Size RP, pRC::Size RQ, class T, pRC::Size D, class S>
    std::tuple<pRC::Tensor<T, D, D>,
        std::map<std::string, std::string>,
        std::map<std::string, double>>
    learnTheta(pRC::Tensor<T, D, D> const &theta,
        std::string const &header, std::string const &output,
        std::map<S, T> const &pD, cMHN::Score<T> const &Score,
        cMHN::Regulator<T, D> const &Regulator, T const &toleranceOptimizer,
        T const &toleranceSolverP, T const &toleranceSolverQ)
    {
        using ModeSizes = decltype(getModeSizes<D>());
 
        auto tempTheta = theta;
 
        T score = pRC::zero();
 
        pRC::Index at_iter = 0;
 
        pRC::Float<64> startTime = pRC::getTimeInSeconds();
 
        std::map<std::string, double> logInfoNumbers{{"Score", score()},
            {"Iterations", at_iter},
            {"Time", pRC::getTimeInSeconds()() - startTime()},
            {"Lambda", Regulator.lambda()()}};
 
        std::map<std::string, std::string> logInfoNames{
            {"Score Name", Score.name()}, {"Regulator Name", Regulator.name()}};
 
        writeTheta(output, header, tempTheta, logInfoNames, logInfoNumbers);
 
        std::cout << "cMHN learning started (TT):" << std::endl;
        std::cout << "\tScore Name:\t" << logInfoNames["Score Name"]
                  << std::endl;
        std::cout << "\tRegulator Name:\t" << logInfoNames["Regulator Name"]
                  << std::endl
                  << std::endl;
 
        // get first pInit
        pRC::SeedSequence seq(8, 16);
        pRC::RandomEngine rng(seq);
        pRC::GaussianDistribution<pRC::Float<>> dist;
        auto pInit = round<decltype(getRanks<D, RP>())>(pRC::random<pRC::TensorTrain::Tensor
            <T, ModeSizes, decltype(getRanks<D, RP>())>>(rng, dist));
        pInit = pInit / norm(pInit);
 
        tempTheta = pRC::optimize<pRC::Optimizer::BFGS<>>(
            tempTheta,
            [&output, &at_iter, &score, &pD, &Score, &Regulator, &pInit,
                &toleranceOptimizer, &toleranceSolverP, &toleranceSolverQ](
                pRC::Tensor<T, D, D> const &tempTheta, pRC::Tensor<T, D, D> &g)
            {
                MHNOperator<T, D> op(tempTheta);
 
                auto const pTheta =
                    ::cMHN::calculatePTheta<RP>(op, pInit, toleranceSolverP);
 
                // use this pTheta as the next pInit
                pInit = pTheta;
 
                g = pRC::zero();
                score = pRC::zero();
 
                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 // _OPENMP
                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 = 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;
                g += gT;
 
                // signs b/c pRC::Optimmizer::BFGS minimizes the score function
                g = Regulator.grad(tempTheta) - g;
                return Regulator.score(tempTheta) - score;
            },
            [&output, &header, &score, &at_iter, &startTime, &logInfoNames,
                &logInfoNumbers](auto const &tempTheta)
            {
                logInfoNumbers["Iterations"] = at_iter;
                logInfoNumbers["Score"] = score();
                logInfoNumbers["Time"] =
                    pRC::getTimeInSeconds()() - startTime();
 
                std::cout << "cMHN learning in progress (TT):" << std::endl;
                std::cout << std::defaultfloat;
                std::cout << "\tIteration:\t" << logInfoNumbers["Iterations"]
                          << std::endl;
                std::cout << std::scientific;
                std::cout << "\tLambda:\t\t" << logInfoNumbers["Lambda"]
                          << std::endl;
                std::cout << "\tScore:\t\t" << logInfoNumbers["Score"]
                          << std::endl;
                std::cout << "\tTime:\t\t" << logInfoNumbers["Time"]
                          << std::endl;
                std::cout << std::defaultfloat;
 
                writeTheta(output, header, tempTheta, logInfoNames,
                    logInfoNumbers);
                at_iter++;
            },
            toleranceOptimizer);
 
        return std::make_tuple(tempTheta, logInfoNames, logInfoNumbers);
    }
}
 
#endif // cMHN_TT_LEARN_THETA_H