nuts_massgiven.hpp

Go to the documentation of this file.

#ifndef __STAN__MCMC__NUTS_MASSGIVEN_H__
#define __STAN__MCMC__NUTS_MASSGIVEN_H__

#include <ctime>
#include <cstddef>
#include <iostream>
#include <vector>

#include <boost/random/normal_distribution.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/random/uniform_01.hpp>

#include <stan/mcmc/adaptive_sampler.hpp>
#include <stan/mcmc/dualaverage.hpp>
#include <stan/mcmc/hmc_base.hpp>
#include <stan/mcmc/util.hpp>

#include <Eigen/Dense>

namespace stan {

  namespace mcmc {

    template <class BaseRNG = boost::mt19937>
    class nuts_massgiven : public hmc_base<BaseRNG> {
    private:

      // Stop immediately if H < u - _maxchange
      const double _maxchange;

      // Limit tree depth
      const int _maxdepth;

      // Depth of last sample taken (-1 before any samples)
      int _lastdepth;

      // Yuanjun : now we use a matrix to get the covariance matrix
      // Yuanjun : The cholesky decomposition of _cov_mat
      Eigen::MatrixXd _cov_L;
      // Running statistics to estimate per-coordinate std. deviations.

        
        //Eigen::Map<Eigen::VectorXd> _x_mat;
        //Eigen::Map<Eigen::VectorXd> _m_mat;
        //Eigen::Map<Eigen::VectorXd> _g_mat;
      
      // Next time we should adapt the per-parameter step sizes.
        

        inline bool compute_criterion(std::vector<double>& xplus,
                                      std::vector<double>& xminus,
                                      std::vector<double>& mplus,
                                      std::vector<double>& mminus) {
            std::vector<double> total_direction;
            stan::math::sub(xplus, xminus, total_direction);
            Eigen::Map<Eigen::VectorXd> total_direction_mat(&total_direction[0],total_direction.size());
            Eigen::Map<Eigen::VectorXd> mplus_mat(&mplus[0],mplus.size());
            Eigen::Map<Eigen::VectorXd> mminus_mat(&mminus[0],mminus.size());
            _cov_L.triangularView<Eigen::Lower>().solveInPlace(total_direction_mat);
            return total_direction_mat.dot(mplus_mat) > 0
            && total_direction_mat.dot(mminus_mat) > 0;
        }

    public:

      nuts_massgiven(stan::model::prob_grad& model,
         const std::vector<double>& params_r,
         const std::vector<int>& params_i,
         std::string cov_file,
         int maxdepth = 10,
         double epsilon = -1,
         double epsilon_pm = 0.0,
         bool epsilon_adapt = true,
         double delta = 0.6,
         double gamma = 0.05,
         BaseRNG base_rng = BaseRNG(std::time(0))) 
        : hmc_base<BaseRNG>(model,
                            params_r,
                            params_i,
                            epsilon,
                            epsilon_pm,
                            epsilon_adapt,
                            delta,
                            gamma,
                            base_rng),
          
          _maxchange(-1000),
          _maxdepth(maxdepth),
          _lastdepth(-1)
          //_x_mat(&(this->_x[0]), (this->_x).size())
      {
        // start at 10 * epsilon because NUTS cheaper for larger epsilon
        this->adaptation_init(10.0);
  
        _cov_L = Eigen::MatrixXd::Identity(model.num_params_r(),model.num_params_r());
        read_cov(cov_file, _cov_L);
         // std::cout << _cov_L << "baby" << std::endl;
        
        //_x_mat = Eigen::Map<Eigen::VectorXd>(this->_x[0], this->_x.size());
        //_m_mat = Eigen::Map<Eigen::VectorXd>(this->_m[0], this->_m.size());
        //_g_mat = Eigen::Map<Eigen::VectorXd>(this->_g[0], this->_g.size());
      }

      ~nuts_massgiven() { }

      virtual sample next_impl() {
        // Initialize the algorithm
        std::vector<double> mminus(this->_model.num_params_r());
        for (size_t i = 0; i < mminus.size(); ++i)
          mminus[i] = this->_rand_unit_norm();
        std::vector<double> mplus(mminus);
        // The log-joint probability of the momentum and position terms, i.e.
        // -(kinetic energy + potential energy)
        double H0 = -0.5 * stan::math::dot_self(mminus) + this->_logp;

        std::vector<double> gradminus(this->_g);
        std::vector<double> gradplus(this->_g);
        std::vector<double> xminus(this->_x);
        std::vector<double> xplus(this->_x);

        // Sample the slice variable
        double u = log(this->_rand_uniform_01()) + H0;
        int nvalid = 1;
        int direction = 2 * (this->_rand_uniform_01() > 0.5) - 1;
        bool criterion = true;

        // Repeatedly double the set of points we've visited
        std::vector<double> newx, newgrad, dummy1, dummy2, dummy3;
        double newlogp = -1;
        double prob_sum = -1;
        int newnvalid = -1;
        int n_considered = 0;
        // for-loop with depth outside to set lastdepth
        int depth = 0;

        double epsilon = this->_epsilon;
        // only vary epsilon after done adapting
        if (!this->adapting() && this->varying_epsilon()) { 
          double low = epsilon * (1.0 - this->_epsilon_pm);
          double high = epsilon * (1.0 + this->_epsilon_pm);
          double range = high - low;
          epsilon = low + (range * this->_rand_uniform_01());
        }
        this->_epsilon_last = epsilon; // use epsilon_last in tree build

        while (criterion && (_maxdepth < 0 || depth <= _maxdepth)) {
          direction = 2 * (this->_rand_uniform_01() > 0.5) - 1;
          if (direction == -1)
            build_tree(xminus, mminus, gradminus, u, direction, depth,
                       H0, xminus, mminus, gradminus, dummy1, dummy2, dummy3,
                       newx, newgrad, newlogp, newnvalid, criterion, prob_sum,
                       n_considered);
          else
            build_tree(xplus, mplus, gradplus, u, direction, depth,
                       H0, dummy1, dummy2, dummy3, xplus, mplus, gradplus, 
                       newx, newgrad, newlogp, newnvalid, criterion, prob_sum,
                       n_considered);
          // We can't look at the results of this last doubling if criterion==false
          if (!criterion)
            break;
          criterion = compute_criterion(xplus, xminus, mplus, mminus);
          // Metropolis-Hastings to determine if we can jump to a point in
          // the new half-tree
          if (this->_rand_uniform_01() < float(newnvalid) / (1e-100+float(nvalid))) {
            this->_x = newx;
            this->_g = newgrad;
            this->_logp = newlogp;
          }
          nvalid += newnvalid;
          ++depth;
        }
        _lastdepth = depth;

        // Now we just have to update global (epsilon) and local
        // (step_sizes) step sizes, if adaptation is on.
        double adapt_stat = prob_sum / float(n_considered);
        if (this->adapting()) { 
          // epsilon.
          double adapt_g = adapt_stat - this->_delta;
          std::vector<double> gvec(1, -adapt_g);
          std::vector<double> result;
          this->_da.update(gvec, result);
          this->_epsilon = exp(result[0]);
          // step_sizes. Doesn't happen every step.
          
        }
        std::vector<double> result;
        this->_da.xbar(result);
        double avg_eta = 1.0 / this->n_steps();
        this->update_mean_stat(avg_eta,adapt_stat);

        return mcmc::sample(this->_x, this->_z, this->_logp);
      }

      virtual void write_sampler_param_names(std::ostream& o) {
        o << "treedepth__,";
        if (this->_epsilon_adapt || this->varying_epsilon())
          o << "stepsize__,";
      }

      virtual void write_sampler_params(std::ostream& o) {
        o << _lastdepth << ',';
        if (this->_epsilon_adapt || this->varying_epsilon())
          o << this->_epsilon_last << ',';
      }

      virtual void write_adaptation_params(std::ostream& o) {
        o << "# (mcmc::nuts_massgiven) adaptation finished" << '\n';
        o << "# step size=" << this->_epsilon << '\n';
        o << "# Preset covariance matrix:\n"; // FIXME:  names/delineation requires access to model
        Eigen::MatrixXd _cov_mat = _cov_L * _cov_L.transpose();
        for(int i=0; i<_cov_mat.rows(); i++){
          o << "#";
          for(int j=0; j<_cov_mat.cols(); j++)
            o << _cov_mat(i,j) << ",";
          o << std::endl;
        }
      }

      virtual void get_sampler_param_names(std::vector<std::string>& names) {
        names.clear();
        names.push_back("treedepth__");
        if (this->_epsilon_adapt || this->varying_epsilon())
          names.push_back("stepsize__");
      }

      virtual void get_sampler_params(std::vector<double>& values) {
        values.clear();
        values.push_back(_lastdepth);
        if (this->_epsilon_adapt || this->varying_epsilon())
          values.push_back(this->_epsilon_last);
      }

      void build_tree(const std::vector<double>& x, 
                      const std::vector<double>& m,
                      const std::vector<double>& grad,
                      double u,
                      int direction,
                      int depth, 
                      double H0,
                      std::vector<double>& xminus,
                      std::vector<double>& mminus,
                      std::vector<double>& gradminus,
                      std::vector<double>& xplus,
                      std::vector<double>& mplus,
                      std::vector<double>& gradplus,
                      std::vector<double>& newx,
                      std::vector<double>& newgrad,
                      double& newlogp,
                      int& nvalid, 
                      bool& criterion, 
                      double& prob_sum, 
                      int& n_considered) {
        if (depth == 0) {   // base case
          xminus = x;
          gradminus = grad;
          mminus = m;
          newlogp = nondiag_leapfrog(this->_model, this->_z, _cov_L,  //Yuanjun : implement a new leapfrog function
                                      xminus, mminus, gradminus,
                                      direction * this->_epsilon_last,
                                      this->_error_msgs, this->_output_msgs);
          newx = xminus;
          newgrad = gradminus;
          xplus = xminus;
          mplus = mminus;
          gradplus = gradminus;
          double newH = newlogp - 0.5 * stan::math::dot_self(mminus);   //Yuanjun : No longer need to calculate H with a matrix
          if (newH != newH) // treat nan as -inf
            newH = -std::numeric_limits<double>::infinity();
          nvalid = newH > u;
          criterion = newH - u > _maxchange;
          prob_sum = stan::math::min(1, exp(newH - H0));
          n_considered = 1;
          this->nfevals_plus_eq(1);
          // Update running statistics if point is in slice

        } else {            // depth >= 1
          build_tree(x, m, grad, u, direction, depth-1, H0, xminus, mminus,
                     gradminus, xplus, mplus, gradplus, newx, newgrad, newlogp,
                     nvalid, criterion, prob_sum, n_considered);
          if (criterion) {
            std::vector<double> dummy1, dummy2, dummy3;
            std::vector<double> newx2;
            std::vector<double> newgrad2;
            double newlogp2;
            int nvalid2;
            bool criterion2;
            double prob_sum2;
            int n_considered2;
            if (direction == -1)
              build_tree(xminus, mminus, gradminus, u, direction, depth-1, H0,
                         xminus, mminus, gradminus, dummy1, dummy2, dummy3,
                         newx2, newgrad2, newlogp2, nvalid2, criterion2, 
                         prob_sum2, n_considered2);
            else
              build_tree(xplus, mplus, gradplus, u, direction, depth-1, H0,
                         dummy1, dummy2, dummy3, xplus, mplus, gradplus, 
                         newx2, newgrad2, newlogp2, nvalid2, criterion2,
                         prob_sum2, n_considered2);
            if (criterion && 
                (this->_rand_uniform_01() 
                 < float(nvalid2) / float(nvalid+nvalid2))){
              newx = newx2;
              newgrad = newgrad2;
              newlogp = newlogp2;
            }
            n_considered += n_considered2;
            prob_sum += prob_sum2;
            criterion &= criterion2;
            nvalid += nvalid2;
          }
          criterion &= compute_criterion(xplus, xminus, mplus, mminus);
        }
      }


    };


  }

}

#endif