import os
os.environ["OMP_NUM_THREADS"] = "1"     # silences scikit-learn's KMeans memory-leak warning on Windows
import numpy as np, pandas as pd, matplotlib.pyplot as plt
from sklearn.mixture import GaussianMixture
import hdp as H
rng = np.random.default_rng(1)
BLUE="#2b6cb0"; RED="#c53030"; GREEN="#2f855a"; ORANGE="#dd6b20"; PURP="#6b46c1"; GREY="#718096"
CLS=[BLUE,RED,GREEN,ORANGE,PURP,GREY]
print("HDP = two stacked Dirichlet processes; groups share one global menu of clusters.")

HDP = two stacked Dirichlet processes; groups share one global menu of clusters.

means=[[0,0],[6,6],[12,0]]; covs=[np.eye(2)*0.7]*3
gw=[[0.7,0.3,0.0],[0.0,0.5,0.5],[0.4,0.3,0.3]]
Ys,gs,zt = H.simulate_grouped(gw, means, covs, 130, rng)
sim = H.hdp_gibbs(Ys, gs, rng, draws=1000, burn=800)
print(f"recovered #shared clusters: mode {np.bincount(sim['K']).argmax()} (true 3), ARI {H.adjusted_rand(zt,sim['z']):.3f}")
fig,ax=plt.subplots(1,2,figsize=(12,4.4))
for j in range(3):
    m=gs==j; ax[0].scatter(Ys[m,0],Ys[m,1],s=14,color=CLS[j],alpha=.5,label=f"group {j}")
ax[0].set_title("Three groups (colour = group)"); ax[0].legend(frameon=False); ax[0].set_xlabel("x1"); ax[0].set_ylabel("x2")
for k in np.unique(sim['z']):
    m=sim['z']==k; ax[1].scatter(Ys[m,0],Ys[m,1],s=14,color=CLS[k%len(CLS)],alpha=.6)
ax[1].set_title("HDP shared clusters (colour = inferred cluster)"); ax[1].set_xlabel("x1"); ax[1].set_ylabel("x2")
plt.tight_layout(); plt.show()
print("Per-group weights over the shared clusters (rows = groups):")
print(np.round(sim['group_weights'],2))
print("The HDP finds the three shared components and reconstructs each group's mix -- including the zeros where a")
print("group never uses a component. A separate DP per group could not tell that the groups share components at all.")

recovered #shared clusters: mode 3 (true 3), ARI 1.000

Per-group weights over the shared clusters (rows = groups):
[[0.   0.3  0.7 ]
 [0.47 0.53 0.  ]
 [0.32 0.25 0.42]]
The HDP finds the three shared components and reconstructs each group's mix -- including the zeros where a
group never uses a component. A separate DP per group could not tell that the groups share components at all.

d = pd.read_csv("penguins.csv")
feat = ["bill_length_mm","bill_depth_mm","flipper_length_mm","body_mass_g"]
Y = d[feat].to_numpy()
islands = pd.Categorical(d["island"]); isl = islands.codes; iname = list(islands.categories)
species = pd.Categorical(d["species"]); sp = species.codes; sname = list(species.categories)
fit = H.hdp_gibbs(Y, isl, rng, draws=1500, burn=1000)
Kmode = np.bincount(fit["K"]).argmax()
print(f"HDP shared clusters: mode {Kmode}, mean {fit['K'].mean():.2f}   ARI vs true species: {H.adjusted_rand(sp,fit['z']):.3f}")
print(f"concentrations: alpha {fit['alpha']:.2f} (within-island), gamma {fit['gamma']:.2f} (global)")

zc = fit["z"]; ucl = np.unique(zc)
fig,ax=plt.subplots(1,2,figsize=(13,4.7))
mk = ["o","s","^"]
for j,isln in enumerate(iname):
    for k in ucl:
        m=(isl==j)&(zc==k)
        if m.sum(): ax[0].scatter(d["bill_length_mm"][m], d["flipper_length_mm"][m], s=26, marker=mk[j],
                                  color=CLS[list(ucl).index(k)%len(CLS)], alpha=.75, edgecolor="none")
ax[0].set_xlabel("bill length (mm)"); ax[0].set_ylabel("flipper length (mm)")
ax[0].set_title("HDP clusters (colour) × island (marker)")
ax[0].legend([plt.Line2D([],[],marker=mk[j],ls="",color="k") for j in range(3)], iname, frameon=False, fontsize=8)
# per-island cluster weight heatmap
W = fit["group_weights"]
im=ax[1].imshow(W, cmap="Blues", vmin=0, vmax=1, aspect="auto")
for j in range(W.shape[0]):
    for k in range(W.shape[1]): ax[1].text(k,j,f"{W[j,k]:.2f}",ha="center",va="center",fontsize=8)
ax[1].set_yticks(range(len(iname))); ax[1].set_yticklabels(iname); ax[1].set_xlabel("shared cluster")
ax[1].set_title("Per-island weights over shared clusters"); ax[1].set_xticks(range(W.shape[1]))
plt.colorbar(im,ax=ax[1],fraction=.046); plt.tight_layout(); plt.show()
shared=[k for k in range(W.shape[1]) if (W[:,k]>0.05).sum()>=2]
print(f"cluster(s) present on 2+ islands (shared): {shared} -- this is the Adelie cluster, found on every island.")
print("Gentoo (Biscoe-only) and Chinstrap (Dream-only) each load a single island. The HDP recovers the species")
print(f"almost perfectly (ARI {H.adjusted_rand(sp,fit['z']):.2f}) AND the cross-island sharing structure.")

HDP shared clusters: mode 4, mean 3.89   ARI vs true species: 0.998
concentrations: alpha 0.17 (within-island), gamma 1.07 (global)

cluster(s) present on 2+ islands (shared): [1] -- this is the Adelie cluster, found on every island.
Gentoo (Biscoe-only) and Chinstrap (Dream-only) each load a single island. The HDP recovers the species
almost perfectly (ARI 1.00) AND the cross-island sharing structure.

print("Independent per-island Gaussian mixtures (BIC-selected K):")
Ys_all=(Y-Y.mean(0))/Y.std(0)
for j,isln in enumerate(iname):
    Yj=Ys_all[isl==j]
    bic=[GaussianMixture(k,covariance_type='full',random_state=0).fit(Yj).bic(Yj) for k in range(1,5)]
    kbest=np.argmin(bic)+1
    print(f"  {isln:10s}: BIC picks K={kbest}  (its clusters have island-local labels, unlinked to other islands)")
# pooled mixture ignores the grouping
pool=GaussianMixture(3,covariance_type='full',random_state=0).fit(Ys_all)
print(f"\nPooled mixture (grouping discarded), K=3: ARI vs species {H.adjusted_rand(sp,pool.predict(Ys_all)):.3f}")
print(f"HDP (shares clusters, keeps grouping):        ARI vs species {H.adjusted_rand(sp,fit['z']):.3f}")
print("\nIndependent fits can't say Torgersen's Adelie = Biscoe's Adelie; the pooled fit forgets which island a bird")
print("came from. Only the HDP does both -- one shared set of species AND island-specific proportions.")

Independent per-island Gaussian mixtures (BIC-selected K):
  Biscoe    : BIC picks K=2  (its clusters have island-local labels, unlinked to other islands)
  Dream     : BIC picks K=2  (its clusters have island-local labels, unlinked to other islands)
  Torgersen : BIC picks K=1  (its clusters have island-local labels, unlinked to other islands)

Pooled mixture (grouping discarded), K=3: ARI vs species 0.960
HDP (shares clusters, keeps grouping):        ARI vs species 0.998

Independent fits can't say Torgersen's Adelie = Biscoe's Adelie; the pooled fit forgets which island a bird
came from. Only the HDP does both -- one shared set of species AND island-specific proportions.

# --- Fig: why "cluster 1" means nothing across independently-fitted islands ---
loc_lab = np.full(len(Y), -1); loc_key = []
for j, isln in enumerate(iname):
    Yj = Ys_all[isl == j]
    bic = [GaussianMixture(k, covariance_type="full", random_state=0).fit(Yj).bic(Yj) for k in range(1, 5)]
    kb = int(np.argmin(bic)) + 1
    lab = GaussianMixture(kb, covariance_type="full", random_state=0).fit(Yj).predict(Yj)
    loc_lab[isl == j] = len(loc_key) + lab
    loc_key += [f"{isln[:3]}-{c+1}" for c in range(kb)]

bl, fl = d["bill_length_mm"].to_numpy(), d["flipper_length_mm"].to_numpy()
fig, ax = plt.subplots(1, 2, figsize=(13, 4.7))
mk = ["o", "s", "^"]

# left: independent fits -- every island gets its own, unrelated label set
for j in range(3):
    for c in np.unique(loc_lab[isl == j]):
        m = (isl == j) & (loc_lab == c)
        ax[0].scatter(bl[m], fl[m], s=28, marker=mk[j], alpha=.8, edgecolor="none",
                      color=CLS[c % len(CLS)])
        ax[0].annotate(loc_key[c], (bl[m].mean(), fl[m].mean()), fontsize=8, fontweight="bold",
                       ha="center", textcoords="offset points", xytext=(-40 + 40 * j, 16 - 13 * j),
                       arrowprops=dict(arrowstyle="-", lw=.7, color="0.5"),
                       bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="0.7", alpha=.9))
ax[0].set_title(f"Independent per-island mixtures — {len(loc_key)} unrelated labels")
ax[0].set_xlabel("bill length (mm)"); ax[0].set_ylabel("flipper length (mm)")

# right: the HDP -- one shared label set spanning islands
for j in range(3):
    for k in np.unique(fit["z"]):
        m = (isl == j) & (fit["z"] == k)
        if m.sum(): ax[1].scatter(bl[m], fl[m], s=28, marker=mk[j], alpha=.8, edgecolor="none",
                                  color=CLS[list(np.unique(fit["z"])).index(k) % len(CLS)])
for k in np.unique(fit["z"]):
    m = fit["z"] == k
    if m.sum() > 5:
        ax[1].annotate(f"shared {k}", (bl[m].mean(), fl[m].mean()), fontsize=8, fontweight="bold",
                       ha="center", bbox=dict(boxstyle="round,pad=0.2", fc="white", ec="0.7", alpha=.85))
ax[1].set_title("HDP — one shared menu, island-specific proportions")
ax[1].set_xlabel("bill length (mm)"); ax[1].set_ylabel("flipper length (mm)")
for a in ax:
    a.legend([plt.Line2D([], [], marker=mk[j], ls="", color="k") for j in range(3)], iname,
             frameon=False, fontsize=8, loc="lower right")
    a.grid(alpha=.25)
plt.tight_layout(); plt.show()

print(f"Left: {len(loc_key)} labels ({', '.join(loc_key)}) and no way to connect them. Torgersen's birds and")
print("Biscoe's may be the same species or different ones -- three separate fits cannot express the question,")
print("let alone answer it. The colours repeat across panels only because a palette was reused; nothing in the")
print("left-hand model says a Biscoe cluster and a Dream cluster are related.")
print(f"Right: one label set spanning all three islands, ARI {H.adjusted_rand(sp, fit['z']):.3f} against species.")
print("That is the difference the HDP buys, and it is a difference in what can be ASKED, not in fit quality --")
print(f"the pooled mixture scores {H.adjusted_rand(sp, pool.predict(Ys_all)):.3f}, marginally better, by throwing the islands away.")

Left: 5 labels (Bis-1, Bis-2, Dre-1, Dre-2, Tor-1) and no way to connect them. Torgersen's birds and
Biscoe's may be the same species or different ones -- three separate fits cannot express the question,
let alone answer it. The colours repeat across panels only because a palette was reused; nothing in the
left-hand model says a Biscoe cluster and a Dream cluster are related.
Right: one label set spanning all three islands, ARI 0.998 against species.
That is the difference the HDP buys, and it is a difference in what can be ASKED, not in fit quality --
the pooled mixture scores 0.960, marginally better, by throwing the islands away.

import pymc as pm, pytensor.tensor as pt
Ys = (Y-Y.mean(0))/Y.std(0); N,Dd = Ys.shape; J=3; Kmax=6
with pm.Model() as mod:
    gamma = pm.Gamma("gamma", 2, 2); alpha = pm.Gamma("alpha", 2, 2)
    v  = pm.Beta("v", 1, gamma, shape=Kmax)
    beta = pm.Deterministic("beta", v*pt.concatenate([[1.0], pt.cumprod(1-v)[:-1]]))
    beta = beta/pt.sum(beta)
    pi = pm.Dirichlet("pi", a=alpha*beta+1e-6, shape=(J,Kmax))       # per-island weights reuse global beta
    mu = pm.Normal("mu", 0, 2, shape=(Kmax,Dd),
                   initval=np.repeat(np.linspace(-1.5,1.5,Kmax)[:,None],Dd,1))
    chol,_,_ = pm.LKJCholeskyCov("chol", n=Dd, eta=2.0, sd_dist=pm.HalfNormal.dist(1.0), compute_corr=True)
    comps = pm.MvNormal.dist(mu, chol=chol)
    w_obs = pi[isl]                                                  # (N,Kmax) each bird's island weights
    pm.Mixture("obs", w=w_obs, comp_dists=comps, observed=Ys)
    idata = pm.sample(700, tune=1200, chains=4, target_accept=0.9, random_seed=3,
                      init="adapt_diag", progressbar=False)

# components label-switch across draws, so averaging mu is meaningless -- read the clustering
# from a SINGLE best-fitting draw (highest mixture log-likelihood among a thinned set).
from scipy.stats import multivariate_normal as mvn
from scipy.special import logsumexp
post = idata.posterior; nc, nd = post.sizes["chain"], post.sizes["draw"]
def _params(ci,di):
    pi1=post["pi"].values[ci,di]; mu1=post["mu"].values[ci,di]
    st=post["chol_stds"].values[ci,di]; co=post["chol_corr"].values[ci,di]
    return pi1, mu1, np.diag(st)@co@np.diag(st)
best=None
for ci in range(nc):
    for di in range(0, nd, max(1, nd//25)):
        pi1,mu1,C1=_params(ci,di)
        L=np.column_stack([np.log(pi1[isl,k]+1e-12)+mvn.logpdf(Ys,mu1[k],C1,allow_singular=True) for k in range(Kmax)])
        ll=logsumexp(L,axis=1).sum()
        if best is None or ll>best[0]: best=(ll,pi1,L)
_,pi1,L=best; zpm=L.argmax(1); active=np.unique(zpm)
print(f"PyMC truncated HDP (best-fitting draw): {len(active)} active shared clusters (from-scratch mode {Kmode})")
ari_pm = H.adjusted_rand(sp, zpm); ari_fs = H.adjusted_rand(sp, fit["z"])
print("PyMC clustering ARI vs species: %.3f   (from-scratch %.3f)" % (ari_pm, ari_fs))
Wpm = pi1[:,active]; Wpm = Wpm/Wpm.sum(1,keepdims=True)
print("PyMC per-island weights over active shared clusters (rows=islands):")
print(np.round(Wpm,2))
ndiv = int(idata.sample_stats["diverging"].values.sum())
print("\nSampler diagnostics: %d divergences. R-hat and ESS on the component parameters are not usable" % ndiv)
print("here -- a mixture is invariant to relabelling its clusters, so chains exploring different labellings")
print("of the same fit look unconverged by construction. Judge this cross-check on the clustering, not on")
print("the per-parameter diagnostics.\n")
print("The qualitative story does carry over: a few shared clusters, one of them present on every island.")
print("The QUANTITATIVE agreement does not, and the gap is worth naming rather than softening. ARI %.3f"
      % ari_pm)
print("against the from-scratch sampler's %.3f is a substantial drop, not a rounding difference -- this" % ari_fs)
print("truncated model shares ONE covariance across all components, so clusters that differ in shape can")
print("only be covered by splitting them, which is exactly what the extra active cluster is doing. The")
print("lesson is the same one the DP-mixture project ends on: the cluster count and the partition are")
print("properties of the model you chose, and two correct implementations of 'an HDP' need not agree on them.")

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

Initializing NUTS using adapt_diag...

Multiprocess sampling (4 chains in 4 jobs)

NUTS: [gamma, alpha, v, pi, mu, chol]

Sampling 4 chains for 1_200 tune and 700 draw iterations (4_800 + 2_800 draws total) took 196 seconds.

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

Chain 0 reached the maximum tree depth. Increase `max_treedepth`, increase `target_accept` or reparameterize.

Chain 1 reached the maximum tree depth. Increase `max_treedepth`, increase `target_accept` or reparameterize.

Chain 3 reached the maximum tree depth. Increase `max_treedepth`, 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 Hierarchical Dirichlet Process¶

1. Does it work? — recovering shared clusters across groups¶

2. The penguins — species shared across islands¶

4. Cross-check in PyMC — a truncated HDP¶

5. Summary¶

The Hierarchical Dirichlet Process¶

Sharing clusters across groups — one menu, many restaurants¶

1. Does it work? — recovering shared clusters across groups¶

2. The penguins — species shared across islands¶

3. The frequentist counterpart — per-group mixtures can't share¶

4. Cross-check in PyMC — a truncated HDP¶

5. Summary¶