Dirichlet-Process & Pitman–Yor Mixtures

Python · PyMC · R  ·  Download DP-mixture module

A Prior on the Mixing Distribution

A finite mixture model asks you to name KK before you look. The Dirichlet process removes that requirement by putting a prior on the mixing distribution itself — a prior over discrete random probability measures, constructed by stick-breaking: break a unit stick at a Beta(1,α)\text{Beta}(1,\alpha) point, keep the piece, break what remains, and repeat. The concentration α\alpha controls how fast the pieces shrink, and that decay is the clustering prior: small α\alpha puts nearly all the mass on a few atoms, large α\alpha spreads it toward the base measure G0G_0.

G=∑k=1∞wk δθk,wk=vk∏j<k(1−vj),vk∼Beta(1−d,  α+kd)G=\sum_{k=1}^{\infty}w_k\,\delta_{\theta_k},\qquad w_k=v_k\prod_{j<k}(1-v_j),\qquad v_k\sim\text{Beta}(1-d,\;\alpha+kd)

Old Faithful, and what a KDE cannot tell you

Applied to the Old Faithful geyser — 272 eruptions, duration against waiting time — the mixture recovers the two obvious clouds plus a small third component bridging them, and returns the whole posterior over how many there are rather than a single count. The frequentist comparison makes the trade explicit. A kernel density estimate renders the same two-bump surface perfectly well; what it cannot do is hand back a clustering, a component count, or any statement of uncertainty about either. KDE is smoothing; a DP mixture is a generative model that happens to smooth — and the extra structure is paid for with a model you have to believe.

Pitman–Yor: when clusters follow a power law

Pitman–Yor generalises the construction with a discount dd that makes each new cluster cheaper than the last. The consequence is asymptotic: the DP's expected cluster count grows like αlog⁡n\alpha\log n, Pitman–Yor's like ndn^{d}. On a small tidy dataset that mostly adds a few tiny components, and the notebooks say so rather than overselling it. Its real domain is heavy-tailed clustering — word frequencies, species abundance — where the number of distinct types keeps climbing with the sample and logarithmic growth simply cannot keep up.

Three Engines, Three Cluster Counts

The most useful result in this project is one that looks at first like agreement. Three engines fit a DP mixture to the same 272 points and report three different cluster counts: the from-scratch collapsed Gibbs sampler says 3, R's dirichletprocess says 2, and PyMC's truncated stick-breaking says 4. Each write-up had framed its own number as confirming the others.

engineconstructioncovariancewhat it countsKK
from-scratch Gibbscollapsed Chinese restaurantper cluster (NIW)occupied clusters3
R dirichletprocessconjugate base, α\alpha sampledper clusteroccupied clusters2
PyMCtruncated stick-breaking, Kmax⁡=8K_{\max}=8sharedweight >1/n>1/n4

They are not measuring the same quantity. The collapsed sampler counts occupied clusters — components with data actually assigned. The PyMC cell counts sticks whose weight exceeds 1/n=0.00371/n = 0.0037, a threshold that near-empty sticks clear without owning a single point, and that model also shares one covariance matrix across all eight components, so it needs extra components to cover two clouds of different shape. R's package uses a third base measure and samples α\alpha rather than fixing it. The obvious suspect is ruled out by direct test: holding everything else fixed and sweeping α\alpha from 0.5 to 2.0 moves the prior expected count from about 4 to about 10 while the posterior mode stays at 3. The likelihood, not the concentration parameter, is doing the work.

concentration α\alphaprior E[K]\mathbb{E}[K]posterior modeposterior mean
0.53.832.91
1.06.233.25
2.010.433.94

So the honest conclusion is not that three engines agree, but that the number of clusters in a DP mixture is a property of the model and the estimand rather than a fact about the data. What does agree is the thing that ought to: the fitted densities from the collapsed Chinese-restaurant and truncated stick-breaking constructions match to a total variation distance of 0.059. The density is label-invariant; the cluster count is not. That is also why R-hat on the component parameters is reported here as uninformative — a mixture likelihood is invariant to relabelling, so chains exploring different labellings of the same density look divergent by construction.

Notebooks

Python Stick-breaking the random measure, a bivariate DP mixture of Old Faithful with the posterior over K, the KDE comparison, Pitman–Yor power-law clustering, and a PyMC stick-breaking cross-check (Python) Open → R Independent R implementation — stick-breaking, mclust and a kernel density estimate as the frequentist baselines, the dirichletprocess package's posterior over K, and DP versus Pitman–Yor growth (R) Open →

Downloads

bnp_dpmix.py Stick-breaking and expected cluster counts; a collapsed Chinese-restaurant Gibbs sampler with normal-inverse-Wishart components; Pitman–Yor CRP simulation and the adjusted Rand index (NumPy / SciPy) faithful.csv Old Faithful geyser — 272 eruptions, duration and waiting time

DP-Mixture Module — Source Code

"""
dpmix.py -- DIRICHLET-PROCESS & PITMAN-YOR MIXTURES, the random-measure view (from scratch).

Backs the notebooks in  "Dirichlet-Process & Pitman-Yor Mixtures".

This project is the MULTIVARIATE / random-measure companion to the univariate DP-mixture
notebook in the variable-selection arc (`Dirichlet-Process Mixtures -- How Many Components?`,
which fits the 1-D galaxy data and asks "how many components?"). Here we instead:

  1. show what a Dirichlet process IS -- a prior over DISCRETE random probability measures,
     built by STICK-BREAKING:  G = sum_k w_k delta_{theta_k},  theta_k ~ G0,
     w_k = v_k prod_{j<k}(1-v_j),  v_k ~ Beta(1-d, alpha + k d);
  2. fit a BIVARIATE DP mixture (Normal-InverseWishart base) for 2-D density estimation and
     clustering (Old Faithful), with a multivariate Student-t predictive;
  3. generalise the DP to the two-parameter PITMAN-YOR process (discount d), whose cluster
     sizes follow a POWER law -- more, smaller clusters than the DP.

Stick-breaking weights, DP vs Pitman-Yor:
    d = 0                -> Dirichlet process,  E[#clusters] ~ alpha * log n
    0 < d < 1            -> Pitman-Yor,         E[#clusters] ~ (alpha/d?) n^d   (power law)

The mixture sampler is the collapsed CRP Gibbs of Neal (2000, Algorithm 3): the component
means/covariances (mu_k, Sigma_k) are integrated out under the conjugate Normal-InverseWishart
base, so a point's predictive under a component is a multivariate Student-t and we resample only
the labels. Data are standardised per dimension so the NIW hyperparameters are portable; the
density is mapped back to the original scale.
"""

import numpy as np
from scipy.special import gammaln, multigammaln


# --------------------------------------------------------------------------- #
#  Stick-breaking: DRAW a random probability measure G ~ DP / Pitman-Yor        #
# --------------------------------------------------------------------------- #

def stick_breaking(alpha, K, rng, d=0.0):
    """Return the first K stick-breaking weights of a DP(alpha) (d=0) or Pitman-Yor(alpha,d).
    v_k ~ Beta(1-d, alpha + k*d);  w_k = v_k * prod_{j<k}(1-v_j)."""
    k = np.arange(K)
    v = rng.beta(1.0 - d, alpha + (k + 1) * d)          # (index k+1 so first is Beta(1-d, alpha+d))
    v[-1] = 1.0                                          # truncate: last stick takes the remainder
    w = v * np.concatenate([[1.0], np.cumprod(1.0 - v)[:-1]])
    return w

def expected_clusters(alpha, n, d=0.0):
    """E[number of occupied clusters] after n draws from a DP / Pitman-Yor CRP (exact recursion)."""
    ek = 0.0
    for i in range(n):
        ek += (alpha + d * ek) / (alpha + i)            # prob draw i+1 starts a new cluster
    return ek


# --------------------------------------------------------------------------- #
#  Multivariate Student-t (the Normal-InverseWishart marginal)                  #
# --------------------------------------------------------------------------- #

def _mvt_logpdf(X, df, loc, scale):
    """log pdf of a multivariate Student-t at rows of X. scale is the D x D scale matrix."""
    X = np.atleast_2d(X); D = X.shape[1]
    L = np.linalg.cholesky(scale); logdet = 2 * np.log(np.diag(L)).sum()
    z = np.linalg.solve(L, (X - loc).T)                 # (D, n)
    maha = (z * z).sum(0)
    return (gammaln((df + D) / 2) - gammaln(df / 2) - 0.5 * D * np.log(df * np.pi)
            - 0.5 * logdet - (df + D) / 2 * np.log1p(maha / df))

def _niw_pred(n, s, S, m0, k0, nu0, Psi0, D):
    """multivariate-t predictive (df, loc, scale) for a cluster with count n, sum vector s,
    sum-of-outer-products S; n=0 gives the prior predictive under the NIW base."""
    if n == 0:
        kn, nun, mn, Psin = k0, nu0, m0, Psi0
    else:
        xbar = s / n
        kn = k0 + n; nun = nu0 + n
        mn = (k0 * m0 + s) / kn
        Cov = S - np.outer(s, xbar) - np.outer(xbar, s) + n * np.outer(xbar, xbar)   # scatter
        Psin = Psi0 + Cov + (k0 * n / kn) * np.outer(xbar - m0, xbar - m0)
    df = nun - D + 1
    scale = Psin * (kn + 1) / (kn * df)
    return df, mn, scale


# --------------------------------------------------------------------------- #
#  Collapsed CRP Gibbs for a multivariate DP / Pitman-Yor mixture               #
# --------------------------------------------------------------------------- #

def dpmix_gibbs(Y, rng, alpha=1.0, d=0.0, m0=None, k0=0.1, nu0=None, Psi0=None,
                draws=2000, burn=1000, grid=None):
    """Collapsed Gibbs for a DP(alpha) / Pitman-Yor(alpha,d) multivariate Normal mixture.
    Y:(N,D) standardised internally. Returns the posterior of the number of clusters K, a
    representative label vector, the cluster occupancy, and -- if a 2-D `grid` (G,2) is given --
    the posterior-predictive density on that grid (original units)."""
    Y = np.asarray(Y, float); N, D = Y.shape
    mu_, sd_ = Y.mean(0), Y.std(0)
    Ys = (Y - mu_) / sd_
    if m0 is None: m0 = np.zeros(D)
    if nu0 is None: nu0 = D + 2.0
    if Psi0 is None: Psi0 = np.eye(D)
    if grid is not None:
        grid = np.asarray(grid, float); gs = (grid - mu_) / sd_; dens = np.zeros(len(grid))
    z = np.zeros(N, int)
    cnt = [float(N)]; sm = [Ys.sum(0)]; SS = [Ys.T @ Ys]
    Ks = []; nsave = 0
    for it in range(draws + burn):
        for i in range(N):
            yi = Ys[i]; c = z[i]
            cnt[c] -= 1; sm[c] -= yi; SS[c] -= np.outer(yi, yi)
            if cnt[c] == 0:
                del cnt[c]; del sm[c]; del SS[c]; z[z > c] -= 1
            K = len(cnt); logp = np.empty(K + 1)
            for k in range(K):
                df, loc, sc = _niw_pred(cnt[k], sm[k], SS[k], m0, k0, nu0, Psi0, D)
                logp[k] = np.log(cnt[k] - d) + _mvt_logpdf(yi, df, loc, sc)[0]
            df, loc, sc = _niw_pred(0, None, None, m0, k0, nu0, Psi0, D)
            logp[K] = np.log(alpha + d * K) + _mvt_logpdf(yi, df, loc, sc)[0]
            logp -= logp.max(); p = np.exp(logp); p /= p.sum()
            new = rng.choice(K + 1, p=p)
            if new == K:
                cnt.append(0.0); sm.append(np.zeros(D)); SS.append(np.zeros((D, D)))
            z[i] = new; cnt[new] += 1; sm[new] += yi; SS[new] += np.outer(yi, yi)
        if it >= burn:
            Ks.append(len(cnt)); nsave += 1
            if grid is not None:
                K = len(cnt); denom = N + alpha; f = np.zeros(len(grid))
                for k in range(K):
                    df, loc, sc = _niw_pred(cnt[k], sm[k], SS[k], m0, k0, nu0, Psi0, D)
                    f += (cnt[k] - d) * np.exp(_mvt_logpdf(gs, df, loc, sc))
                df, loc, sc = _niw_pred(0, None, None, m0, k0, nu0, Psi0, D)
                f += (alpha + d * K) * np.exp(_mvt_logpdf(gs, df, loc, sc))
                dens += (f / denom) / np.prod(sd_)
    out = dict(K=np.array(Ks), z=z.copy(), counts=np.array(cnt))
    if grid is not None:
        out["grid"] = grid; out["density"] = dens / nsave
    return out


# --------------------------------------------------------------------------- #
#  simulation helpers                                                           #
# --------------------------------------------------------------------------- #

def simulate_mv_mixture(N, weights, means, covs, rng):
    """draw N points from a finite multivariate Normal mixture (the 'truth')."""
    weights = np.asarray(weights, float); weights /= weights.sum()
    k = rng.choice(len(weights), size=N, p=weights)
    means = np.asarray(means); covs = np.asarray(covs)
    return np.array([rng.multivariate_normal(means[kk], covs[kk]) for kk in k]), k

def crp_simulate_K(alpha, n, rng, d=0.0, reps=200):
    """simulate the number of clusters formed by n CRP draws (DP if d=0, else Pitman-Yor),
    returning the mean over `reps` -- used to contrast DP (log n) with PY (power-law) growth."""
    Ks = np.empty(reps)
    for r in range(reps):
        counts = []
        for i in range(n):
            probs = np.array([c - d for c in counts] + [alpha + d * len(counts)])
            probs /= probs.sum()
            j = rng.choice(len(counts) + 1, p=probs)
            if j == len(counts): counts.append(1.0)
            else: counts[j] += 1.0
        Ks[r] = len(counts)
    return Ks.mean()

def adjusted_rand(a, b):
    """adjusted Rand index between two label vectors (clustering-recovery check)."""
    a = np.asarray(a); b = np.asarray(b); n = len(a)
    ca = {v: i for i, v in enumerate(np.unique(a))}; cb = {v: i for i, v in enumerate(np.unique(b))}
    M = np.zeros((len(ca), len(cb)))
    for i in range(n): M[ca[a[i]], cb[b[i]]] += 1
    su = M.sum(1); sv = M.sum(0)
    idx = (M * (M - 1) / 2).sum(); ei = (su * (su - 1) / 2).sum(); ej = (sv * (sv - 1) / 2).sum()
    exp = ei * ej / (n * (n - 1) / 2); mx = 0.5 * (ei + ej)
    return (idx - exp) / (mx - exp) if mx != exp else 1.0

References