bnp

commit ca36a3f2f866e4876afe3dbe029776a2edb9c6f2
parent 2642db77d9ff7cb7cbabaaf4430f9b93824634a9
Author: Jared Tobin <jared@jtobin.ca>
Date:   Mon, 15 Feb 2016 21:48:56 +1300

Add Pitman-Yor style CRP.

Diffstat:
A
chinese-restaurant-process/src/generalized_crp.r
 | 
29
+++++++++++++++++++++++++++++

1 file changed, 29 insertions(+), 0 deletions(-)
diff --git a/chinese-restaurant-process/src/generalized_crp.r b/chinese-restaurant-process/src/generalized_crp.r
@@ -0,0 +1,29 @@
+generalized_crp = function(n, a, b) {
+  restaurant = data.frame(table = 1, customers = 1)
+  for (j in seq(n - 1)) {
+    restaurant = arrival(restaurant, a, b)
+    }
+  restaurant
+  }
+
+arrival = function(r, a, b) {
+  k = nrow(r)
+  p = 1 - (b + k * a) / (sum(r$customers) + b)
+  if (rbinom(1, 1, p)) {
+    join_table(r, a)
+    } else {
+    start_table(r)
+    }
+  }
+
+join_table = function(r, a) {
+  probs = r$customers / sum(r$customers)
+  table = sample(1:nrow(r), size = 1, prob = probs)
+  r[table, 'customers'] = r[table, 'customers'] + 1
+  r
+  }
+
+start_table = function(r) {
+  rbind(r, c(nrow(r) + 1, 1))
+  }
+