1 files changed, 112 insertions, 7 deletions

diff --git a/gi/clda/src/ccrp.h b/gi/clda/src/ccrp.h
index eeccce1a..74d5be29 100644
--- a/gi/clda/src/ccrp.h
+++ b/gi/clda/src/ccrp.h
@@ -1,6 +1,7 @@
 #ifndef _CCRP_H_
 #define _CCRP_H_
 
+#include <numeric>
 #include <cassert>
 #include <cmath>
 #include <list>
@@ -9,6 +10,7 @@
 #include <tr1/unordered_map>
 #include <boost/functional/hash.hpp>
 #include "sampler.h"
+#include "slice_sampler.h"
 
 // Chinese restaurant process (Pitman-Yor parameters) with explicit table
 // tracking.
@@ -16,7 +18,36 @@
 template <typename Dish, typename DishHash = boost::hash<Dish> >
 class CCRP {
  public:
-  CCRP(double disc, double conc) : num_tables_(), num_customers_(), discount_(disc), concentration_(conc) {}
+  CCRP(double disc, double conc) :
+    num_tables_(),
+    num_customers_(),
+    discount_(disc),
+    concentration_(conc),
+    discount_prior_alpha_(std::numeric_limits<double>::quiet_NaN()),
+    discount_prior_beta_(std::numeric_limits<double>::quiet_NaN()),
+    concentration_prior_shape_(std::numeric_limits<double>::quiet_NaN()),
+    concentration_prior_rate_(std::numeric_limits<double>::quiet_NaN()) {}
+
+  CCRP(double d_alpha, double d_beta, double c_shape, double c_rate, double d = 0.1, double c = 10.0) :
+    num_tables_(),
+    num_customers_(),
+    discount_(d),
+    concentration_(c),
+    discount_prior_alpha_(d_alpha),
+    discount_prior_beta_(d_beta),
+    concentration_prior_shape_(c_shape),
+    concentration_prior_rate_(c_rate) {}
+
+  double discount() const { return discount_; }
+  double concentration() const { return concentration_; }
+
+  bool has_discount_prior() const {
+    return !std::isnan(discount_prior_alpha_);
+  }
+
+  bool has_concentration_prior() const {
+    return !std::isnan(concentration_prior_shape_);
+  }
 
   void clear() {
     num_tables_ = 0;
@@ -115,24 +146,90 @@ class CCRP {
     }
   }
 
+  double log_crp_prob() const {
+    return log_crp_prob(discount_, concentration_);
+  }
+
+  static double log_beta_density(const double& x, const double& alpha, const double& beta) {
+    assert(x > 0.0);
+    assert(x < 1.0);
+    assert(alpha > 0.0);
+    assert(beta > 0.0);
+    const double lp = (alpha-1)*log(x)+(beta-1)*log(1-x)+lgamma(alpha+beta)-lgamma(alpha)-lgamma(beta);
+    return lp;
+  }
+
+  static double log_gamma_density(const double& x, const double& shape, const double& rate) {
+    assert(x >= 0.0);
+    assert(shape > 0.0);
+    assert(rate > 0.0);
+    const double lp = (shape-1)*log(x) - shape*log(rate) - x/rate - lgamma(shape);
+    return lp;
+  }
+
   // taken from http://en.wikipedia.org/wiki/Chinese_restaurant_process
   // does not include P_0's
-  double log_crp_prob() const {
+  double log_crp_prob(const double& discount, const double& concentration) const {
     double lp = 0.0;
+    if (has_discount_prior())
+      lp = log_beta_density(discount, discount_prior_alpha_, discount_prior_beta_);
+    if (has_concentration_prior())
+      lp += log_gamma_density(concentration, concentration_prior_shape_, concentration_prior_rate_);
+    assert(lp <= 0.0);
     if (num_customers_) {
-      const double r = lgamma(1.0 - discount_);
-      lp = lgamma(concentration_) - lgamma(concentration_ + num_customers_)
-           + num_tables_ * discount_ + lgamma(concentration_ / discount_ + num_tables_)
-           - lgamma(concentration_ / discount_);
+      const double r = lgamma(1.0 - discount);
+      lp += lgamma(concentration) - lgamma(concentration + num_customers_)
+           + num_tables_ * discount + lgamma(concentration / discount + num_tables_)
+           - lgamma(concentration / discount);
+      assert(std::isfinite(lp));
       for (typename std::tr1::unordered_map<Dish, DishLocations, DishHash>::const_iterator it = dish_locs_.begin();
            it != dish_locs_.end(); ++it) {
         const DishLocations& cur = it->second;
-        lp += lgamma(cur.total_dish_count_ - discount_) - r;
+        for (std::list<unsigned>::const_iterator ti = cur.table_counts_.begin(); ti != cur.table_counts_.end(); ++ti) {
+          lp += lgamma(*ti - discount) - r;
+        }
       }
     }
+    assert(std::isfinite(lp));
     return lp;
   }
 
+  void resample_hyperparameters(MT19937* rng) {
+    assert(has_discount_prior() || has_concentration_prior());
+    DiscountResampler dr(*this);
+    ConcentrationResampler cr(*this);
+    const int niterations = 10;
+    double gamma_upper = std::numeric_limits<double>::infinity();
+    for (int iter = 0; iter < 5; ++iter) {
+      if (has_concentration_prior()) {
+        concentration_ = slice_sampler1d(cr, concentration_, *rng, 0.0,
+                               gamma_upper, 0.0, niterations, 100*niterations);
+      }
+      if (has_discount_prior()) {
+        discount_ = slice_sampler1d(dr, discount_, *rng, std::numeric_limits<double>::min(),
+                               1.0, 0.0, niterations, 100*niterations);
+      }
+    }
+    concentration_ = slice_sampler1d(cr, concentration_, *rng, 0.0,
+                             gamma_upper, 0.0, niterations, 100*niterations);
+  }
+
+  struct DiscountResampler {
+    DiscountResampler(const CCRP& crp) : crp_(crp) {}
+    const CCRP& crp_;
+    double operator()(const double& proposed_discount) const {
+      return crp_.log_crp_prob(proposed_discount, crp_.concentration_);
+    }
+  };
+
+  struct ConcentrationResampler {
+    ConcentrationResampler(const CCRP& crp) : crp_(crp) {}
+    const CCRP& crp_;
+    double operator()(const double& proposed_concentration) const {
+      return crp_.log_crp_prob(crp_.discount_, proposed_concentration);
+    }
+  };
+
   struct DishLocations {
     DishLocations() : total_dish_count_() {}
     unsigned total_dish_count_;        // customers at all tables with this dish
@@ -166,6 +263,14 @@ class CCRP {
 
   double discount_;
   double concentration_;
+
+  // optional beta prior on discount_ (NaN if no prior)
+  double discount_prior_alpha_;
+  double discount_prior_beta_;
+
+  // optional gamma prior on concentration_ (NaN if no prior)
+  double concentration_prior_shape_;
+  double concentration_prior_rate_;
 };
 
 template <typename T,typename H>