SupervisedMStep.hpp

#ifndef _LDAPLUSPLUS_EM_SUPERVISEDMSTEP_HPP_
#define _LDAPLUSPLUS_EM_SUPERVISEDMSTEP_HPP_

#include <vector>

#include "ldaplusplus/em/UnsupervisedMStep.hpp"

namespace ldaplusplus {
namespace em {


template <typename Scalar>
class SupervisedMStep : public UnsupervisedMStep<Scalar>
{
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> MatrixX;
    typedef Eigen::Matrix<Scalar, Eigen::Dynamic, 1> VectorX;

    public:
        SupervisedMStep(
            size_t m_step_iterations = 10,
            Scalar m_step_tolerance = 1e-2,
            Scalar regularization_penalty = 1e-2
        ) : m_step_iterations_(m_step_iterations),
            m_step_tolerance_(m_step_tolerance),
            regularization_penalty_(regularization_penalty),
            docs_(0)
        {}

        virtual void m_step(
            std::shared_ptr<parameters::Parameters> parameters
        ) override;

        virtual void doc_m_step(
            const std::shared_ptr<corpus::Document> doc,
            const std::shared_ptr<parameters::Parameters> v_parameters,
            std::shared_ptr<parameters::Parameters> m_parameters
        ) override;

    private:
        // The maximum number of iterations in M-step
        size_t m_step_iterations_;
        // The convergence tolerance for the maximazation of the ELBO w.r.t.
        // eta in M-step
        Scalar m_step_tolerance_;
        // The regularization penalty for the multinomial logistic regression
        Scalar regularization_penalty_;

        // Number of documents processed so far
        int docs_;
        MatrixX phi_scaled;
        MatrixX expected_z_bar_;
        std::vector<MatrixX> variance_z_bar_;
        Eigen::VectorXi y_;
};

}  // namespace em
}  // namespace ldaplusplus

#endif  // _LDAPLUSPLUS_EM_SUPERVISEDMSTEP_HPP_