mcmc.h

Go to the documentation of this file.

#ifndef MCMC_H
#define MCMC_H
 
#include "cmp_defines.h"
#include "distribution.h"
#include <algorithm>
#include <cmath>
#include <iostream>
#include <iomanip>
#include <vector>
#include <limits>
#include <random>
#include <Eigen/Dense>
 
namespace cmp::mcmc {
 
// ==============================================================================
// HELPER FUNCTIONS
// ==============================================================================
 
// Numerically safe log(1 - α)
inline double log1m(double a) {
    return (a < 1.0) ? std::log1p(-a) : -std::numeric_limits<double>::infinity();
}
 
// ==============================================================================
// MARKOV CHAIN CLASS TEMPLATE
// ==============================================================================
 
 
template <typename ProposalType>
class MarkovChain {
  protected:
    double score_{0.0};                            
    size_t dim_;                                   
    size_t nSteps_{0};                             
    size_t nAccepts_{0};                           
 
    Eigen::VectorXd mean_;                         
    Eigen::MatrixXd cov_;                          
 
    ProposalType proposal_;                        
 
    std::uniform_real_distribution<double> distU_{0.0, 1.0}; 
    std::default_random_engine rng_;               
 
    double scale_;                                 
    double targetAcceptanceRatio_;                 
 
  public:
    MarkovChain(const ProposalType &proposal,
                std::default_random_engine &rng,
                const double &targetAcceptanceRatio = 0.234)
        : rng_(rng),
          proposal_(proposal),
          score_(-std::numeric_limits<double>::infinity()),
          dim_(proposal.get().size()),
          targetAcceptanceRatio_(targetAcceptanceRatio),
          mean_(Eigen::VectorXd::Zero(dim_)),
          cov_(Eigen::MatrixXd::Zero(dim_, dim_)),
          scale_(5.6644 / static_cast<double>(dim_)) // 2.38^2 = 5.6644
    {}
 
    MarkovChain() = default;
    ~MarkovChain() = default;
 
    void increaseSteps() {
        nSteps_++;
    }
 
    bool accept(const Eigen::Ref<const Eigen::VectorXd> &par, double score) {
        if((score - score_) > std::log(distU_(rng_))) {
            proposal_.set(par);
            score_ = score;
            nAccepts_++;
            return true;
        }
        return false;
    }
 
    void step(const score_t &getScore, const bool &delayedRejection = false, const double& gamma = 0.1) {
        increaseSteps();
 
        // ========================================================================
        // STAGE 1: Standard Proposal
        // ========================================================================
        Eigen::VectorXd cand_prev = proposal_.sample(rng_, 1.0);
        double score_prev = getScore(cand_prev);
 
        // Calculate Stage 1 continuous acceptance probability (alpha_1_fwd)
        double log_alpha_1_fwd = std::min(0.0, score_prev - score_);
        double alpha_1_fwd = std::exp(log_alpha_1_fwd);
 
        bool accepted = accept(cand_prev, score_prev);
 
        // Robbins-Monro Global Scale Update (Tuning to First-Stage efficiency)
        double rm_gamma = 1.0 / std::sqrt(static_cast<double>(nSteps_));
        scale_ *= std::exp(rm_gamma * (alpha_1_fwd - targetAcceptanceRatio_));
 
        // ========================================================================
        // STAGE 2: Delayed Rejection (DRAM Detailed Balance)
        // ========================================================================
        if(!accepted && delayedRejection) {
 
            // Use the first gamma for the narrower DR backup proposal
            Eigen::VectorXd cand_new = proposal_.sample(rng_, gamma);
            double score_new = getScore(cand_new);
 
            // 1. The Backward Rejection Probability
            double log_alpha_1_bwd = std::min(0.0, score_prev - score_new);
            double alpha_1_bwd = std::exp(log_alpha_1_bwd);
 
            double log_rej_fwd = log1m(alpha_1_fwd);
            double log_rej_bwd = log1m(alpha_1_bwd);
 
            // 2. The Proposal Asymmetry
            Eigen::VectorXd jump_fwd = cand_prev - getCurrent();
            Eigen::VectorXd jump_bwd = cand_prev - cand_new;
 
            double log_q1_fwd = -0.5 * proposal_.squaredMahalanobis(jump_fwd);
            double log_q1_bwd = -0.5 * proposal_.squaredMahalanobis(jump_bwd);
 
            // 3. The True DRAM Stage-2 Acceptance Ratio
            double log_ratio_stage2 = (score_new - score_) +         // Pi ratio
                                      (log_q1_bwd - log_q1_fwd) +    // Proposal ratio
                                      (log_rej_bwd - log_rej_fwd);   // Rejection ratio
 
            double alpha_2 = std::exp(std::min(0.0, log_ratio_stage2));
 
            // Final Stage-2 Coin Flip
            if(std::log(distU_(rng_)) < std::log(alpha_2)) {
                proposal_.set(cand_new);
                score_ = score_new;
                accepted = true;
            }
        }
 
        update();
    }
 
    void update() {
        Eigen::VectorXd par = proposal_.get();
 
        if(nSteps_ == 1) {
            mean_ = par;
            cov_.setZero();
        } else {
            // Welford's stable 1-pass multi-dimensional algorithm
            Eigen::VectorXd delta_old = par - mean_;
            mean_ += delta_old / static_cast<double>(nSteps_);
            Eigen::VectorXd delta_new = par - mean_;
 
            cov_ += delta_old * delta_new.transpose();
        }
    }
 
    Eigen::MatrixXd getAdaptedCovariance() const {
        double epsilon = 1e-6; // Regularization
        Eigen::MatrixXd identity = Eigen::MatrixXd::Identity(dim_, dim_);
        return scale_ * getCovariance() + scale_ * epsilon * identity;
    }
 
    void reset() {
        nSteps_ = 0;
        nAccepts_ = 0;
        mean_.setZero();
        cov_.setZero();
        scale_ = 5.6644 / static_cast<double>(dim_);
    }
 
    Eigen::VectorXd getCurrent() const {
        return proposal_.get();
    }
 
    double getScore() const {
        return score_;
    }
 
    size_t getDim() const {
        return dim_;
    }
 
    size_t getSteps() const {
        return nSteps_;
    }
 
    double getAcceptanceRatio() const {
        return static_cast<double>(nAccepts_) / static_cast<double>(nSteps_);
    }
 
    Eigen::VectorXd getMean() const {
        return mean_;
    }
 
    Eigen::MatrixXd getCovariance() const {
        if(nSteps_ <= 1) {
            return Eigen::MatrixXd::Zero(dim_, dim_);
        }
        return cov_ / (static_cast<double>(nSteps_) - 1.0);
    }
 
    void info() const {
        std::cout << "run " << nSteps_ << " steps\n"
                  << "acceptance ratio: " << std::fixed << std::setprecision(3) << getAcceptanceRatio() << "\n"
                  << "Data covariance: \n" << getCovariance() << "\n"
                  << "Data mean: \n" << getMean() << std::endl;
    }
};
 
// ==============================================================================
// CHAIN DIAGNOSTICS (Also templated)
// ==============================================================================
 
template <typename ProposalType>
double multiChainDiagnosis(const std::vector<MarkovChain<ProposalType>> &chains) {
    if(chains.empty()) return 0.0;
 
    size_t dim_chain = chains[0].getDim();
    size_t num_chains = chains.size();
 
    Eigen::VectorXd chainwise_mean = Eigen::VectorXd::Zero(dim_chain);
    Eigen::VectorXd chainwise_cov = Eigen::VectorXd::Zero(dim_chain);
    Eigen::VectorXd chainwise_mean_cov = Eigen::VectorXd::Zero(dim_chain);
 
    for(const auto &chain : chains) {
        Eigen::VectorXd current_mean = chain.getMean();
        Eigen::MatrixXd current_cov = chain.getCovariance();
 
        chainwise_mean += current_mean;
        chainwise_cov += current_cov.diagonal();
        chainwise_mean_cov += current_mean.cwiseProduct(current_mean);
    }
 
    chainwise_mean /= static_cast<double>(num_chains);
    chainwise_cov /= static_cast<double>(num_chains);
 
    chainwise_mean_cov = chainwise_mean_cov / static_cast<double>(num_chains) - chainwise_mean.cwiseProduct(chainwise_mean);
 
    if(num_chains > 1) {
        chainwise_mean_cov *= static_cast<double>(num_chains) / static_cast<double>(num_chains - 1);
    }
 
    Eigen::VectorXd var_chain = chainwise_cov + chainwise_mean_cov;
 
    // Prevent division by zero if variance is perfectly flat
    Eigen::VectorXd r_hat = Eigen::VectorXd::Zero(dim_chain);
    for(size_t i = 0; i < dim_chain; ++i) {
        if(chainwise_cov(i) > 0.0) {
            r_hat(i) = std::sqrt(var_chain(i) / chainwise_cov(i));
        } else {
            r_hat(i) = 1.0;
        }
    }
 
    return r_hat.maxCoeff();
}
 
 
class EvolutionaryMarkovChain {
  protected:
    size_t nChains_{0};                                  
    size_t dim_{0};                                      
    std::vector<Eigen::VectorXd> chainSamples_;          
    std::vector<double> chainScores_;                    
 
    // Random distribution to pick a random chain
    std::uniform_int_distribution<size_t> distChain_;    
 
    // Random distribution to pick a random number
    std::uniform_real_distribution<double> distU_{0, 1}; 
    std::normal_distribution<double> distN_{0, 1};       
 
    double nugget_{1e-6};                               
 
  public:
    EvolutionaryMarkovChain() = default;
 
 
    EvolutionaryMarkovChain(std::vector<Eigen::VectorXd> initialSamples, std::vector<double> initialScores, double nugget = 1e-6);
 
    void step(const score_t &getScore, std::default_random_engine &rng, double gamma = 0.2);
 
    std::vector<Eigen::VectorXd> getCurrent() const {
        return chainSamples_;
    }
 
    std::vector<double> getScores() const {
        return chainScores_;
    }
};
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
}
 
#endif