<commit_before>#ifndef HMM_MESSAGES_H
#define HMM_MESSAGES_H

#include <Eigen/Core>
#include <stdint.h> // int32_t
#include <limits> // infinity

#include "nptypes.h"
#include "util.h"

// NOTE: HMM_TEMPS_ON_HEAP is for robustness. for the OpenMP case, where dynamic
// allocation has locks and there might be false sharing, I force some
// temporaries onto the stack using compiler directives. It may be a dumb idea
// though; I haven't profiled it and I don't know how compiler-dependent it is.

// NOTE: HMM_NOT_ROBUST switch needs to be benchmarked and removed if it doesn't
// make a difference

namespace hmm
{
    using namespace std;
    using namespace Eigen;
    using namespace nptypes;

    // Messages

    template <typename Type>
    void messages_backwards_log(
            bool hetero, int M, int T, Type *As, Type *aBl, Type *betal)
    {
        NPMatrix<Type> eAs(As, hetero ? (T-1)*M : M, M);
        NPMatrix<Type> eaBl(aBl,T,M);

        NPMatrix<Type> ebetal(betal,T,M);

#ifdef HMM_TEMPS_ON_HEAP
        Matrix<Type,Dynamic,1> thesum(M);
#else
        Type thesum_buf[M] __attribute__((aligned(16)));
        NPVector<Type> thesum(thesum_buf,M);
#endif
        Type cmax;

        ebetal.row(T-1).setZero();
        for (int t=T-2; t>=0; t--) {
            thesum = (eaBl.row(t+1) + ebetal.row(t+1)).transpose();
            cmax = thesum.maxCoeff();
            ebetal.row(t) = (eAs.block(t*M*hetero,0,M,M)
                * (thesum.array() - cmax).exp().matrix()).array().log() + cmax;
        }
    }

    template <typename Type>
    void messages_forwards_log(
            bool hetero, int M, int T, Type *As, Type *pi0, Type *aBl, Type *alphal)
    {
        NPMatrix<Type> eAs(As, hetero ? (T-1)*M : M, M);
        NPArray<Type> epi0(pi0,1,M);
        NPArray<Type> eaBl(aBl,T,M);

        NPArray<Type> ealphal(alphal,T,M);

        Type cmax;

        ealphal.row(0) = epi0.log() + eaBl.row(0);
        for (int t=0; t<T-1; t++) {
            cmax = ealphal.row(t).maxCoeff();
#ifndef HMM_NOT_ROBUST
            if (likely(util::is_finite(cmax))) {
#endif
                ealphal.row(t+1) = ((ealphal.row(t) - cmax).exp().matrix()
                    * eAs.block(t*M*hetero,0,M,M)).array().log()
                    + cmax + eaBl.row(t+1);
#ifndef HMM_NOT_ROBUST
            } else {
                ealphal.block(t+1,0,T-(t+1),M).setConstant(-numeric_limits<Type>::infinity());
                return;
            }
#endif
        }
    }

    template <typename Type>
    Type expected_statistics_log(
            bool hetero, int M, int T,
            Type *log_trans_potential, Type *log_likelihood_potential,
            Type *alphal, Type *betal,
            Type *expected_states, Type *expected_transcounts)
    {
        NPArray<Type> eAs(log_trans_potential, hetero ? (T-1)*M : M, M);
        NPArray<Type> eaBl(log_likelihood_potential,T,M);
        NPArray<Type> ebetal(betal,T,M);
        NPArray<Type> ealphal(alphal,T,M);

        NPArray<Type> eexpected_states(expected_states,T,M);
        NPArray<Type> eexpected_transcounts(expected_transcounts,hetero ? (T-1)*M : M,M);

#ifdef HMM_TEMPS_ON_HEAP
        Array<Type,Dynamic,Dynamic> pair(M,M);
#else
        Type pair_buf[M*M] __attribute__((aligned(16)));
        NPArray<Type> pair(pair_buf,M,M);
#endif

        Type cmax;
        cmax = ealphal.row(T-1).maxCoeff();
        Type log_normalizer = log((ealphal.row(T-1) - cmax).exp().sum()) + cmax;

        for (int t=0; t<T-1; t++) {
            pair = eAs.block(t*M*hetero,0,M,M) - log_normalizer;
            pair.colwise() += ealphal.row(t).transpose().array();
            pair.rowwise() += ebetal.row(t+1) + eaBl.row(t+1);

            eexpected_transcounts.block(t*M*hetero,0,M,M) += pair.exp();
            eexpected_states.row(t) += (ealphal.row(t) + ebetal.row(t) - log_normalizer).exp();
        }
        eexpected_states.row(T-1) += (ealphal.row(T-1) - log_normalizer).exp();

        return log_normalizer;
    }

    template <typename Type>
    Type messages_forwards_normalized(
            bool hetero, int M, int T, Type *As, Type *pi0, Type *aBl, Type *alphan)
    {
        NPMatrix<Type> eAs(As, hetero ? (T-1)*M : M, M);
        NPArray<Type> eaBl(aBl,T,M);

        NPMatrix<Type> ealphan(alphan,T,M);

        Type logtot = 0.;
        Type cmax, norm;

#ifdef HMM_TEMPS_ON_HEAP
        Matrix<Type,1,Dynamic> ein_potential(1,M);
#else
        Type in_potential_buf[M] __attribute__((aligned(16)));
        NPRowVector<Type> ein_potential(in_potential_buf,M);
#endif

        ein_potential = NPRowVector<Type>(pi0,M);
        for (int t=0; t<T; t++) {
            cmax = eaBl.row(t).maxCoeff();
            ealphan.row(t) = ein_potential.array() * (eaBl.row(t) - cmax).exp();
            norm = ealphan.row(t).sum();
#ifndef HMM_NOT_ROBUST
            if (likely(norm != 0)) {
#endif
                ealphan.row(t) /= norm;
                logtot += log(norm) + cmax;
#ifndef HMM_NOT_ROBUST
            } else {
                ealphan.block(t,0,T-t,M).setZero();
                return -numeric_limits<Type>::infinity();
            }
#endif
            ein_potential = ealphan.row(t) * eAs.block(t*M*hetero,0,M,M);
        }
        return logtot;
    }

    // Sampling

    template <typename FloatType, typename IntType>
    void sample_forwards_log(
            bool hetero, int M, int T, FloatType *As, FloatType *pi0,
            FloatType *aBl, FloatType *betal,
            IntType *stateseq, FloatType *randseq)
    {
        NPMatrix<FloatType> eAs(As, hetero ? (T-1)*M : M, M);<commit_msg>add check for last iteration of forward message passing, needed for inhomog<commit_after>#ifndef HMM_MESSAGES_H
#define HMM_MESSAGES_H

#include <Eigen/Core>
#include <stdint.h> // int32_t
#include <limits> // infinity

#include "nptypes.h"
#include "util.h"

// NOTE: HMM_TEMPS_ON_HEAP is for robustness. for the OpenMP case, where dynamic
// allocation has locks and there might be false sharing, I force some
// temporaries onto the stack using compiler directives. It may be a dumb idea
// though; I haven't profiled it and I don't know how compiler-dependent it is.

// NOTE: HMM_NOT_ROBUST switch needs to be benchmarked and removed if it doesn't
// make a difference

namespace hmm
{
    using namespace std;
    using namespace Eigen;
    using namespace nptypes;

    // Messages

    template <typename Type>
    void messages_backwards_log(
            bool hetero, int M, int T, Type *As, Type *aBl, Type *betal)
    {
        NPMatrix<Type> eAs(As, hetero ? (T-1)*M : M, M);
        NPMatrix<Type> eaBl(aBl,T,M);

        NPMatrix<Type> ebetal(betal,T,M);

#ifdef HMM_TEMPS_ON_HEAP
        Matrix<Type,Dynamic,1> thesum(M);
#else
        Type thesum_buf[M] __attribute__((aligned(16)));
        NPVector<Type> thesum(thesum_buf,M);
#endif
        Type cmax;

        ebetal.row(T-1).setZero();
        for (int t=T-2; t>=0; t--) {
            thesum = (eaBl.row(t+1) + ebetal.row(t+1)).transpose();
            cmax = thesum.maxCoeff();
            ebetal.row(t) = (eAs.block(t*M*hetero,0,M,M)
                * (thesum.array() - cmax).exp().matrix()).array().log() + cmax;
        }
    }

    template <typename Type>
    void messages_forwards_log(
            bool hetero, int M, int T, Type *As, Type *pi0, Type *aBl, Type *alphal)
    {
        NPMatrix<Type> eAs(As, hetero ? (T-1)*M : M, M);
        NPArray<Type> epi0(pi0,1,M);
        NPArray<Type> eaBl(aBl,T,M);

        NPArray<Type> ealphal(alphal,T,M);

        Type cmax;

        ealphal.row(0) = epi0.log() + eaBl.row(0);
        for (int t=0; t<T-1; t++) {
            cmax = ealphal.row(t).maxCoeff();
#ifndef HMM_NOT_ROBUST
            if (likely(util::is_finite(cmax))) {
#endif
                ealphal.row(t+1) = ((ealphal.row(t) - cmax).exp().matrix()
                    * eAs.block(t*M*hetero,0,M,M)).array().log()
                    + cmax + eaBl.row(t+1);
#ifndef HMM_NOT_ROBUST
            } else {
                ealphal.block(t+1,0,T-(t+1),M).setConstant(-numeric_limits<Type>::infinity());
                return;
            }
#endif
        }
    }

    template <typename Type>
    Type expected_statistics_log(
            bool hetero, int M, int T,
            Type *log_trans_potential, Type *log_likelihood_potential,
            Type *alphal, Type *betal,
            Type *expected_states, Type *expected_transcounts)
    {
        NPArray<Type> eAs(log_trans_potential, hetero ? (T-1)*M : M, M);
        NPArray<Type> eaBl(log_likelihood_potential,T,M);
        NPArray<Type> ebetal(betal,T,M);
        NPArray<Type> ealphal(alphal,T,M);

        NPArray<Type> eexpected_states(expected_states,T,M);
        NPArray<Type> eexpected_transcounts(expected_transcounts,hetero ? (T-1)*M : M,M);

#ifdef HMM_TEMPS_ON_HEAP
        Array<Type,Dynamic,Dynamic> pair(M,M);
#else
        Type pair_buf[M*M] __attribute__((aligned(16)));
        NPArray<Type> pair(pair_buf,M,M);
#endif

        Type cmax;
        cmax = ealphal.row(T-1).maxCoeff();
        Type log_normalizer = log((ealphal.row(T-1) - cmax).exp().sum()) + cmax;

        for (int t=0; t<T-1; t++) {
            pair = eAs.block(t*M*hetero,0,M,M) - log_normalizer;
            pair.colwise() += ealphal.row(t).transpose().array();
            pair.rowwise() += ebetal.row(t+1) + eaBl.row(t+1);

            eexpected_transcounts.block(t*M*hetero,0,M,M) += pair.exp();
            eexpected_states.row(t) += (ealphal.row(t) + ebetal.row(t) - log_normalizer).exp();
        }
        eexpected_states.row(T-1) += (ealphal.row(T-1) - log_normalizer).exp();

        return log_normalizer;
    }

    template <typename Type>
    Type messages_forwards_normalized(
            bool hetero, int M, int T, Type *As, Type *pi0, Type *aBl, Type *alphan)
    {
        NPMatrix<Type> eAs(As, hetero ? (T-1)*M : M, M);
        NPArray<Type> eaBl(aBl,T,M);

        NPMatrix<Type> ealphan(alphan,T,M);

        Type logtot = 0.;
        Type cmax, norm;

#ifdef HMM_TEMPS_ON_HEAP
        Matrix<Type,1,Dynamic> ein_potential(1,M);
#else
        Type in_potential_buf[M] __attribute__((aligned(16)));
        NPRowVector<Type> ein_potential(in_potential_buf,M);
#endif

        ein_potential = NPRowVector<Type>(pi0,M);
        for (int t=0; t<T; t++) {
            cmax = eaBl.row(t).maxCoeff();
            ealphan.row(t) = ein_potential.array() * (eaBl.row(t) - cmax).exp();
            norm = ealphan.row(t).sum();
#ifndef HMM_NOT_ROBUST
            if (likely(norm != 0)) {
#endif
                ealphan.row(t) /= norm;
                logtot += log(norm) + cmax;
#ifndef HMM_NOT_ROBUST
            } else {
                ealphan.block(t,0,T-t,M).setZero();
                return -numeric_limits<Type>::infinity();
            }
#endif
            if (likely(t != T-1)) {
                ein_potential = ealphan.row(t) * eAs.block(t*M*hetero,0,M,M);
            }
        }
        return logtot;
    }

    // Sampling

    template <typename FloatType, typename IntType>
    void sample_forwards_log(
            bool hetero, int M, int T, FloatType *As, FloatType *pi0,
            FloatType *aBl, FloatType *betal,
            IntType *stateseq, FloatType *randseq)
    {
        NPMatrix<FloatType> eAs(As, hetero ? (T-1)*M : M, M);
{'fraction_non_alphanumeric': 0.14556737588652482, 'fraction_numerical': 0.010017730496453901, 'mean_word_length': 2.400964727163099, 'pattern_counts': {'":': 0, '<': 77, '<?xml version=': 0, '>': 73, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1859333', 'n_tokens_mistral': 4268, 'n_tokens_neox': 4093, 'n_words': 831}