replace pp_dirichlet_process_{rv,base}.py

Author: dustinvtran

examples/pp_dirichlet_process_base.py

@@ -2,7 +2,8 @@
 """Dirichlet process.
 
 We sample from a Dirichlet process (with inputted base distribution)
-via its stick breaking construction.
+via its stick breaking construction. For details on the
+implementation, see the source for ``DirichletProcess``.
 
 References
 ----------
@@ -14,69 +15,24 @@
 
 import tensorflow as tf
 
-from edward.models import Bernoulli, Beta, Normal
+from edward.models import DirichletProcess, Exponential, Normal
 
+base_cls = Normal
+kwargs = {'mu': 0.0, 'sigma': 1.0}
+dp = DirichletProcess(0.1, base_cls, **kwargs)
+print(dp)
 
-def dirichlet_process(alpha, base_cls, sample_n=50, *args, **kwargs):
-  """Dirichlet process DP(``alpha``, ``base_cls(*args, **kwargs)``).
+# ``theta`` is the distribution indirectly returned by the DP.
+theta = base_cls(value=tf.cast(dp, tf.float32), **kwargs)
+print(theta)
 
-  Only works for scalar alpha and scalar base distribution.
-
-  Parameters
-  ----------
-  alpha : tf.Tensor
-    Concentration parameter. Its shape determines the batch shape of the DP.
-  base_cls : RandomVariable
-    Class of base distribution. Its shape (when instantiated)
-    determines the event shape of the DP.
-  sample_n : int, optional
-    Number of samples for each DP in the batch shape.
-  *args, **kwargs : optional
-    Arguments passed into ``base_cls``.
-
-  Returns
-  -------
-  tf.Tensor
-    A ``tf.Tensor`` of shape ``[sample_n] + batch_shape + event_shape``,
-    where ``sample_n`` is the number of samples for each DP,
-    ``batch_shape`` is the number of independent DPs, and
-    ``event_shape`` is the shape of the base distribution.
-  """
-  def cond(k, beta_k, draws, bools):
-    # Proceed if at least one bool is True.
-    return tf.reduce_any(bools)
-
-  def body(k, beta_k, draws, bools):
-    k = k + 1
-    beta_k = beta_k * Beta(a=1.0, b=alpha).sample(sample_n)
-    theta_k = base_cls(*args, **kwargs).sample(sample_n)
-
-    # Assign True samples to the new theta_k.
-    draws = tf.where(bools, theta_k, draws)
-
-    flips = tf.cast(Bernoulli(p=beta_k), tf.bool)
-    bools = tf.logical_and(flips, tf.equal(draws, theta_k))
-    return k, beta_k, draws, bools
-
-  k = 0
-  beta_k = Beta(a=1.0, b=alpha).sample(sample_n)
-  theta_k = base_cls(*args, **kwargs).sample(sample_n)
-
-  # Initialize all samples as theta_k.
-  draws = theta_k
-  # Flip ``sample_n`` coins, one for each sample.
-  flips = tf.cast(Bernoulli(p=beta_k), tf.bool)
-  # Get boolean tensor for samples that return heads
-  # and are currently equal to theta_k.
-  bools = tf.logical_and(flips, tf.equal(draws, theta_k))
-
-  total_sticks, _, samples, _ = tf.while_loop(
-      cond, body, loop_vars=[k, beta_k, draws, bools])
-  return total_sticks, samples
-
-
-dp = dirichlet_process(0.1, Normal, mu=0.0, sigma=1.0)
+# Fetching theta is the same as fetching the Dirichlet process.
 sess = tf.Session()
-print(sess.run(dp))
-print(sess.run(dp))
-print(sess.run(dp))
+print(sess.run([dp, theta]))
+print(sess.run([dp, theta]))
+
+# This also works for non-scalar base distributions.
+base_cls = Exponential
+kwargs = {'lam': tf.ones([5, 2])}
+dp = DirichletProcess(0.1, base_cls, **kwargs)
+print(dp)