The Hierarchical Dirichlet Process

Sharing Clusters Across Groups

A single Dirichlet process handles one exchangeable population. Real data usually arrives in groups — patients in hospitals, documents in corpora, birds on islands — and the interesting question is whether those groups draw on the same underlying set of components in different proportions. Fitting a separate DP per group cannot answer it, because each group invents its own labels and nothing connects them. Pooling the data answers it by destroying the grouping. The hierarchical Dirichlet process stacks two DPs: a global one draws a shared menu of clusters, and each group draws its own mixing proportions over that same menu.

Verified against ground truth first

The construction is verified against known ground truth before it is trusted on real data. Three simulated groups drawn from three shared components in deliberately different proportions — one group never using the third component at all — are recovered exactly: cluster count mode 3, ARI 1.000 — the adjusted Rand index scores agreement between two clusterings on a scale where chance is 0 and an exact match is 1 — and a per-group weight table that reproduces the generating proportions including the structural zeros. That table is the deliverable. A separate DP per group could not even represent it.

The penguins, and the biogeography that falls out

On the Palmer penguins — 342 birds across three islands, four morphological measurements — the HDP recovers the species almost perfectly (ARI 0.998) and, simultaneously, the biogeography. One cluster carries weight on all three islands: that is Adélie, the species actually found on every island. Gentoo loads Biscoe alone and Chinstrap loads Dream alone, which is exactly the true distribution. The fitted concentrations tell the same story from the other side — $\alpha=0.17$ within islands (each island uses few clusters) against $\gamma=1.07$ globally.

A Table Normalised Along the Wrong Axis

Getting that table right depends on normalising along the correct axis, and the arithmetic makes a slip obvious. Rows reading 0.62 + 0.00 + 0.28 = 0.90 and 0.00 + 0.67 + 0.42 = 1.09 — impossible for mixing proportions — while the columns summed to 1.00. Transposing the count matrix and then normalising divides each cluster across groups; normalising after the transpose divides each group across clusters, which is what mixing proportions mean. The same matrix drives the figure and the test for which cluster appears on two or more islands, so the sharing claim depends on it as well. Normalised within group, the simulated recovery reproduces the generating weights: (0.69, 0.00, 0.31) against a true (0.7, 0.3, 0.0) up to a label permutation, and every row sums to one.

What the comparison actually shows

The comparison against the alternatives is worth arguing carefully. The pooled mclust fit scores ARI 0.960 against the HDP's 0.958 — the HDP is marginally worse — and the reason that does not matter is that on a dataset this well separated almost anything recovers three species, so the case for the HDP is not that it clusters better. It is that it answers a question the alternatives cannot pose: independent fits cannot say Torgersen's Adélie is Biscoe's Adélie, and the pooled fit can only link them by discarding the island, after which it can no longer say Gentoo is absent from Torgersen. The weight table, not the ARI, is the output. Judge a model by what it can express, not only by a scalar score.

Where the cross-check parts company

The PyMC cross-check disagrees by more than the qualitative story suggests. A truncated stick-breaking HDP finds 5 active clusters against the collapsed sampler's 4, with ARI 0.781 against 0.998 — a 0.22 drop, which is not a rounding difference. The cause is structural: that model shares one covariance matrix across all components, so clusters differing in shape can only be covered by splitting them, which is precisely what the surplus cluster does. It is the same lesson the DP-mixture project ends on — the partition and the cluster count are properties of the model you chose, and two correct implementations of "an HDP" need not agree on them.

Notebooks

Downloads

HDP Module — Source Code

"""
hdp.py -- HIERARCHICAL DIRICHLET PROCESS for grouped data (from scratch).

Backs the notebooks in  "The Hierarchical Dirichlet Process".

A single Dirichlet-process mixture clusters ONE dataset. When data come in GROUPS (documents,
islands, patients, ...) we usually want the groups to share ONE set of clusters but with their
own mixing proportions -- Adelie penguins appear on every island, Gentoo only on some. The
Hierarchical Dirichlet Process (Teh, Jordan, Beal & Blei 2006) does exactly this by stacking
two DPs so the group-level measures share a common, DISCRETE global measure (hence shared atoms):

    G0 ~ DP(gamma, H)                      (global measure -- the shared menu of clusters)
    G_j ~ DP(alpha, G0)   for each group j (group measure -- reuses G0's atoms, own weights)
    theta_ji ~ G_j ,  x_ji ~ F(theta_ji)   (observations in group j)

Because G0 is discrete, the G_j necessarily reuse the SAME atoms: clusters are shared across
groups, and the number of them is inferred. The Chinese-restaurant analogue is a FRANCHISE: each
group is a restaurant with its own tables (a CRP with concentration alpha), but every table
serves a dish from a single global menu shared by all restaurants (a CRP with concentration
gamma). A dish (cluster) can be served in several restaurants -- that is the sharing.

From scratch we use Teh et al.'s DIRECT-ASSIGNMENT sampler: keep global mixing weights beta over
the shared clusters (a stick-breaking measure with concentration gamma), assign each observation
to a cluster with probability proportional to (n_jk + alpha*beta_k) * f_k(x), and after each
sweep resample beta from the number of TABLES via the Antoniak (Chinese-restaurant) distribution.
Cluster parameters use a conjugate Normal-Inverse-Wishart base, so f_k is a multivariate
Student-t and the parameters integrate out.  H = NIW(m0, kappa0, nu0, Psi0).
"""

import numpy as np
from scipy.special import gammaln


# --------------------------------------------------------------------------- #
#  NIW multivariate-t predictive (shared with the DP-mixture project)           #
# --------------------------------------------------------------------------- #

def _mvt_logpdf(x, df, loc, scale):
    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); 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 _pred(n, s, S, m0, k0, nu0, Psi0, D):
    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)
        Psin = Psi0 + Cov + (k0 * n / kn) * np.outer(xbar - m0, xbar - m0)
    df = nun - D + 1
    return df, mn, Psin * (kn + 1) / (kn * df)


def _antoniak(n, w, rng):
    """number of tables when n customers sit in a CRP with mass w (= alpha*beta_k):
    m = sum_{c=0}^{n-1} Bernoulli( w / (w + c) )."""
    if n == 0:
        return 0
    c = np.arange(n)
    return int((rng.random(n) < w / (w + c)).sum())


# --------------------------------------------------------------------------- #
#  Direct-assignment HDP Gibbs sampler                                          #
# --------------------------------------------------------------------------- #

def hdp_gibbs(Y, g, rng, alpha=1.0, gamma=1.0, m0=None, k0=0.1, nu0=None, Psi0=None,
              draws=2000, burn=1000, update_conc=True):
    """Direct-assignment HDP Gibbs for grouped multivariate-normal data.
    Y:(N,D) standardised internally; g:(N,) integer group labels 0..J-1.
    Returns the posterior of the number of shared clusters K, a representative global label
    per observation, and the per-group cluster-weight matrix (groups x clusters)."""
    Y = np.asarray(Y, float); g = np.asarray(g, int); N, D = Y.shape
    J = g.max() + 1
    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)
    # initialise: one shared cluster
    z = np.zeros(N, int)
    Nk = [float(N)]; sm = [Ys.sum(0)]; SS = [Ys.T @ Ys]
    njk = [np.bincount(g, minlength=J).astype(float)]      # per-group counts, one per cluster (list of length-J arrays)
    beta = np.array([0.5, 0.5])                            # [active..., unused]
    Ks = []; nsave = 0
    a_c, b_c = 1.0, 1.0                                    # Gamma hyperpriors for alpha, gamma

    def predict(k, yi):
        return _pred(Nk[k], sm[k], SS[k], m0, k0, nu0, Psi0, D)

    for it in range(draws + burn):
        for i in range(N):
            yi = Ys[i]; j = g[i]; k = z[i]
            Nk[k] -= 1; sm[k] -= yi; SS[k] -= np.outer(yi, yi); njk[k][j] -= 1
            if Nk[k] == 0:                                 # cluster died -> remove, fold mass into unused
                beta[-1] += beta[k]
                del Nk[k]; del sm[k]; del SS[k]; del njk[k]
                beta = np.delete(beta, k); z[z > k] -= 1
            K = len(Nk); logp = np.empty(K + 1)
            for kk in range(K):
                df, loc, sc = _pred(Nk[kk], sm[kk], SS[kk], m0, k0, nu0, Psi0, D)
                logp[kk] = np.log(njk[kk][j] + alpha * beta[kk]) + _mvt_logpdf(yi, df, loc, sc)[0]
            df, loc, sc = _pred(0, None, None, m0, k0, nu0, Psi0, D)
            logp[K] = np.log(alpha * beta[-1] + 1e-300) + _mvt_logpdf(yi, df, loc, sc)[0]
            logp -= logp.max(); p = np.exp(logp); p /= p.sum()
            knew = rng.choice(K + 1, p=p)
            if knew == K:                                  # new shared cluster: break the unused stick
                b = rng.beta(1.0, gamma); bu = beta[-1]
                beta = np.concatenate([beta[:-1], [bu * b], [bu * (1 - b)]])
                Nk.append(0.0); sm.append(np.zeros(D)); SS.append(np.zeros((D, D))); njk.append(np.zeros(J))
            z[i] = knew; Nk[knew] += 1; sm[knew] += yi; SS[knew] += np.outer(yi, yi); njk[knew][j] += 1
        # ---- resample global weights beta from table counts (Antoniak) ----
        K = len(Nk); m = np.zeros(K)
        for kk in range(K):
            for j in range(J):
                m[kk] += _antoniak(int(round(njk[kk][j])), alpha * beta[kk], rng)
        beta = rng.dirichlet(np.concatenate([m, [gamma]]))
        # ---- concentration updates (Teh 2006 appendix) ----
        if update_conc:
            mtot = m.sum(); nj = np.bincount(g, minlength=J).astype(float)
            # gamma via Escobar-West on (total tables, #clusters)
            eta = rng.beta(gamma + 1, mtot); pi = (a_c + K - 1) / (a_c + K - 1 + mtot * (b_c - np.log(eta)))
            gamma = (rng.gamma(a_c + K, 1 / (b_c - np.log(eta))) if rng.random() < pi
                     else rng.gamma(a_c + K - 1, 1 / (b_c - np.log(eta))))
            # alpha via auxiliary w_j, s_j
            wj = rng.beta(alpha + 1, nj); sj = (rng.random(J) < nj / (nj + alpha)).astype(float)
            alpha = rng.gamma(a_c + mtot - sj.sum(), 1 / (b_c - np.log(wj).sum()))
        if it >= burn:
            Ks.append(len(Nk)); nsave += 1
            if it == draws + burn - 1:
                W = np.array([njk[k] for k in range(len(Nk))]).T   # (J,K)
                W = W / W.sum(1, keepdims=True)
    return dict(K=np.array(Ks), z=z.copy(), group_weights=W, mu=mu_, sd=sd_,
                cluster_counts=np.array(Nk), alpha=alpha, gamma=gamma)


# --------------------------------------------------------------------------- #
#  simulate grouped data that shares a global pool of clusters                  #
# --------------------------------------------------------------------------- #

def simulate_grouped(group_weights, means, covs, n_per_group, rng):
    """group_weights:(J,K) mixing weights of each group over the SHARED components; means/covs
    are the K shared Gaussian components. Returns Y:(N,D), g:(N,), z_true:(N,)."""
    means = np.asarray(means); covs = np.asarray(covs); J, K = np.asarray(group_weights).shape
    Y = []; g = []; zt = []
    for j in range(J):
        w = np.asarray(group_weights[j], float); w = w / w.sum()
        ks = rng.choice(K, size=n_per_group, p=w)
        for k in ks:
            Y.append(rng.multivariate_normal(means[k], covs[k])); g.append(j); zt.append(k)
    return np.array(Y), np.array(g), np.array(zt)


def adjusted_rand(a, b):
    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

• Teh, Y. W., Jordan, M. I., Beal, M. J. & Blei, D. M. (2006). Hierarchical Dirichlet processes. JASA 101(476), 1566–1581. — the model, and the direct-assignment sampler implemented here
• Antoniak, C. E. (1974). Mixtures of Dirichlet processes with applications to Bayesian nonparametric problems. Annals of Statistics 2(6), 1152–1174. — the table-count distribution used to update the global weights
• Teh, Y. W. & Jordan, M. I. (2010). Hierarchical Bayesian nonparametric models with applications. In Bayesian Nonparametrics, Cambridge University Press. — the Chinese-restaurant-franchise picture
• Horst, A. M., Hill, A. P. & Gorman, K. B. (2020). palmerpenguins: Palmer Archipelago penguin data. — the grouped dataset used here
• Hubert, L. & Arabie, P. (1985). Comparing partitions. Journal of Classification 2, 193–218. — the adjusted Rand index
• Scrucca, L. et al. (2016). mclust 5: clustering, classification and density estimation. The R Journal 8(1), 289–317. — the per-group and pooled EM baselines