""" ============================================================================= CS236 · 第 5 讲共享模块：toy2d.py —— 密度已知的二维玩具分布 ============================================================================= 【这个模块是干什么的】 像素世界（第 1 讲）是离散的；从第 5 讲开始，很多模型（高斯混合、VAE、流模型、GAN、得分模型） 处理的是**连续**数据。为了能在终端里"看见"分布，并且能把模型和真相精确比较，这里提供几个二维分布： 每个都能 (1) 采样，(2) 精确计算 log p(x)。 GaussianMixture(weights, means, covs) 任意维数的高斯混合（全协方差）：采样、log p(x)、后验 p(k | x) ring8() 8 个高斯排成一圈（半径 2）——"模式很多、彼此分开" blobs3() 3 个形状不同的高斯团——聚类的例子 moons() 两个交错的半月牙。这里把它**定义成**沿两条弧线排开的 200 个小高斯的混合，所以密度精确已知 Banana() 香蕉形：x₁ = z₁，x₂ = 0.5 z₂ + 0.4 (z₁² − 1)——一个三角形的可逆变换（第 7 讲的主角） DATASETS 名字 → (构造函数, 画图的 x 范围, y 范围) 以及两个工具： density_map(log_prob_fn, xlim, ylim) 把一个**密度函数**画成字符图（text_plot 画的是**样本**） kl_mc(p, log_q_fn, n, generator) 用 p 的 n 个样本估计 KL(p ‖ q) = E_p[log p − log q] 及其标准误 后面的讲次（第 6–10、12–14、16 讲）会只读地导入它。 【运行方式】 python3 toy2d.py # 自检：数值积分 ∫p = 1、样本的均值 vs 解析均值、画出 4 个分布（约 2 秒） 【PyTorch 小知识】 - torch.linalg.cholesky(Σ)：把正定矩阵分解成 Σ = L Lᵀ（L 下三角）。采样时 x = μ + L ε（ε ~ N(0, I)） 的协方差正好是 L Lᵀ = Σ；求密度时 log|Σ| = 2 Σ log L_ii，马氏距离用三角方程求解 L⁻¹(x − μ)。 - torch.linalg.solve_triangular(L, B, upper=False)：解下三角方程 L Y = B，比求逆矩阵快又稳。 """ import math import torch LOG2PI = math.log(2 * math.pi) class GaussianMixture: """ 高斯混合 p(x) = Σ_k π_k N(x; μ_k, Σ_k)，D 维、全协方差。既用来**造数据**（真实分布），也用作**模型**（第 5 讲的 EM）。 【参数】 weights：(K,) 混合权重 π，非负、和为 1 means：(K, D) 各成分的均值 μ_k covs：(K, D, D) 协方差矩阵（必须正定），或 (K,) / 标量——表示各向同性的标准差 σ_k（Σ_k = σ_k² I） 【属性】K、D、weights、means、covs、L（Cholesky 因子） 【边界】协方差不正定时 cholesky 会报错；权重之和偏离 1 超过 1e-9 时报错。 【DGM 联系】"浅层"隐变量模型：z ~ Categorical(π)，x | z = k ~ N(μ_k, Σ_k)（原课第 5 讲 slide 7）。 """ def __init__(self, weights, means, covs): self.weights = torch.as_tensor(weights, dtype=torch.float64) self.means = torch.as_tensor(means, dtype=torch.float64) self.K, self.D = self.means.shape covs = torch.as_tensor(covs, dtype=torch.float64) if covs.dim() <= 1: # 各向同性：给的是标准差 s = covs.expand(self.K) if covs.dim() == 1 else covs.repeat(self.K) covs = (s ** 2).reshape(-1, 1, 1) * torch.eye(self.D, dtype=torch.float64) if abs(self.weights.sum().item() - 1.0) > 1e-9 or (self.weights < 0).any(): raise ValueError("weights 必须非负、和为 1") self.covs = covs self.L = torch.linalg.cholesky(covs) # (K, D, D) def log_joint(self, X): """ log p(x, z = k) = log π_k + log N(x; μ_k, Σ_k)，形状 (N, K)。 【思路】对每个成分：r = L_k⁻¹ (x − μ_k)（解三角方程），马氏距离 = ‖r‖²； log N = −½‖r‖² − ½ D log 2π − Σ log diag(L_k)。 """ X = torch.as_tensor(X, dtype=torch.float64) diff = X.unsqueeze(0) - self.means.unsqueeze(1) # (K, N, D) r = torch.linalg.solve_triangular(self.L, diff.transpose(1, 2), upper=False) # (K, D, N) maha = (r ** 2).sum(dim=1) # (K, N) logdet = 2 * torch.log(torch.diagonal(self.L, dim1=1, dim2=2)).sum(dim=1) # (K,) logN = -0.5 * maha - 0.5 * self.D * LOG2PI - 0.5 * logdet.unsqueeze(1) return (logN + self.weights.log().unsqueeze(1)).T # (N, K) def log_prob(self, X): """精确的 log p(x) = logsumexp_k log p(x, k)，形状 (N,)。""" return torch.logsumexp(self.log_joint(X), dim=1) def posterior(self, X): """后验（"责任"）p(z = k | x)，形状 (N, K)，每行和为 1。""" return torch.softmax(self.log_joint(X), dim=1) def sample(self, n, generator, return_z=False): """ 祖先采样：先 z ~ Categorical(π)，再 x = μ_z + L_z ε。 【返回】(n, D) 张量；return_z=True 时同时返回成分编号 (n,) """ z = torch.multinomial(self.weights, n, replacement=True, generator=generator) eps = torch.randn(n, self.D, 1, generator=generator, dtype=torch.float64) x = self.means[z] + (self.L[z] @ eps).squeeze(-1) return (x, z) if return_z else x def n_params(self): """自由参数个数：(K − 1) 个权重 + K·D 个均值 + K·D(D+1)/2 个协方差元素。""" return (self.K - 1) + self.K * self.D + self.K * self.D * (self.D + 1) // 2 # --------------------------------------------------------------------------- # 几个具体的二维分布 # --------------------------------------------------------------------------- def ring8(radius=2.0, std=0.25): """8 个等权的各向同性高斯，均匀排在半径 radius 的圆上。""" ang = torch.arange(8, dtype=torch.float64) * (2 * math.pi / 8) means = torch.stack([radius * ang.cos(), radius * ang.sin()], dim=1) return GaussianMixture(torch.full((8,), 1 / 8, dtype=torch.float64), means, torch.full((8,), std, dtype=torch.float64)) def blobs3(): """3 个形状、大小、权重都不同的高斯团（用于聚类演示：一个圆、一个斜长条、一个小而密的团）。""" weights = torch.tensor([0.45, 0.35, 0.20], dtype=torch.float64) means = torch.tensor([[-1.6, 0.9], [1.3, 0.6], [0.0, -1.5]], dtype=torch.float64) covs = torch.tensor([[[0.30, 0.00], [0.00, 0.30]], [[0.50, 0.38], [0.38, 0.40]], [[0.06, 0.00], [0.00, 0.06]]], dtype=torch.float64) return GaussianMixture(weights, means, covs) def moons(n_per_arc=100, noise=0.15): """ 两个交错的半月牙（类似 sklearn 的 make_moons），定义成 2 × n_per_arc 个等权小高斯的混合。 【思路】上弧：(cos t, sin t)，t ∈ [0, π]；下弧：(1 − cos t, 0.5 − sin t)。平移、放大 1.5 倍使中心在原点附近。 沿弧线等距放 n_per_arc 个中心，每个中心是一个标准差 noise 的高斯。相邻中心的间距（≈ 0.047）远小于 noise， 所以密度沿弧线是平滑的——它和"均匀地在弧上取点再加高斯噪声"几乎一样，但密度可以精确计算。 """ t = (torch.arange(n_per_arc, dtype=torch.float64) + 0.5) * (math.pi / n_per_arc) upper = torch.stack([t.cos(), t.sin()], dim=1) lower = torch.stack([1 - t.cos(), 0.5 - t.sin()], dim=1) means = (torch.cat([upper, lower]) - torch.tensor([0.5, 0.25], dtype=torch.float64)) * 1.5 K = 2 * n_per_arc return GaussianMixture(torch.full((K,), 1 / K, dtype=torch.float64), means, torch.full((K,), noise, dtype=torch.float64)) class Banana: """ 香蕉形分布：z ~ N(0, I)，x₁ = z₁，x₂ = s·z₂ + b·(z₁² − 1)。 【思路】这是一个**三角形**的可逆变换：x₁ 只依赖 z₁，x₂ 依赖 (z₁, z₂)。反过来 z₁ = x₁，z₂ = (x₂ − b(x₁² − 1)) / s， 雅可比矩阵是下三角的，行列式 = 1 · s。由变量替换公式（第 7 讲）： log p(x) = log N(x₁; 0, 1) + log N(z₂; 0, 1) − log s 【参数】s：x₂ 方向的噪声尺度；b：弯曲程度 """ def __init__(self, s=0.5, b=0.4): self.s, self.b, self.D = s, b, 2 def sample(self, n, generator): z = torch.randn(n, 2, generator=generator, dtype=torch.float64) return torch.stack([z[:, 0], self.s * z[:, 1] + self.b * (z[:, 0] ** 2 - 1)], dim=1) def log_prob(self, X): X = torch.as_tensor(X, dtype=torch.float64) z2 = (X[:, 1] - self.b * (X[:, 0] ** 2 - 1)) / self.s return -0.5 * (X[:, 0] ** 2 + z2 ** 2) - LOG2PI - math.log(self.s) DATASETS = { "ring8": (ring8, (-3.2, 3.2), (-3.2, 3.2)), "blobs3": (blobs3, (-3.5, 3.5), (-2.8, 2.6)), "moons": (moons, (-2.6, 2.6), (-1.6, 1.6)), "banana": (Banana, (-3.5, 3.5), (-1.8, 3.6)), } # --------------------------------------------------------------------------- # 工具 # --------------------------------------------------------------------------- def grid(xlim, ylim, w, h): """ 在 [xlim] × [ylim] 上取 w × h 个格子的中心点。 【返回】(pts, cell_area)：pts 形状 (h·w, 2)，按"从上到下、从左到右"排列（第一行是 y 最大的那一行）； cell_area 是每个格子的面积（数值积分时用：∫f ≈ Σ f(格点) × 面积） """ xs = xlim[0] + (torch.arange(w, dtype=torch.float64) + 0.5) * (xlim[1] - xlim[0]) / w ys = ylim[1] - (torch.arange(h, dtype=torch.float64) + 0.5) * (ylim[1] - ylim[0]) / h Y, X = torch.meshgrid(ys, xs, indexing="ij") area = (xlim[1] - xlim[0]) / w * (ylim[1] - ylim[0]) / h return torch.stack([X.reshape(-1), Y.reshape(-1)], dim=1), area def density_map(log_prob_fn, xlim, ylim, w=56, h=20, title=None): """ 把一个二维**密度函数**画成字符图（每个字符 = 一个格子中心的密度，越黑越大）。 【参数】log_prob_fn：函数 (n, 2) 张量 → (n,) 的 log 密度；xlim、ylim：范围；w、h：字符数 【返回】多行字符串 【思路】和 text_plot.text_density2d 一样用 √(密度 / 最大密度) 选字符，这样"样本图"和"密度图"可以直接对比。 密度低于最大值的 0.2% 的格子留空。 """ shades = " .:-=+*#%@" pts, _ = grid(xlim, ylim, w, h) with torch.no_grad(): p = log_prob_fn(pts).double().exp().reshape(h, w) top = p.max().item() lines = [] if title is None else [" " + title] lines.append(" +" + "-" * w + "+") for r in range(h): row = "" for c in range(w): v = p[r, c].item() / top if top > 0 else 0.0 row += " " if v < 0.002 else shades[min(len(shades) - 1, 1 + int(math.sqrt(v) * (len(shades) - 2)))] lines.append(" |" + row + "|") lines.append(" +" + "-" * w + "+") lines.append(f" x ∈ [{xlim[0]}, {xlim[1]}]，y ∈ [{ylim[0]}, {ylim[1]}]") return "\n".join(lines) def kl_mc(p, log_q_fn, n, generator): """ 蒙特卡洛估计 KL(p ‖ q) = E_{x~p}[log p(x) − log q(x)]（单位 nat），以及它的标准误。 【参数】p：有 sample 和 log_prob 的真实分布；log_q_fn：模型的 log 密度函数；n：样本数 【返回】(估计值, 标准误) 【边界】q 在某个样本处的密度为 0（log q = −∞）时结果是 +∞——前向 KL 惩罚"漏掉的质量"（第 4 讲 §2.3）。 【DGM 联系】这正是第 4 讲的"KL = 交叉熵 − 熵"：只有真实分布已知时才能算；对真实数据只能比较交叉熵。 """ x = p.sample(n, generator) with torch.no_grad(): d = p.log_prob(x) - log_q_fn(x).double() return d.mean().item(), (d.std() / math.sqrt(n)).item() def _selftest(): torch.set_num_threads(1) g = torch.Generator().manual_seed(0) print("自检：每个分布在画图范围外扩 3 倍的网格上数值积分 ∫p(x)dx，应接近 1") for name, (make, xl, yl) in DATASETS.items(): dist = make() big_x = (3 * xl[0], 3 * xl[1]) big_y = (3 * yl[0], 3 * yl[1]) pts, area = grid(big_x, big_y, 600, 600) total = (dist.log_prob(pts).exp().sum() * area).item() x = dist.sample(200000, g) mc_mean = x.mean(0) num_mean = (dist.log_prob(pts).exp().unsqueeze(1) * pts).sum(0) * area print(f" {name:7s} ∫p = {total:.5f} 样本均值 ({mc_mean[0]:+.3f}, {mc_mean[1]:+.3f}) " f"数值积分均值 ({num_mean[0]:+.3f}, {num_mean[1]:+.3f})") for name, (make, xl, yl) in DATASETS.items(): print() print(density_map(make().log_prob, xl, yl, w=48, h=14, title=name)) if __name__ == "__main__": _selftest()