Dirichlet-Process & Pitman-Yor Mixtures¶

What a random measure is — plus bivariate density estimation and clustering¶

The variable-selection arc already fit a univariate Dirichlet-process mixture to the galaxy data to ask "how many components?". This notebook takes the modelling view instead. Three things the earlier notebook did not do:

  1. What a Dirichlet process actually is — a prior over discrete random probability measures, constructed by stick-breaking;
  2. a bivariate DP mixture for 2-D density estimation and clustering (Old Faithful), with a Normal-Inverse-Wishart base and a multivariate-$t$ predictive;
  3. the two-parameter Pitman-Yor generalisation, whose clusters follow a power law.

We also show the frequentist counterpart at each step — kernel density estimation for the density, EM/mclust for the mixture — so the Bayesian nonparametric method sits beside its classical analog.

The model. A mixture with an unknown number of components: $$y_i\mid z_i=k \sim \mathcal N(\mu_k,\Sigma_k),\quad (\mu_k,\Sigma_k)\sim G_0=\text{NIW}(m_0,\kappa_0,\nu_0,\Psi_0),\quad z\sim\text{CRP}(\alpha,d).$$ The Normal-Inverse-Wishart base is conjugate, so $(\mu_k,\Sigma_k)$ integrate out (Neal 2000, Algorithm 3) and a point's predictive under a component is a multivariate Student-$t$; we resample only the labels.

In [1]:
import numpy as np, pandas as pd, matplotlib.pyplot as plt
from scipy.stats import gaussian_kde, norm
import dpmix as D
rng = np.random.default_rng(4)
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; PURP="#6b46c1"; GREY="#718096"
print("Dirichlet process = a prior over discrete random probability measures.")

Dirichlet process = a prior over discrete random probability measures.

1. What is a Dirichlet process? — stick-breaking a random measure¶

A draw $G\sim\text{DP}(\alpha,G_0)$ is a discrete distribution, even when the base $G_0$ is continuous. Sethuraman's stick-breaking builds it explicitly: break a unit stick with proportions $v_k\sim\text{Beta}(1,\alpha)$, put atoms at locations $\theta_k\sim G_0$, and stack them, $$w_k=v_k\prod_{j<k}(1-v_j),\qquad G=\sum_{k=1}^{\infty}w_k\,\delta_{\theta_k}.$$ Small $\alpha$ concentrates the weight on a few atoms (strong clustering); large $\alpha$ spreads it over many (so $G$ approaches $G_0$). Below: draws of the weights and of the random measure itself, for three concentrations, with the base $G_0=\mathcal N(0,1)$.

In [2]:
fig,ax=plt.subplots(2,3,figsize=(13,6.2))
for j,al in enumerate([1.0,5.0,25.0]):
    w=D.stick_breaking(al,40,rng); th=rng.standard_normal(40)
    ax[0,j].bar(range(1,21), w[:20], color=BLUE); ax[0,j].set_title(f"stick weights $w_k$   ($\\alpha={al:.0f}$)")
    ax[0,j].set_xlabel("atom k"); ax[0,j].set_ylabel("weight")
    # the random measure vs the base CDF
    o=np.argsort(th); ax[1,j].vlines(th, 0, w, color=GREY, alpha=.7)
    ax[1,j].scatter(th, w, s=14, color=RED, zorder=3)
    xx=np.linspace(-3,3,200); ax[1,j].plot(xx, 0.06*norm.pdf(xx), color=GREEN, lw=1.5, label="$G_0$ density (scaled)")
    ax[1,j].set_title(f"random measure $G$   ($E[K]\\approx{D.expected_clusters(al,200):.0f}$/200)")
    ax[1,j].set_xlabel(r"$\theta$"); ax[1,j].set_ylabel("weight"); ax[1,j].legend(frameon=False,fontsize=7)
plt.tight_layout(); plt.show()
print("As alpha grows the stick weights spread over more atoms and G fills in toward the base G0;")
print("as alpha shrinks a few atoms carry almost all the mass -- that concentration IS the clustering prior.")
No description has been provided for this image 
As alpha grows the stick weights spread over more atoms and G fills in toward the base G0;
as alpha shrinks a few atoms carry almost all the mass -- that concentration IS the clustering prior.

2. A bivariate DP mixture — density and clustering of Old Faithful¶

The Old Faithful geyser: 272 eruptions, each an (eruption duration, waiting time) pair — famously two clouds (short/quick vs long/slow). We fit a bivariate DP mixture: it estimates the joint density and, as a by-product, a clustering with the number of groups inferred rather than fixed.

In [3]:
F = pd.read_csv("faithful.csv").to_numpy()          # columns: eruptions, waiting
gx,gy = np.meshgrid(np.linspace(1.4,5.4,45), np.linspace(41,98,45))
grid = np.column_stack([gx.ravel(), gy.ravel()])
fit = D.dpmix_gibbs(F, rng, alpha=1.0, draws=1500, burn=1000, grid=grid)
dens = fit["density"].reshape(gx.shape)
Kmode = np.bincount(fit["K"]).argmax()
print(f"posterior number of clusters K: mode {Kmode}, mean {fit['K'].mean():.2f}")
print("K posterior:", {int(k):round(v,2) for k,v in zip(*np.unique(fit['K'],return_counts=True)/np.array([1,len(fit['K'])]))} if False else dict(zip(*[list(map(int,np.unique(fit['K']))), np.round(np.bincount(fit['K'])[np.unique(fit['K'])]/len(fit['K']),2)])))

fig,ax=plt.subplots(1,2,figsize=(12.5,4.8))
ax[0].scatter(F[:,0],F[:,1],s=12,color=GREY,alpha=.6)
ax[0].contour(gx,gy,dens,levels=8,cmap="viridis")
ax[0].set_xlabel("eruption duration (min)"); ax[0].set_ylabel("waiting to next (min)")
ax[0].set_title(f"DP-mixture joint density (K mode = {Kmode})")
# hard clustering from the representative labels
z=fit["z"]; labs=np.unique(z); cols=[BLUE,RED,GREEN,ORANGE,PURP,GREY]
for i,l in enumerate(labs): m=z==l; ax[1].scatter(F[m,0],F[m,1],s=16,color=cols[i%len(cols)],label=f"cluster {i+1} (n={m.sum()})")
ax[1].set_xlabel("eruption duration (min)"); ax[1].set_ylabel("waiting to next (min)")
ax[1].set_title("Representative clustering"); ax[1].legend(frameon=False,fontsize=8)
plt.tight_layout(); plt.show()
print("The DP mixture recovers the two main clouds plus a small third component bridging them, and")
print("unlike a fixed-K fit it returns the UNCERTAINTY in that count rather than a single number.")
print("\nHow much of that count is the data and how much is the prior? The concentration alpha is the")
print("usual suspect, so it is worth checking rather than assuming:")
for a_ in (0.5, 1.0, 2.0):
    f_ = D.dpmix_gibbs(F, np.random.default_rng(0), alpha=a_, draws=400, burn=300)
    print("   alpha=%.1f  ->  K mode %d, mean %.2f   (prior E[K] = %.1f)"
          % (a_, np.bincount(f_["K"]).argmax(), f_["K"].mean(), D.expected_clusters(a_, len(F))))
print("The mode is stable at 3 across a fourfold range of alpha even though the PRIOR expectation")
print("moves from about 4 to about 10 clusters -- so on this dataset the likelihood, not the")
print("concentration parameter, is doing the work. Keep that in mind at the PyMC cross-check below,")
print("where a different count appears for reasons that have nothing to do with alpha.")

posterior number of clusters K: mode 3, mean 3.36
K posterior: {2: np.float64(0.14), 3: np.float64(0.47), 4: np.float64(0.29), 5: np.float64(0.08), 6: np.float64(0.02), 7: np.float64(0.0)}
No description has been provided for this image 
The DP mixture recovers the two main clouds plus a small third component bridging them, and
unlike a fixed-K fit it returns the UNCERTAINTY in that count rather than a single number.

How much of that count is the data and how much is the prior? The concentration alpha is the
usual suspect, so it is worth checking rather than assuming:

   alpha=0.5  ->  K mode 3, mean 2.91   (prior E[K] = 3.8)

   alpha=1.0  ->  K mode 3, mean 3.25   (prior E[K] = 6.2)

   alpha=2.0  ->  K mode 3, mean 3.94   (prior E[K] = 10.4)
The mode is stable at 3 across a fourfold range of alpha even though the PRIOR expectation
moves from about 4 to about 10 clusters -- so on this dataset the likelihood, not the
concentration parameter, is doing the work. Keep that in mind at the PyMC cross-check below,
where a different count appears for reasons that have nothing to do with alpha.

3. The frequentist counterpart — kernel density estimation¶

The classical nonparametric density estimator is the kernel density estimate: place a small Gaussian kernel on every observation and sum. It smooths beautifully but does only smoothing — no clustering, no number-of-groups, no uncertainty, and it hinges on a bandwidth choice. Here the 2-D KDE and the DP-mixture density side by side.

In [4]:
kde = gaussian_kde(F.T)                              # Scott's-rule bandwidth
kd = kde(grid.T).reshape(gx.shape)
fig,ax=plt.subplots(1,2,figsize=(12.5,4.7))
for a,(Z,ttl) in zip(ax,[(kd,"Kernel density estimate (frequentist)"),(dens,"DP-mixture density (Bayesian NP)")]):
    a.contourf(gx,gy,Z,levels=12,cmap="viridis"); a.scatter(F[:,0],F[:,1],s=8,color="white",alpha=.4)
    a.set_xlabel("eruption (min)"); a.set_ylabel("waiting (min)"); a.set_title(ttl)
plt.tight_layout(); plt.show()
print("Both render the two-bump density well. The difference is what you GET: the KDE is a smooth surface and")
print("nothing more; the DP mixture also delivers a clustering, a posterior over the number of groups, and full")
print("uncertainty -- at the cost of a model. KDE = smoothing; DP mixture = a generative model that smooths.")
No description has been provided for this image 
Both render the two-bump density well. The difference is what you GET: the KDE is a smooth surface and
nothing more; the DP mixture also delivers a clustering, a posterior over the number of groups, and full
uncertainty -- at the cost of a model. KDE = smoothing; DP mixture = a generative model that smooths.

4. Pitman-Yor — when clusters follow a power law¶

The Pitman-Yor process adds a discount $d\in[0,1)$: an existing cluster of size $n_k$ is joined with weight $\propto n_k-d$, a new cluster opened with weight $\propto \alpha+dK$. The discount rewards new clusters, so the number of clusters grows like a power law $n^d$ rather than the DP's $\log n$ — the right prior when there are many rare types (words, species, surnames). The growth curves make the difference vivid.

In [5]:
ns=np.array([10,25,50,100,200,400,800,1600])
fig,ax=plt.subplots(1,2,figsize=(12,4.4))
for (al,dd,c,lab) in [(1.0,0.0,BLUE,"DP  $\\alpha=1,d=0$"),(1.0,0.5,RED,"PY  $\\alpha=1,d=0.5$"),(1.0,0.7,ORANGE,"PY  $\\alpha=1,d=0.7$")]:
    ek=[D.expected_clusters(al,int(n),dd) for n in ns]; ax[0].plot(ns,ek,'o-',color=c,label=lab)
ax[0].set_xscale("log"); ax[0].set_xlabel("n (log scale)"); ax[0].set_ylabel("E[ #clusters ]")
ax[0].set_title("DP grows like log n; Pitman-Yor like a power law $n^d$"); ax[0].legend(frameon=False)
# fit PY vs DP to faithful and compare cluster counts
kDP=[D.dpmix_gibbs(F,rng,alpha=1.0,d=0.0,draws=600,burn=400)["K"] for _ in range(1)][0]
kPY=[D.dpmix_gibbs(F,rng,alpha=1.0,d=0.6,draws=600,burn=400)["K"] for _ in range(1)][0]
ax[1].hist(kDP,bins=np.arange(1,12)-.5,alpha=.6,color=BLUE,label=f"DP  (mode {np.bincount(kDP).argmax()})",density=True)
ax[1].hist(kPY,bins=np.arange(1,12)-.5,alpha=.6,color=RED,label=f"PY d=0.6 (mode {np.bincount(kPY).argmax()})",density=True)
ax[1].set_xlabel("number of clusters on Old Faithful"); ax[1].set_ylabel("posterior"); ax[1].set_title("DP vs Pitman-Yor on the geyser"); ax[1].legend(frameon=False)
plt.tight_layout(); plt.show()
print("Same data, more clusters under Pitman-Yor: the discount d makes new clusters cheaper, so PY prefers more,")
print("smaller groups. On a small tidy dataset like Old Faithful that mostly adds tiny extra components; its real")
print("value is heavy-tailed clustering (texts, species-abundance) where the DP's log-n growth is too slow.")
No description has been provided for this image 
Same data, more clusters under Pitman-Yor: the discount d makes new clusters cheaper, so PY prefers more,
smaller groups. On a small tidy dataset like Old Faithful that mostly adds tiny extra components; its real
value is heavy-tailed clustering (texts, species-abundance) where the DP's log-n growth is too slow.

5. Cross-check in PyMC — truncated stick-breaking, bivariate¶

A truncated stick-breaking DP with multivariate-normal components, marginalised over labels by pm.Mixture. We cap the sticks at $K_{\max}=8$, put a Normal base on the means and an LKJ prior on a shared component covariance (shared for a stable, well-initialised fit — the collapsed sampler let every cluster have its own), and compare the resulting density.

In [6]:
import pymc as pm, pytensor.tensor as pt
Fs = (F - F.mean(0))/F.std(0); Kmax=8
with pm.Model() as mod:
    alpha = pm.Gamma("alpha", 2, 2)
    beta  = pm.Beta("beta", 1, alpha, shape=Kmax)
    stick = beta*pt.concatenate([[1.0], pt.cumprod(1-beta)[:-1]])
    w     = pm.Deterministic("w", stick/pt.sum(stick))          # normalise the truncated sticks
    mu    = pm.Normal("mu", 0, 2, shape=(Kmax,2),
                      initval=np.linspace(-1.5,1.5,Kmax)[:,None]*np.ones((1,2)))
    chol,_,_ = pm.LKJCholeskyCov("chol", n=2, eta=2.0, sd_dist=pm.HalfNormal.dist(1.0), compute_corr=True)
    comps = pm.MvNormal.dist(mu, chol=chol)                     # batched over Kmax (shared covariance)
    pm.Mixture("obs", w=w, comp_dists=comps, observed=Fs)
    idata = pm.sample(800, tune=1200, chains=4, target_accept=0.95, random_seed=6,
                      init="adapt_diag", progressbar=False)

W=idata.posterior["w"].values.reshape(-1,Kmax)
ndiv = int(idata.sample_stats["diverging"].values.sum())
print("Sampler diagnostics: %d divergences out of %d draws." % (ndiv, W.shape[0]))
print("R-hat on the component parameters is NOT a usable diagnostic here: a mixture likelihood is")
print("invariant to relabelling its components, so chains that explore different labellings of the")
print("SAME density look divergent by construction. That is why the comparison below is made on the")
print("fitted DENSITY, which is label-invariant, rather than on mu, w or the cluster count.\n")
print("PyMC effective clusters (weight>1/n): mean %.1f, mode %d"
      % ((W>1/len(F)).sum(1).mean(), np.bincount((W>1/len(F)).sum(1)).argmax()))
# posterior-mean density on the grid (average the mixture over draws)
mu_s=idata.posterior["mu"].values.reshape(-1,Kmax,2)
gsz=(grid-F.mean(0))/F.std(0); S=W.shape[0]; idx=rng.choice(S,300,replace=False)
from scipy.stats import multivariate_normal as mvn
stds=idata.posterior["chol_stds"].values.reshape(-1,2); corr=idata.posterior["chol_corr"].values.reshape(-1,2,2)
dens_pm=np.zeros(len(grid))
for s in idx:
    C=np.diag(stds[s])@corr[s]@np.diag(stds[s])
    for k in range(Kmax):
        if W[s,k]>1e-3: dens_pm+=W[s,k]*mvn.pdf(gsz,mean=mu_s[s,k],cov=C)
dens_pm=(dens_pm/len(idx))/np.prod(F.std(0))
fig,ax=plt.subplots(1,2,figsize=(12.5,4.7))
for a,(Z,ttl) in zip(ax,[(dens,"from-scratch collapsed Gibbs"),(dens_pm.reshape(gx.shape),"PyMC stick-breaking")]):
    a.contourf(gx,gy,Z,levels=12,cmap="viridis"); a.scatter(F[:,0],F[:,1],s=8,color="white",alpha=.4)
    a.set_xlabel("eruption (min)"); a.set_ylabel("waiting (min)"); a.set_title("DP density: "+ttl)
plt.tight_layout(); plt.show()
l1 = np.abs(dens/dens.sum() - dens_pm.reshape(gx.shape)/dens_pm.sum()).sum()
print("The two DP constructions -- collapsed Chinese-restaurant and truncated stick-breaking -- recover")
print("the same two-cloud density (total variation distance between the normalised grids: %.3f)." % (0.5*l1))
print("\nThe CLUSTER COUNTS, however, do not match, and the reason is worth being explicit about because")
print("it is the single most misread number in nonparametric clustering:")
print("  * the collapsed Gibbs above counts OCCUPIED clusters -- components with data assigned;")
print("  * this cell counts sticks with weight > 1/n = %.4f, which is a different estimand." % (1/len(F)))
print("    With Kmax=%d, near-empty sticks clear that threshold without owning a single point." % Kmax)
print("  * and this model shares ONE covariance across all components, while the from-scratch sampler")
print("    gives each cluster its own. Two clouds of different shape need more shared-shape components")
print("    to cover them, so this construction is structurally biased toward a larger count.")
print("\nThe R engine, with a third prior, reports a third number again. None of them is wrong. The count")
print("of clusters in a DP mixture is a property of the MODEL AND THE ESTIMAND, not a fact about the")
print("data -- which is exactly why the posterior over K, and not a single K, is the thing to report.")

g++ not available, if using conda: `conda install gxx`

Initializing NUTS using adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [alpha, beta, mu, chol]

Sampling 4 chains for 1_200 tune and 800 draw iterations (4_800 + 3_200 draws total) took 34 seconds.

There were 32 divergences after tuning. Increase `target_accept` or reparameterize.

The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details

The effective sample size per chain is smaller than 100 for some parameters.  A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details

Sampler diagnostics: 32 divergences out of 3200 draws.
R-hat on the component parameters is NOT a usable diagnostic here: a mixture likelihood is
invariant to relabelling its components, so chains that explore different labellings of the
SAME density look divergent by construction. That is why the comparison below is made on the
fitted DENSITY, which is label-invariant, rather than on mu, w or the cluster count.

PyMC effective clusters (weight>1/n): mean 4.4, mode 4
No description has been provided for this image 
The two DP constructions -- collapsed Chinese-restaurant and truncated stick-breaking -- recover
the same two-cloud density (total variation distance between the normalised grids: 0.059).

The CLUSTER COUNTS, however, do not match, and the reason is worth being explicit about because
it is the single most misread number in nonparametric clustering:
  * the collapsed Gibbs above counts OCCUPIED clusters -- components with data assigned;
  * this cell counts sticks with weight > 1/n = 0.0037, which is a different estimand.
    With Kmax=8, near-empty sticks clear that threshold without owning a single point.
  * and this model shares ONE covariance across all components, while the from-scratch sampler
    gives each cluster its own. Two clouds of different shape need more shared-shape components
    to cover them, so this construction is structurally biased toward a larger count.

The R engine, with a third prior, reports a third number again. None of them is wrong. The count
of clusters in a DP mixture is a property of the MODEL AND THE ESTIMAND, not a fact about the
data -- which is exactly why the posterior over K, and not a single K, is the thing to report.

6. Summary¶

A Dirichlet process is a prior over discrete random measures: stick-breaking makes that concrete, and the concentration $\alpha$ is exactly a clustering dial. Fitting a bivariate DP mixture to Old Faithful estimated the joint density and recovered the clustering with the number of groups inferred — the two clouds plus a small bridging component, matching mclust's $K=3$, but with the count's uncertainty attached. The Pitman-Yor discount turns the DP's $\log n$ cluster growth into a power law, the right prior for heavy-tailed, many-rare-types data.

Beside each Bayesian tool sat its frequentist analog: kernel density estimation renders the same two-bump surface but delivers only smoothing — no clusters, no $K$, no uncertainty — the classic trade of a model for pure flexibility. The from-scratch collapsed sampler, the PyMC stick-breaking mixture, and (in the companion notebook) R's dirichletprocess/BNPmix and mclust all agree.

This is the multivariate, random-measure face of the DP; the variable-selection arc's DP notebook is its univariate "how many components" face, and the discrete version is the latent-class arc's DP-LCA. The next step, the hierarchical Dirichlet process, lets several groups share one set of atoms — the bridge to topic models and grouped density estimation.