hmc.h

Go to the documentation of this file.

#include <functional>
#include <random>
#include <Eigen/Dense>
#include <cmp_defines.h>
 
namespace cmp::mcmc {
 
 
class HamiltonianMarkovChain {
  private:
    std::default_random_engine& rng_; 
 
    Eigen::VectorXd currentPosition_; 
    Eigen::VectorXd currentMomentum_; 
    double currentScore_;             
    size_t dim_;                      
 
    // keep track of the mean and covariance of the samples
    Eigen::VectorXd mean_;            
    Eigen::MatrixXd cov_;             
 
    double epsilon_;                  
 
    size_t nSteps_ = 0;               
    double sumAcceptanceProb_ = 0.0;  
 
    // Dual Averaging parameters for step size adaptation
    double mu_;                       
    double gamma_ = 0.05;             
    double t0_ = 10.0;                
    double kappa_ = 0.75;             
 
    // Dual Averaging State Trackers
    double H_bar_ = 0.0;              
    double log_epsilon_bar_ = 0.0;    
    double target_accept_ = 0.80;     
 
 
    // Proposal distributions for momentum and uniform random numbers
    std::normal_distribution<double> distN_{0.0, 1.0};      
    std::uniform_real_distribution<double> distU_{0.0, 1.0}; 
 
 
    void leapfrog(Eigen::VectorXd& q, Eigen::VectorXd& p, Eigen::VectorXd& grad, double epsilon, const gradient_t& getGradient) const;
 
 
    struct TreeState {
        Eigen::VectorXd q_minus;    
        Eigen::VectorXd p_minus;    
        Eigen::VectorXd grad_minus; 
 
        Eigen::VectorXd q_plus;     
        Eigen::VectorXd p_plus;     
        Eigen::VectorXd grad_plus;  
 
        Eigen::VectorXd q_prime;    
        size_t n;                   
        int s;                      
        double alpha;               
        size_t n_alpha;             
    };
 
    TreeState build_tree(const Eigen::VectorXd& q, const Eigen::VectorXd& p, const Eigen::VectorXd& grad, double log_u, int v, int j, double epsilon, const score_t& getScore, const gradient_t& getGradient);
 
    void update();
 
  public:
    HamiltonianMarkovChain(Eigen::VectorXd initialState, std::default_random_engine& rng, double epsilon);
 
    void step(const score_t& getScore, const gradient_t& getGradient, bool adapt = false);
    void reset();
 
    Eigen::VectorXd getCurrent() const;
    size_t getSteps() const;
    double getAcceptanceRatio() const;
 
    double getStepSize() const {
        return epsilon_;
    }
 
    Eigen::VectorXd getMean() const {
        return mean_;
    }
 
    Eigen::MatrixXd getCovariance() const {
        return cov_ / static_cast<double>(std::max(nSteps_ - 1, (size_t)1));
    }
 
    void info() const;
};
 
}