I've just red the great 2012 blog post of Edwin Chen about Dirichlet Process with companion code in R and Ruby. Then I'm trying to translate the Stick-Breaking Process from R to Python.

I've got this small piece of code. Is-it good?

# A direct translation of the R code
import numpy as np
def Stick_Breaking(num_weights,alpha):
    betas = np.random.beta(1,alpha, size=num_weights)
    remaining_stick_lengths =[1]+list(np.cumprod(1-betas))[0:num_weights-1]
    return remaining_stick_lengths * betas

# A nice correction suggested by Tomáš Tunys
def Stick_Breaking(num_weights,alpha):
    betas = np.random.beta(1,alpha, size=num_weights) 
    betas[1:] *= np.cumprod(1 - betas[:-1])       
    return betas

In order to show small distribution histograms:

import matplotlib.pyplot as plt

for _ in range(5):
    num_weights = 10
    alpha = 1
    weights = Stick_Breaking(num_weights,alpha)
    plt.axis([0, num_weights+1, 0, 1])

enter image description here


2 Answers 2


Correct me if I am wrong, but shouldn't betas sum to 1.0 to make a valid distribution?

I have changed your code snippet just to do that (mind that the last sample from beta distribution is not used):

def Stick_Breaking(num_weights, alpha):
    betas = np.random.beta(1, alpha, size=num_weights)
    betas[1:] *= np.cumprod(1 - betas[:-1])
    return betas

For example, one sample from Stick_Breaking(10, 1.0) give array([0.8698612 , 0.11382917, 0.01005555, 0.00410595, 0.00137072]) and np.sum(Stick_Breaking(10, 1.0)) gives values close to 1.0.

EDIT: I stand corrected that the outcomes is not a distribution for any finite num_weights, but the code should be correct.

  • $\begingroup$ Nice! Indeed, it should sum up to one. I was suspecting something was wrong, but you put your finger on it. I would correct it mention you! Many thanks. $\endgroup$ Commented Mar 9, 2019 at 19:09
  • $\begingroup$ That said with big alpha parameter (like alpha=10) the spreading effect is obvious but the sum up is under one. It's more complicated than I figure out at the beginning. $\endgroup$ Commented Mar 9, 2019 at 19:57
  • 1
    $\begingroup$ @ClaudeCOULOMBE True, I was wrong assuming the result should be a distribution. It actually is when num_weights tends to $inf$, but it has nothing to do with it. I should remind myself some of this stuff :). If you are interested into DP I highly recommend this talk by Michael Jordan on YT: youtube.com/watch?v=yfLoxwjCGNY. Have fun! ;) $\endgroup$ Commented Mar 10, 2019 at 11:58
  • $\begingroup$ @ Tomáš Tunys Thanks Tomáš, great talk by a great researcher! Thanks for the sharing. By the way, you help me a lot. I'm fascinated by all those China restaurant, Indian Buffet, Urn-Draw and Stick Breaking. I was wrong, you were imprecise, we're correct now. It's how the science and open source software progress and share. Keep contributing. $\endgroup$ Commented Mar 10, 2019 at 18:31

Let me also add that you can compute this without using the Numpy library. I am currently working on a journal that requires me to use stick-breaking as a prior for my autoencoder.

Method 1:

def compute_stick_segment(self, v):
    n_dims = v.size()[1]
    pi = torch.ones(size=v.size()).to(device)
    for idx in range(n_dims):
        product = 1
        for sub_idx in range(idx):
            product *=1-v[:,sub_idx]
        pi[:,idx] = v[:,idx] * product
    return pi

Method 2

from numpy.testing import assert_almost_equal

def set_v_K_to_one(self, v):
        # set Kth fraction v_i,K to one to ensure the stick segments sum to one
        if v.ndim > 2:
            v = v.squeeze() 
        v0 = v[:, -1].pow(0).reshape(v.shape[0], 1) 
        v1 = torch.cat([v[:, :latent_ndims - 1], v0], dim=1) 
        return v1.to(device)

    def get_stick_segments(self, v):
        n_samples = v.size()[0]
        n_dims = v.size()[1]
        pi = torch.zeros((n_samples, n_dims))

        for k in range(n_dims): 
            if k == 0:
                pi[:, k] = v[:, k]
                pi[:, k] = v[:, k] * torch.stack([(1 - v[:, j]) for j in range(n_dims) if j < k]).prod(axis=0)

        # ensure stick segments sum to 1
        assert_almost_equal(torch.ones(n_samples), pi.sum(axis=1).detach().numpy(),
                            decimal=2, err_msg='stick segments do not sum to 1')
        return pi.to(device)
z = self.set_v_K_to_one(v)
pi = self.get_stick_segments(z)

Any of the methods works fine.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.