# Coding a simple Stick-Breaking Process in Python

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])
plt.bar(range(1,num_weights+1),weights)
plt.show()
print(sum(weights))


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.

• 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. Commented Mar 9, 2019 at 19:09
• 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. Commented Mar 9, 2019 at 19:57
• @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! ;) Commented Mar 10, 2019 at 11:58
• @ 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. 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]
else:
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)
#Usage:
z = self.set_v_K_to_one(v)
pi = self.get_stick_segments(z)


Any of the methods works fine.