Stan  1.3
probability, sampling & optimization
 All Classes Namespaces Files Functions Variables Typedefs Enumerator Friends Macros Pages
nuts_nondiag.hpp
Go to the documentation of this file.
1 #ifndef __STAN__MCMC__NUTS_NONDIAG_H__
2 #define __STAN__MCMC__NUTS_NONDIAG_H__
3 
4 #include <ctime>
5 #include <cstddef>
6 #include <iostream>
7 #include <vector>
8 
9 #include <boost/random/normal_distribution.hpp>
10 #include <boost/random/mersenne_twister.hpp>
11 #include <boost/random/variate_generator.hpp>
12 #include <boost/random/uniform_01.hpp>
13 
14 #include <stan/mcmc/adaptive_sampler.hpp>
15 #include <stan/mcmc/dualaverage.hpp>
16 #include <stan/mcmc/hmc_base.hpp>
17 #include <stan/mcmc/util.hpp>
18 
19 #include <Eigen/Dense>
20 
21 namespace stan {
22 
23  namespace mcmc {
24 
33  template <class BaseRNG = boost::mt19937>
34  class nuts_nondiag : public hmc_base<BaseRNG> {
35  private:
36 
37  // Stop immediately if H < u - _maxchange
38  const double _maxchange;
39 
40  // Limit tree depth
41  const int _maxdepth;
42 
43  // Depth of last sample taken (-1 before any samples)
44  int _lastdepth;
45 
46  int _drop_warm;
47 
48  // Yuanjun : now we use a matrix to get the covariance matrix
49  Eigen::MatrixXd _cov_mat;
50  // Yuanjun : The cholesky decomposition of _cov_mat
51  Eigen::MatrixXd _cov_L;
52  // Running statistics to estimate per-coordinate std. deviations.
53  Eigen::MatrixXd _x_sum;
54  Eigen::MatrixXd _xsq_sum;
55 
56  double _x_sum_n;
57  // Next time we should adapt the per-parameter step sizes.
58  int _next_diag_adapt;
59 
60 
68  inline bool compute_criterion(std::vector<double>& xplus,
69  std::vector<double>& xminus,
70  std::vector<double>& mplus,
71  std::vector<double>& mminus) {
72  std::vector<double> total_direction;
73  stan::math::sub(xplus, xminus, total_direction);
74  Eigen::Map<Eigen::VectorXd> total_direction_mat(&total_direction[0],total_direction.size());
75  Eigen::Map<Eigen::VectorXd> mplus_mat(&mplus[0],mplus.size());
76  Eigen::Map<Eigen::VectorXd> mminus_mat(&mminus[0],mminus.size());
77  _cov_L.triangularView<Eigen::Lower>().solveInPlace(total_direction_mat);
78  return total_direction_mat.dot(mplus_mat) > 0
79  && total_direction_mat.dot(mminus_mat) > 0;
80  }
81 
82  public:
83 
110  nuts_nondiag(stan::model::prob_grad& model,
111  const std::vector<double>& params_r,
112  const std::vector<int>& params_i,
113  int maxdepth = 10,
114  double epsilon = -1,
115  double epsilon_pm = 0.0,
116  bool epsilon_adapt = true,
117  double delta = 0.6,
118  double gamma = 0.05,
119  BaseRNG base_rng = BaseRNG(std::time(0)))
120  : hmc_base<BaseRNG>(model,
121  params_r,
122  params_i,
123  epsilon,
124  epsilon_pm,
125  epsilon_adapt,
126  delta,
127  gamma,
128  base_rng),
129  _maxchange(-1000),
130  _maxdepth(maxdepth),
131  _lastdepth(-1),
132  _drop_warm(10),
133  _x_sum_n(0),
134  _next_diag_adapt(10 + model.num_params_r())
135  {
136  // start at 10 * epsilon because NUTS cheaper for larger epsilon
137  this->adaptation_init(10.0);
138  _x_sum = Eigen::MatrixXd::Zero(model.num_params_r(), 1);
139  _xsq_sum = Eigen::MatrixXd::Zero(model.num_params_r(),model.num_params_r());
140  _cov_mat = Eigen::MatrixXd::Identity(model.num_params_r(),model.num_params_r());
141  //use the same initialization method as diagonal one
142  _cov_L = Eigen::MatrixXd::Identity(model.num_params_r(),model.num_params_r());
143  }
144 
150  ~nuts_nondiag() { }
151 
157  virtual sample next_impl() {
158  // Initialize the algorithm
159  std::vector<double> mminus(this->_model.num_params_r());
160  for (size_t i = 0; i < mminus.size(); ++i)
161  mminus[i] = this->_rand_unit_norm();
162  std::vector<double> mplus(mminus);
163  // The log-joint probability of the momentum and position terms, i.e.
164  // -(kinetic energy + potential energy)
165  double H0 = -0.5 * stan::math::dot_self(mminus) + this->_logp;
166 
167  std::vector<double> gradminus(this->_g);
168  std::vector<double> gradplus(this->_g);
169  std::vector<double> xminus(this->_x);
170  std::vector<double> xplus(this->_x);
171 
172  // Sample the slice variable
173  double u = log(this->_rand_uniform_01()) + H0;
174  int nvalid = 1;
175  int direction = 2 * (this->_rand_uniform_01() > 0.5) - 1;
176  bool criterion = true;
177 
178  // Repeatedly double the set of points we've visited
179  std::vector<double> newx, newgrad, dummy1, dummy2, dummy3;
180  double newlogp = -1;
181  double prob_sum = -1;
182  int newnvalid = -1;
183  int n_considered = 0;
184  // for-loop with depth outside to set lastdepth
185  int depth = 0;
186 
187  double epsilon = this->_epsilon;
188  // only vary epsilon after done adapting
189  if (!this->adapting() && this->varying_epsilon()) {
190  double low = epsilon * (1.0 - this->_epsilon_pm);
191  double high = epsilon * (1.0 + this->_epsilon_pm);
192  double range = high - low;
193  epsilon = low + (range * this->_rand_uniform_01());
194  }
195  this->_epsilon_last = epsilon; // use epsilon_last in tree build
196 
197  while (criterion && (_maxdepth < 0 || depth <= _maxdepth)) {
198  direction = 2 * (this->_rand_uniform_01() > 0.5) - 1;
199  if (direction == -1)
200  build_tree(xminus, mminus, gradminus, u, direction, depth,
201  H0, xminus, mminus, gradminus, dummy1, dummy2, dummy3,
202  newx, newgrad, newlogp, newnvalid, criterion, prob_sum,
203  n_considered);
204  else
205  build_tree(xplus, mplus, gradplus, u, direction, depth,
206  H0, dummy1, dummy2, dummy3, xplus, mplus, gradplus,
207  newx, newgrad, newlogp, newnvalid, criterion, prob_sum,
208  n_considered);
209  // We can't look at the results of this last doubling if criterion==false
210  if (!criterion)
211  break;
212  criterion = compute_criterion(xplus, xminus, mplus, mminus);
213  // Metropolis-Hastings to determine if we can jump to a point in
214  // the new half-tree
215  if (this->_rand_uniform_01() < float(newnvalid) / (1e-100+float(nvalid))) {
216  this->_x = newx;
217  this->_g = newgrad;
218  this->_logp = newlogp;
219  }
220  if (this->_n_adapt_steps > _drop_warm) {
221  _x_sum_n++;
222 
223  for (size_t i = 0; i < newx.size(); i++) {
224  _x_sum(i) += this->_x[i];
225  for (size_t j = 0; j < newx.size(); j++){
226  _xsq_sum(i,j) += this->_x[i] * this->_x[j];
227  }
228  }
229  }
230  nvalid += newnvalid;
231  ++depth;
232  }
233  _lastdepth = depth;
234 
235  // Now we just have to update global (epsilon) and local
236  // (step_sizes) step sizes, if adaptation is on.
237  double adapt_stat = prob_sum / float(n_considered);
238  if (this->adapting()) {
239  // epsilon.
240  double adapt_g = adapt_stat - this->_delta;
241  std::vector<double> gvec(1, -adapt_g);
242  std::vector<double> result;
243  this->_da.update(gvec, result);
244  this->_epsilon = exp(result[0]);
245 
246  if (this->_n_adapt_steps == _next_diag_adapt) {
247  _next_diag_adapt *= 2;
248  //double step_size_sq_sum = 0;
249 
250  _x_sum = _x_sum / _x_sum_n; //Yuanjun : change summation to expectation
251  _xsq_sum = _xsq_sum / _x_sum_n; //Yuanjun : change summation to expectation
252  _cov_mat = _xsq_sum - _x_sum * _x_sum.transpose();
253  double _norm = _cov_mat.trace() / _cov_mat.rows();
254  if(_norm == 0)
255  _cov_mat = Eigen::MatrixXd::Identity(_cov_mat.rows(),_cov_mat.cols());
256  else{
257  _cov_mat = _cov_mat * (_x_sum_n / _norm / (_x_sum_n + 5)) ;
258  for(int i = 0; i < _cov_mat.rows(); i++)
259  _cov_mat(i,i) = _cov_mat(i,i) * ((_x_sum_n + 2) / _x_sum_n) + (3 / (_x_sum_n + 5));
260  }
261  _cov_L = _cov_mat.selfadjointView<Eigen::Upper>().llt().matrixL();
262  _x_sum *= 0;
263  _xsq_sum *= 0;
264  _x_sum_n *= 0;
265  }
266  }
267  std::vector<double> result;
268  this->_da.xbar(result);
269  double avg_eta = 1.0 / this->n_steps();
270  this->update_mean_stat(avg_eta,adapt_stat);
271 
272  return mcmc::sample(this->_x, this->_z, this->_logp);
273  }
274 
275  Eigen::Matrix<double,Eigen::Dynamic,Eigen::Dynamic> get_mass_matrix() {
276  return _cov_mat;
277  }
278 
279  virtual void write_sampler_param_names(std::ostream& o) {
280  o << "treedepth__,";
281  if (this->_epsilon_adapt || this->varying_epsilon())
282  o << "stepsize__,";
283  }
284 
285  virtual void write_sampler_params(std::ostream& o) {
286  o << _lastdepth << ',';
287  if (this->_epsilon_adapt || this->varying_epsilon())
288  o << this->_epsilon_last << ',';
289  }
290 
291  virtual void write_adaptation_params(std::ostream& o) {
292  o << "# (mcmc::nuts_nondiag) adaptation finished" << '\n';
293  o << "# step size=" << this->_epsilon << '\n';
294  o << "# estimated covariance matrix:\n"; // FIXME: names/delineation requires access to model
295  for(int i=0; i<_cov_mat.rows(); i++){
296  o << "#";
297  for(int j=0; j<_cov_mat.cols(); j++)
298  o << _cov_mat(i,j) << ",";
299  o << std::endl;
300  }
301  }
302 
303  virtual void get_sampler_param_names(std::vector<std::string>& names) {
304  names.clear();
305  names.push_back("treedepth__");
306  if (this->_epsilon_adapt || this->varying_epsilon())
307  names.push_back("stepsize__");
308  }
309 
310  virtual void get_sampler_params(std::vector<double>& values) {
311  values.clear();
312  values.push_back(_lastdepth);
313  if (this->_epsilon_adapt || this->varying_epsilon())
314  values.push_back(this->_epsilon_last);
315  }
316 
351  void build_tree(const std::vector<double>& x,
352  const std::vector<double>& m,
353  const std::vector<double>& grad,
354  double u,
355  int direction,
356  int depth,
357  double H0,
358  std::vector<double>& xminus,
359  std::vector<double>& mminus,
360  std::vector<double>& gradminus,
361  std::vector<double>& xplus,
362  std::vector<double>& mplus,
363  std::vector<double>& gradplus,
364  std::vector<double>& newx,
365  std::vector<double>& newgrad,
366  double& newlogp,
367  int& nvalid,
368  bool& criterion,
369  double& prob_sum,
370  int& n_considered) {
371  if (depth == 0) { // base case
372  xminus = x;
373  gradminus = grad;
374  mminus = m;
375  newlogp = nondiag_leapfrog(this->_model, this->_z, _cov_L,
376  xminus, mminus, gradminus,
377  direction * this->_epsilon_last,
378  this->_error_msgs, this->_output_msgs);
379  newx = xminus;
380  newgrad = gradminus;
381  xplus = xminus;
382  mplus = mminus;
383  gradplus = gradminus;
384  double newH = newlogp - 0.5 * stan::math::dot_self(mminus);
385  if (newH != newH) // treat nan as -inf
386  newH = -std::numeric_limits<double>::infinity();
387  nvalid = newH > u;
388  criterion = newH - u > _maxchange;
389  prob_sum = stan::math::min(1, exp(newH - H0));
390  n_considered = 1;
391  this->nfevals_plus_eq(1);
392  // Update running statistics if point is in slice
393  /*if (nvalid & this->_n_adapt_steps > _drop_warm) {
394  _x_sum_n++;
395 
396  for (size_t i = 0; i < newx.size(); i++) {
397  _x_sum(i) += newx[i];
398  for (size_t j = 0; j < newx.size(); j++){
399  _xsq_sum(i,j) += newx[i] * newx[j];
400  }
401  }
402  }*/
403  } else { // depth >= 1
404  build_tree(x, m, grad, u, direction, depth-1, H0, xminus, mminus,
405  gradminus, xplus, mplus, gradplus, newx, newgrad, newlogp,
406  nvalid, criterion, prob_sum, n_considered);
407  if (criterion) {
408  std::vector<double> dummy1, dummy2, dummy3;
409  std::vector<double> newx2;
410  std::vector<double> newgrad2;
411  double newlogp2;
412  int nvalid2;
413  bool criterion2;
414  double prob_sum2;
415  int n_considered2;
416  if (direction == -1)
417  build_tree(xminus, mminus, gradminus, u, direction, depth-1, H0,
418  xminus, mminus, gradminus, dummy1, dummy2, dummy3,
419  newx2, newgrad2, newlogp2, nvalid2, criterion2,
420  prob_sum2, n_considered2);
421  else
422  build_tree(xplus, mplus, gradplus, u, direction, depth-1, H0,
423  dummy1, dummy2, dummy3, xplus, mplus, gradplus,
424  newx2, newgrad2, newlogp2, nvalid2, criterion2,
425  prob_sum2, n_considered2);
426  if (criterion &&
427  (this->_rand_uniform_01()
428  < float(nvalid2) / float(nvalid+nvalid2))){
429  newx = newx2;
430  newgrad = newgrad2;
431  newlogp = newlogp2;
432  }
433  n_considered += n_considered2;
434  prob_sum += prob_sum2;
435  criterion &= criterion2;
436  nvalid += nvalid2;
437  }
438  criterion &= compute_criterion(xplus, xminus, mplus, mminus);
439  }
440  }
441 
442 
443  };
444 
445 
446  }
447 
448 }
449 
450 #endif
     [ Stan Home Page ] © 2011–2013, Stan Development Team.