""" 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