流模型 Normalizing Flow

# 变量变换公式的数值验证
import torch
from torch.autograd import grad

def change_of_variables(f, x, base_dist):
    z, log_det = f.inverse_and_log_det(x)
    log_p_z = base_dist.log_prob(z).sum(dim=-1)
    log_p_x = log_p_z + log_det
    return log_p_x

class SimpleFlow:
    def __init__(self, dim, a=0.5):
        self.a = a  # z = a * x (可逆缩放)

    def inverse_and_log_det(self, x):
        z = x / self.a
        log_det = -x.shape[-1] * torch.log(torch.tensor(abs(self.a)))
        return z, log_det.expand(x.shape[0])

# 验证：1D 缩放变换
flow = SimpleFlow(dim=1, a=2.0)
x = torch.randn(1000, 1)
log_p_x = change_of_variables(flow, x, torch.distributions.Normal(0, 1))
print(f'平均 log p(x): {log_p_x.mean().item():.4f}')

# 雅可比行列式计算复杂度对比
import numpy as np
import time

def benchmark_jacobian_det(dim):
    A_dense = np.random.randn(dim, dim)
    A_tri = np.tril(np.random.randn(dim, dim))  # 下三角

    start = time.time()
    for _ in range(100):
        sign, logdet = np.linalg.slogdet(A_dense)
    dense_time = (time.time() - start) / 100

    start = time.time()
    for _ in range(100):
        tri_logdet = np.sum(np.log(np.abs(np.diag(A_tri))))
    tri_time = (time.time() - start) / 100

    return {'dim': dim, 'dense': f'{dense_time*1000:.3f}ms',
            'triangular': f'{tri_time*1000:.4f}ms'}

for d in [10, 100, 1000]:
    result = benchmark_jacobian_det(d)
    print(f"dim={d:4d} | dense: {result['dense']:>10s} | "
          f"triangular: {result['triangular']}")

# Planar Flow 实现
import torch
import torch.nn as nn
import torch.nn.functional as F

class PlanarFlow(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.w = nn.Parameter(torch.randn(dim))
        self.u = nn.Parameter(torch.randn(dim))
        self.b = nn.Parameter(torch.zeros(1))
        self._ensure_invertibility()

    def _ensure_invertibility(self):
        # 确保 w^T u > -1 (可逆性条件)
        with torch.no_grad():
            wtu = torch.dot(self.w, self.u)
            if wtu <= -1:
                m = -1 + torch.log1p(torch.exp(wtu))
                self.u.data = self.u + (m - wtu) * self.w / self.w.norm()2

    def forward(self, z):
        # 前向变换 + log|det J|
        linear = torch.matmul(z, self.w) + self.b
        h_val = torch.tanh(linear)
        h_prime = 1 - h_val  2  # tanh 导数
        u_hat = self.u + (torch.log1p(torch.exp(
            torch.dot(self.w, self.u))) - 1 - torch.dot(self.w, self.u)) \
            * self.w / self.w.norm()**2
        psi = h_prime * self.w
        log_det = torch.log1p(torch.abs(torch.matmul(psi, u_hat)))
        return z + self.u * h_val.unsqueeze(-1), log_det

# Radial Flow 实现
class RadialFlow(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.z_ref = nn.Parameter(torch.randn(dim))
        self.alpha = nn.Parameter(torch.rand(1))  # > 0
        self.beta = nn.Parameter(torch.randn(1))

    def forward(self, z):
        # 径向变换
        diff = z - self.z_ref
        r = diff.norm(dim=-1, keepdim=True)
        alpha_soft = torch.log1p(torch.exp(self.alpha))  # 保证 > 0
        h_val = 1.0 / (alpha_soft + r)
        h_prime = -1.0 / (alpha_soft + r) ** 2

        # f(z) = z + beta * h(r) * (z - z_ref)
        z_new = z + self.beta * h_val * diff

        # log|det J| = (D-1)*log(1+beta*h) + log(1+beta*h+beta*h'*r)
        term1 = (z.shape[-1] - 1) * torch.log1p(self.beta * h_val).squeeze(-1)
        term2 = torch.log1p(self.beta * h_val.squeeze(-1) + \
                           self.beta * h_prime.squeeze(-1) * r.squeeze(-1))
        log_det = term1 + term2
        return z_new, log_det

# 堆叠多层 Planar Flow
class PlanarFlowStack(nn.Module):
    def __init__(self, dim, n_flows=16):
        super().__init__()
        self.flows = nn.ModuleList([
            PlanarFlow(dim) for _ in range(n_flows)
        ])

    def forward(self, z_0):
        # 正向变换，返回最终输出和总 log|det|
        z = z_0
        total_log_det = z.new_zeros(z.shape[0])

        for flow in self.flows:
            z, log_det = flow(z)
            total_log_det = total_log_det + log_det

        return z, total_log_det

    def inverse(self, x):
        # 逆变换（数值求解）
        z = x
        for flow in reversed(self.flows):
            z = flow.inverse(z)  # 数值迭代求解
        return z

# 可视化 Planar Flow 对 2D 高斯的变换
import matplotlib.pyplot as plt
import numpy as np

def visualize_flow_transformation(flow_stack, n_samples=10000):
    # 采样基础分布
    z_0 = torch.randn(n_samples, 2)

    # 应用 Flow 变换
    z_k, _ = flow_stack(z_0)

    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    axes[0].scatter(z_0[:, 0].numpy(), z_0[:, 1].numpy(),
                    s=1, alpha=0.3)
    axes[0].set_title('Base: N(0, I)')
    axes[0].set_xlim(-4, 4)
    axes[0].set_ylim(-4, 4)

    axes[1].scatter(z_k[:, 0].numpy(), z_k[:, 1].numpy(),
                    s=1, alpha=0.3)
    axes[1].set_title('After Flow')
    axes[1].set_xlim(-4, 4)
    axes[1].set_ylim(-4, 4)

    # 密度估计对比
    from scipy.stats import gaussian_kde
    kde = gaussian_kde(z_k.numpy().T)
    grid_x, grid_y = np.mgrid[-4:4:100j, -4:4:100j]
    grid = np.vstack([grid_x.ravel(), grid_y.ravel()])
    density = kde(grid).reshape(100, 100)
    axes[2].contourf(grid_x, grid_y, density, levels=20, cmap='viridis')
    axes[2].set_title('Estimated Density')
    plt.tight_layout()

# RealNVP 仿射耦合层
class AffineCouplingLayer(nn.Module):
    def __init__(self, dim, mask_type='channel', hidden_dim=256):
        super().__init__()
        self.dim = dim
        self.split_dim = dim // 2
        self.mask_type = mask_type

        # s 和 t 网络
        self.s_net = nn.Sequential(
            nn.Linear(self.split_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, self.split_dim)
        )
        self.t_net = nn.Sequential(
            nn.Linear(self.split_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, self.split_dim)
        )
        # 缩放稳定：tanh 限制输出范围
        self.s_net[-1].weight.data.mul_(0.01)
        self.s_net[-1].bias.data.zero_()

    def forward(self, z):
        # 前向：z -> x
        z_a, z_b = z[:, :self.split_dim], z[:, self.split_dim:]
        s = self.s_net(z_a)
        t = self.t_net(z_a)
        x_b = z_b * torch.exp(s) + t
        log_det = s.sum(dim=-1)
        return torch.cat([z_a, x_b], dim=-1), log_det

    def inverse(self, x):
        # 逆向：x -> z (直接计算，无需迭代)
        x_a, x_b = x[:, :self.split_dim], x[:, self.split_dim:]
        s = self.s_net(x_a)
        t = self.t_net(x_a)
        z_b = (x_b - t) * torch.exp(-s)
        return torch.cat([x_a, z_b], dim=-1)

# RealNVP 完整模型（多耦合层 + 交替掩码）
class RealNVP(nn.Module):
    def __init__(self, dim, n_flows=8, hidden_dim=256):
        super().__init__()
        self.layers = nn.ModuleList()
        for i in range(n_flows):
            self.layers.append(AffineCouplingLayer(
                dim, hidden_dim=hidden_dim
            ))
            # 交替分割：奇数层交换 z_a 和 z_b 的角色
            if i % 2 == 1:
                self.layers.append(ReverseLayer(dim))

    def forward(self, z):
        total_log_det = z.new_zeros(z.shape[0])
        for layer in self.layers:
            z, log_det = layer(z)
            total_log_det += log_det
        return z, total_log_det

    def sample(self, n_samples, base_dist):
        z_0 = base_dist.sample((n_samples,))
        x, _ = self.forward(z_0)
        return x

    def log_prob(self, x):
        z, log_det = self.inverse_and_log_det(x)
        log_p_z = torch.distributions.Normal(0, 1).log_prob(z).sum(-1)
        return log_p_z + log_det

class ReverseLayer(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
    def forward(self, z):
        return torch.flip(z, dims=[-1]), z.new_zeros(z.shape[0])
    def inverse(self, x):
        return torch.flip(x, dims=[-1])

# 可逆 1x1 卷积（LU 分解实现）
class InvertibleConv1x1(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.channels = channels
        # LU 分解参数化
        # W = P @ L @ (U + diag(s))
        self.register_buffer('P', torch.eye(channels))
        self.L = nn.Parameter(torch.tril(torch.randn(channels, channels), -1))
        self.U_diag = nn.Parameter(torch.randn(channels))  # log 对角线
        self.U_off = nn.Parameter(torch.triu(torch.randn(channels, channels), 1))

    def _get_W(self):
        L = torch.eye(self.channels) + self.L
        U = torch.diag(torch.exp(self.U_diag)) + self.U_off
        return self.P @ L @ U

    def _get_log_det(self):
        return self.U_diag.sum()  # log|det W| = sum(log|U_ii|)

    def forward(self, z):
        W = self._get_W()
        # z: [..., channels]
        z_out = torch.einsum('...i,ij->...j', z, W)
        return z_out, self._get_log_det()

    def inverse(self, x):
        W = self._get_W()
        W_inv = torch.inverse(W)  # DxD 矩阵求逆
        return torch.einsum('...i,ij->...j', x, W_inv)

# ActNorm（可逆批量归一化）
class ActNorm(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
        self.log_scale = nn.Parameter(torch.zeros(dim))
        self.bias = nn.Parameter(torch.zeros(dim))
        self.initialized = False

    def forward(self, z):
        if not self.initialized and self.training:
            # 数据驱动初始化：使输出零均值单位方差
            with torch.no_grad():
                self.log_scale.data = -torch.log(z.std(dim=0) + 1e-6)
                self.bias.data = -z.mean(dim=0) * torch.exp(self.log_scale.data)
            self.initialized = True

        out = z * torch.exp(self.log_scale) + self.bias
        log_det = self.log_scale.sum() * z.shape[0]
        return out, log_det

    def inverse(self, x):
        return (x - self.bias) * torch.exp(-self.log_scale)

# MADE 掩码生成（自回归约束的核心）
def make_made_masks(input_dim, hidden_dims, output_multiplier=1):
    # 生成 MADE 的连接掩码
    # 给每个输入分配度数 (1 到 input_dim)
    input_degrees = torch.arange(1, input_dim + 1)

    masks = []
    prev_degrees = input_degrees

    for hidden_dim in hidden_dims:
        # 隐藏单元的度数 >= 前一层的最小度数
        hidden_degrees = torch.randint(
            low=prev_degrees.min(), high=input_dim,
            size=(hidden_dim,))
        masks.append((hidden_degrees[:, None] >= prev_degrees[None]).float())
        prev_degrees = hidden_degrees

    # 输出层的度数 = 对应输入的度数
    output_degrees = input_degrees.repeat(output_multiplier)
    masks.append((output_degrees[:, None] > prev_degrees[None]).float())

    return masks

# 示例：输入 4 维，隐藏层 [16, 16]
masks = make_made_masks(input_dim=4, hidden_dims=[16, 16])
print(f'Mask shapes: {[m.shape for m in masks]}')
print(f'Input mask (first):\n{masks[0]}')

# MAF 层实现
class MAFLayer(nn.Module):
    def __init__(self, dim, hidden_dim=64):
        super().__init__()
        self.dim = dim
        masks = make_made_masks(dim, [hidden_dim, hidden_dim], 2)

        # 用线性层 + 掩码实现 MADE
        self.linear1 = nn.Linear(dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, hidden_dim)
        self.linear_out = nn.Linear(hidden_dim, dim * 2)  # mean + log_std

        self.masks = masks
        self._apply_masks()

    def _apply_masks(self):
        self.linear1.weight.data *= self.masks[0]
        self.linear2.weight.data *= self.masks[1]
        self.linear_out.weight.data *= self.masks[2]

    def forward(self, z):
        # MAF: 计算 x 和 log|det J| (密度估计高效)
        h = torch.tanh(self.linear1(z))
        h = torch.tanh(self.linear2(h))
        params = self.linear_out(h)  # [mean, log_std] 拼接
        mean = params[:, :self.dim]
        log_std = params[:, self.dim:]
        log_std = torch.tanh(log_std)  # 稳定化

        x = z * torch.exp(log_std) + mean
        log_det = log_std.sum(dim=-1)
        return x, log_det

    def inverse(self, x):
        # MAF 逆变换需要串行计算
        z = torch.zeros_like(x)
        for i in range(self.dim):
            h = torch.tanh(self.linear1(z))
            h = torch.tanh(self.linear2(h))
            params = self.linear_out(h)
            mean_i = params[:, i]
            log_std_i = torch.tanh(params[:, self.dim + i])
            z[:, i] = (x[:, i] - mean_i) * torch.exp(-log_std_i)
        return z

# Flow 模型训练器
class FlowTrainer:
    def __init__(self, model, lr=1e-3, clip_grad=5.0):
        self.model = model
        self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)
        self.clip_grad = clip_grad
        self.base_dist = torch.distributions.Normal(0, 1)

    def train_step(self, x):
        # 单步训练：计算 NLL 并更新参数
        self.model.train()

        # 计算负对数似然
        z, log_det = self.model.inverse_and_log_det(x)
        log_p_z = self.base_dist.log_prob(z).sum(dim=-1)
        log_p_x = log_p_z + log_det
        loss = -log_p_x.mean()  # 最小化 NLL

        self.optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), self.clip_grad)
        self.optimizer.step()

        return loss.item()

    def sample(self, n_samples):
        # 从 Flow 模型采样
        self.model.eval()
        with torch.no_grad():
            z = self.base_dist.sample((n_samples, self.model.dim))
            x, _ = self.model.forward(z)
        return x

    def evaluate(self, x_test):
        # 评估测试集上的 NLL
        self.model.eval()
        with torch.no_grad():
            z, log_det = self.model.inverse_and_log_det(x_test)
            log_p_z = self.base_dist.log_prob(z).sum(dim=-1)
            log_p_x = log_p_z + log_det
            return -log_p_x.mean().item()

# 完整训练流程（2D Two-Moons 数据集）
import torch
import matplotlib.pyplot as plt
from sklearn.datasets import make_moons

def train_flow_2d(n_flows=8, epochs=500, batch_size=256):
    # 数据：双螺旋流形
    moons_data, _ = make_moons(n_samples=10000, noise=0.1)
    data = torch.FloatTensor(moons_data)
    # 标准化
    data = (data - data.mean(0)) / data.std(0)

    # 模型：RealNVP 风格
    model = RealNVP(dim=2, n_flows=n_flows, hidden_dim=128)
    trainer = FlowTrainer(model, lr=1e-3)

    # 训练
    losses = []
    for epoch in range(epochs):
        indices = torch.randperm(data.shape[0])
        epoch_loss = 0
        n_batches = 0

        for start in range(0, data.shape[0], batch_size):
            batch = data[indices[start:start+batch_size]]
            loss = trainer.train_step(batch)
            epoch_loss += loss
            n_batches += 1

        losses.append(epoch_loss / n_batches)

        if (epoch + 1) % 50 == 0:
            avg_nll = trainer.evaluate(data[:1000])
            print(f'Epoch {epoch+1:3d} | NLL: {avg_nll:.4f}')

            # 可视化生成样本
            samples = trainer.sample(5000)
            plt.figure(figsize=(4, 4))
            plt.scatter(samples[:, 0], samples[:, 1],
                       s=1, alpha=0.3)
            plt.title(f'Generated (epoch {epoch+1})')
            plt.xlim(-3, 3)
            plt.ylim(-3, 3)
            plt.savefig(f'flow_samples_{epoch+1}.png', dpi=150)
            plt.close()

    return model, losses

# 运行训练
model, losses = train_flow_2d(n_flows=8, epochs=200)