""" Regenerates every figure in EM_MIXTURES_claude_1015.tex into images_claude/. Run from this directory: python3 build_em_figs.py Needs numpy, matplotlib, scipy, scikit-learn (the shared d2l venv has all four). Figures produced: mixture_motivation.png -- one feature whose histogram has three humps; one Gaussian fit to it (wrong) vs a 3-component mixture (right). The "fit a density" picture from the GDA_NB note, now with a density that one bump cannot make. kmeans_vs_em.png -- the same 2-D data assigned hard (k-means) and soft (responsibilities). The on-ramp figure: distance -> probability. em_iterations.png -- real EM on HW4's 2gaussian.txt: fitted contours at iterations 0, 1, 5 and at convergence, plus the observed-data log-likelihood per iteration (monotone -- HW4 Problem 6). coins_recovery.png -- binomial-mixture EM (HW4 Problem 4) on well-separated vs overlapping coins, 40 restarts each: why overlapping coins recover weakly. """ import os import numpy as np import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from scipy.stats import norm from scipy.special import logsumexp HERE = os.path.dirname(os.path.abspath(__file__)) OUT = os.path.join(HERE, "images_claude") HW3 = os.path.join(HERE, "..", "..", "HW3") os.makedirs(OUT, exist_ok=True) INK, C0, C1, C2 = "#1A1D21", "#2A6FB5", "#B5452A", "#1E7A5F" plt.rcParams.update({ "font.size": 10, "axes.edgecolor": INK, "axes.labelcolor": INK, "text.color": INK, "xtick.color": INK, "ytick.color": INK, "axes.spines.top": False, "axes.spines.right": False, "figure.dpi": 160, }) # ---------------------------------------------------------------------- # A self-contained 2-D Gaussian-mixture EM -- the same algorithm the note # writes out in pseudocode, so the figures and the text cannot drift apart. # ---------------------------------------------------------------------- def log_gauss(X, mu, Sigma): """log N(x; mu, Sigma) per row of X, via slogdet (never det + log).""" D = mu.shape[0] _, logdet = np.linalg.slogdet(Sigma) diff = X - mu quad = np.einsum("ni,ij,nj->n", diff, np.linalg.inv(Sigma), diff) return -0.5 * (D * np.log(2 * np.pi) + logdet + quad) def em_gmm(X, K, iters=60, seed=0, floor=1e-6, snapshots=()): """Returns (pi, mu, Sigma, loglik_per_iter, {it: (pi,mu,Sigma)}).""" rng = np.random.RandomState(seed) N, D = X.shape pi = np.full(K, 1.0 / K) mu = X[rng.choice(N, K, replace=False)].astype(float) S0 = np.atleast_2d(np.cov(X.T)) # atleast_2d so D=1 still gives a 1x1 matrix Sigma = np.array([S0.copy() for _ in range(K)]) lls, snaps = [], {} for it in range(iters): if it in snapshots: snaps[it] = (pi.copy(), mu.copy(), Sigma.copy()) # ---- E step: r_ik = pi_k N(x_i; mu_k, Sigma_k) / sum_l (...) logw = np.column_stack([np.log(pi[k]) + log_gauss(X, mu[k], Sigma[k]) for k in range(K)]) ll = logsumexp(logw, axis=1) # log p(x_i) per point lls.append(ll.sum()) r = np.exp(logw - ll[:, None]) # ---- M step: weighted mean / weighted covariance / weight share Nk = r.sum(0) pi = Nk / N mu = (r.T @ X) / Nk[:, None] for k in range(K): d = X - mu[k] Sigma[k] = (r[:, k, None] * d).T @ d / Nk[k] + floor * np.eye(D) if iters in snapshots: snaps[iters] = (pi.copy(), mu.copy(), Sigma.copy()) return pi, mu, Sigma, np.array(lls), snaps def draw_ellipses(ax, mu, Sigma, color, lw=1.6, levels=(1.0, 2.0)): """Contours of a fitted Gaussian at 1 and 2 standard deviations.""" vals, vecs = np.linalg.eigh(Sigma) ang = np.degrees(np.arctan2(vecs[1, -1], vecs[0, -1])) from matplotlib.patches import Ellipse for s in levels: w, h = 2 * s * np.sqrt(np.maximum(vals[::-1], 1e-12)) ax.add_patch(Ellipse(mu, w, h, angle=ang, fill=False, edgecolor=color, lw=lw, alpha=0.9)) # ---------------------------------------------------------------------- # Figure 1: one bump is not enough # ---------------------------------------------------------------------- def fig_mixture_motivation(): rng = np.random.RandomState(7) # three sub-populations in one feature -- the instructor's weight-bucket picture x = np.concatenate([ rng.normal(58, 4.0, 300), rng.normal(78, 5.0, 500), rng.normal(102, 6.0, 200), ]) grid = np.linspace(38, 125, 700) single = norm.pdf(grid, x.mean(), x.std()) X = x[:, None] # restarts, keep the highest-likelihood run -- a single random start lands in a # local optimum that merges two of the three humps (note section 9) best = None for s in range(20): out = em_gmm(X, 3, iters=200, seed=s) if best is None or out[3][-1] > best[3][-1]: best = out pi, mu, Sigma = best[0], best[1], best[2] comp = np.array([pi[k] * norm.pdf(grid, mu[k, 0], np.sqrt(Sigma[k, 0, 0])) for k in range(3)]) fig, axes = plt.subplots(1, 2, figsize=(10.6, 3.6), sharey=True) for ax in axes: ax.hist(x, bins=45, density=True, color=C0, alpha=0.20, edgecolor=C0, linewidth=0.6) ax.set_xlabel("weight (kg)") axes[0].plot(grid, single, color=C1, lw=2.2) axes[0].set_title(r"one Gaussian: $p(x)=\mathcal{N}(\mu,\sigma^2)$", fontsize=10.5) axes[0].set_ylabel("density") axes[0].annotate("the peak sits where\nalmost nobody is", xy=(x.mean(), norm.pdf(x.mean(), x.mean(), x.std())), xytext=(40, 0.030), fontsize=9, color=C1, arrowprops=dict(arrowstyle="->", color=C1, lw=1.0)) for k in range(3): axes[1].plot(grid, comp[k], color=C2, lw=1.1, ls="--", alpha=0.85) axes[1].plot(grid, comp.sum(0), color=C2, lw=2.2) axes[1].set_title(r"three Gaussians: $p(x)=\sum_k \pi_k\,\mathcal{N}(\mu_k,\sigma_k^2)$", fontsize=10.5) order = np.argsort(mu[:, 0]) for rank, k in enumerate(order): axes[1].annotate(rf"$\pi_{{{rank+1}}}={pi[k]:.2f}$", xy=(mu[k, 0] - 5, comp[k].max() + 0.0022), fontsize=8.5, color=C2) axes[1].set_ylim(0, 0.046) fig.tight_layout() fig.savefig(os.path.join(OUT, "mixture_motivation.png"), bbox_inches="tight") plt.close(fig) print("wrote mixture_motivation.png") # ---------------------------------------------------------------------- # Figure 2: hard assignment vs soft responsibilities # ---------------------------------------------------------------------- def fig_kmeans_vs_em(): rng = np.random.RandomState(11) # deliberately overlapping, so there is a real band of points the two models # disagree about -- that band is the whole point of the figure X = np.vstack([ rng.multivariate_normal([0, 0], [[1.0, 0.3], [0.3, 1.0]], 150), rng.multivariate_normal([2.6, 1.7], [[1.6, -0.6], [-0.6, 0.8]], 150), ]) pi, mu, Sigma, _, _ = em_gmm(X, 2, iters=120, seed=11) # component 1 = the one nearer the origin, so both panels use the same colours if mu[0, 0] > mu[1, 0]: pi, mu, Sigma = pi[::-1], mu[::-1], Sigma[::-1] logw = np.column_stack([np.log(pi[k]) + log_gauss(X, mu[k], Sigma[k]) for k in range(2)]) r = np.exp(logw - logsumexp(logw, axis=1)[:, None]) # --- plain k-means (hard assignment by distance), initialised at the same means c = mu.copy() for _ in range(50): a = np.argmin(((X[:, None, :] - c[None]) ** 2).sum(-1), axis=1) for k in range(2): if (a == k).any(): c[k] = X[a == k].mean(0) from matplotlib.colors import LinearSegmentedColormap # = 1 must land on component 1's own colour, so the two panels agree cmap = LinearSegmentedColormap.from_list("soft", [C0, "#EDE9E1", C1]) fig, axes = plt.subplots(1, 2, figsize=(10.6, 4.2)) axes[0].scatter(*X.T, s=15, c=[C1 if ai == 0 else C0 for ai in a], linewidths=0, alpha=0.85) # the Voronoi boundary: the perpendicular bisector of the two centres mid, d = c.mean(0), c[1] - c[0] t = np.linspace(-6, 6, 2) axes[0].plot(mid[0] - d[1] * t, mid[1] + d[0] * t, color=INK, lw=1.4, ls="--") axes[0].scatter(*c.T, s=120, c=INK, marker="X", edgecolor="white", linewidths=1.4, zorder=5) axes[0].set_title(r"$k$-means: nearest centre wins, $z_{ik}\in\{0,1\}$", fontsize=10.5) sc = axes[1].scatter(*X.T, s=15, c=r[:, 0], cmap=cmap, vmin=0, vmax=1, linewidths=0, alpha=0.95) for k, col in [(0, C1), (1, C0)]: draw_ellipses(axes[1], mu[k], Sigma[k], col) axes[1].scatter(*mu.T, s=120, c=INK, marker="X", edgecolor="white", linewidths=1.4, zorder=5) # angle brackets match the note: is the expected indicator, i.e. the responsibility axes[1].set_title(r"EM: every point splits itself, " r"$\langle z_{i1}\rangle+\langle z_{i2}\rangle=1$", fontsize=10.5) cb = fig.colorbar(sc, ax=axes[1], fraction=0.046, pad=0.03) cb.set_label(r"$\langle z_{i1}\rangle$", fontsize=10) cb.outline.set_edgecolor(INK) for ax in axes: ax.set_xlabel("$x^1$"); ax.set_ylabel("$x^2$") ax.set_xlim(X[:, 0].min() - 0.6, X[:, 0].max() + 0.6) ax.set_ylim(X[:, 1].min() - 0.6, X[:, 1].max() + 0.6) fig.tight_layout() fig.savefig(os.path.join(OUT, "kmeans_vs_em.png"), bbox_inches="tight") plt.close(fig) print("wrote kmeans_vs_em.png") # ---------------------------------------------------------------------- # Figure 3: EM running on HW4's own data # ---------------------------------------------------------------------- def fig_em_iterations(): X = np.loadtxt(os.path.join(HW3, "2gaussian.txt")) snaps_at = (0, 1, 5, 60) pi, mu, Sigma, lls, snaps = em_gmm(X, 2, iters=60, seed=4, snapshots=snaps_at) fig = plt.figure(figsize=(11.0, 5.6)) gs = fig.add_gridspec(2, 3, width_ratios=[1, 1, 1.15], hspace=0.38, wspace=0.52) panels = [(0, 0), (0, 1), (1, 0), (1, 1)] titles = ["iteration 0 (random init)", "after 1 iteration", "after 5 iterations", "converged (60)"] for (rr, cc), it, title in zip(panels, snaps_at, titles): ax = fig.add_subplot(gs[rr, cc]) ax.scatter(*X.T, s=3, c="#9AA3AB", linewidths=0, alpha=0.5) p, m, S = snaps[it] for k, col in enumerate([C0, C1]): draw_ellipses(ax, m[k], S[k], col, lw=1.5) ax.scatter(*m[k], s=70, c=col, marker="X", edgecolor="white", linewidths=1.2, zorder=5) ax.set_title(f"{title}\n" + r"$\log p(X)=$" + f"{lls[min(it, len(lls)-1)]:,.0f}", fontsize=9) ax.set_xlim(-1, 12); ax.set_ylim(-2, 10) ax.tick_params(labelsize=8) axll = fig.add_subplot(gs[:, 2]) axll.plot(np.arange(len(lls)), lls, color=C2, lw=2.0) axll.scatter([0, 1, 5], lls[[0, 1, 5]], s=36, facecolor="white", edgecolor=C2, zorder=5, linewidths=1.4) axll.set_xlabel("EM iteration") axll.set_title("observed-data log-likelihood " + r"$\sum_i\log p(\mathbf{x}_i)$" + "\nnever decreases — HW4 Problem 6", fontsize=10) axll.annotate("almost all of the gain\nin the first few steps", xy=(4, lls[4]), xytext=(14, lls[0] + 0.55 * (lls[-1] - lls[0])), fontsize=9, color=C2, arrowprops=dict(arrowstyle="->", color=C2, lw=1.0)) fig.savefig(os.path.join(OUT, "em_iterations.png"), bbox_inches="tight") plt.close(fig) print("wrote em_iterations.png (final pi = %s)" % np.array2string(pi, precision=3)) # ---------------------------------------------------------------------- # Figure 4: the coins problem, separated vs overlapping # ---------------------------------------------------------------------- def gen_coins(p, q, pi, n, N, rng): """N sequences of n flips; coin 1 w.p. pi (head-prob p), else coin 2 (head-prob q).""" z = rng.rand(N) < pi heads = np.where(z, rng.binomial(n, p, N), rng.binomial(n, q, N)) return heads def em_binomial(heads, n, iters=300, seed=0): """Mixture of two Binomials. Returns (p, q, pi). C(n,x) cancels in the E step.""" rng = np.random.RandomState(seed) p, q = rng.uniform(0.1, 0.9, 2) pi = rng.uniform(0.2, 0.8) for _ in range(iters): # E step: r_i = posterior that sequence i came from coin 1 l1 = np.log(pi) + heads * np.log(p) + (n - heads) * np.log1p(-p) l2 = np.log1p(-pi) + heads * np.log(q) + (n - heads) * np.log1p(-q) r = 1.0 / (1.0 + np.exp(l2 - l1)) # M step: each coin's head-rate is its own weighted fraction of heads p = np.clip((r * heads).sum() / (n * r.sum()), 1e-6, 1 - 1e-6) q = np.clip(((1 - r) * heads).sum() / (n * (1 - r).sum()), 1e-6, 1 - 1e-6) pi = np.clip(r.mean(), 1e-6, 1 - 1e-6) return p, q, pi def fig_coins_recovery(): n, N, restarts = 10, 400, 40 settings = [(0.80, 0.30, 0.60, "well separated: $p=0.80$, $q=0.30$"), (0.60, 0.50, 0.60, "overlapping: $p=0.60$, $q=0.50$")] fig, axes = plt.subplots(1, 2, figsize=(10.6, 3.9), sharey=True) for ax, (p, q, pi, title) in zip(axes, settings): rng = np.random.RandomState(5) heads = gen_coins(p, q, pi, n, N, rng) ph, qh = [], [] for s in range(restarts): pp, qq, _ = em_binomial(heads, n, seed=s) # label switching: the run may return the two coins in either order, # so canonicalise by head-rate before plotting (note section 13) lo, hi = sorted((pp, qq)) qh.append(lo); ph.append(hi) ax.scatter(ph, qh, s=28, facecolor="white", edgecolor=C0, linewidths=1.2, alpha=0.8, label="EM run (40 restarts)") ax.scatter([max(p, q)], [min(p, q)], s=150, c=C1, marker="*", edgecolor="white", linewidths=0.8, zorder=5, label="truth") ax.plot([0, 1], [0, 1], color=INK, lw=0.7, ls=":", alpha=0.5) ax.set_xlim(0.35, 1.0); ax.set_ylim(0.0, 0.75) ax.set_xlabel("recovered head-rate of the more-biased coin") ax.set_title(title, fontsize=10.5) axes[0].set_ylabel("recovered head-rate\nof the other coin") axes[0].legend(frameon=False, fontsize=9, loc="upper left") axes[1].annotate("runs scatter — the data\nbarely distinguishes the coins", xy=(0.63, 0.47), xytext=(0.70, 0.12), fontsize=9, color=C1, arrowprops=dict(arrowstyle="->", color=C1, lw=1.0)) fig.tight_layout() fig.savefig(os.path.join(OUT, "coins_recovery.png"), bbox_inches="tight") plt.close(fig) print("wrote coins_recovery.png") if __name__ == "__main__": fig_mixture_motivation() fig_kmeans_vs_em() fig_em_iterations() fig_coins_recovery() print("all figures written to", OUT)