流匹配与 Rectified Flow

所属方向：生成范式 ｜ 难度：进阶 ｜ 前置知识：DDPM 的训练目标与采样流程（ddpm）、扩散过程的前向与反向推导（diffusion_math）
关键词：流匹配、flow matching、rectified flow、速度场、直线路径、reflow、NFE

---

01. 为什么需要它

先给三个数字，全部来自文末附录里能直接跑的脚本。

数字一：换个回归目标，预测误差幅度系数差 100 倍，对应平方损失权重差 10000 倍。 把「干净图估计」上的误差记为 $\delta$，在 $t=0.99$（几乎纯噪声）这一时刻：如果模型输出的是噪声 $\epsilon$，loss 看到的误差是 $0.0101\,\delta$；如果输出的是速度 $v$，loss 看到的是 $1.0101\,\delta$。在固定干净图误差的比较下，速度目标给予高噪声端更大的相对权重。 这是速度参数化与噪声参数化的一个权重差异；采用何种目标还取决于路径、预条件、训练时间分布与架构，不能把所有模型选择归因为这一个系数（03 节推导，图 1 画出整条曲线）。

数字二：「直线路径」的 ODE 轨迹一点都不直。 Rectified Flow 最常被转述成"路径是直的，所以快"。实测：把二维八模高斯混合推到数据端，用附录 reflow_lab.py 的 10 万个起点、64 步 RK2（128 NFE）轨迹账本，精确场的弧长/弦长均值为 1.6499（图 2 另用 300 步 RK4 示意）——轨迹比两点连线多走了 65% 的路。训练时那条 $(1-t)x+t\epsilon$ 插值线确实是直线，但模型实际采样走的是边缘速度场的积分曲线，两者不是一回事（02 节的图 2 把这件事画出来了）。

数字三：reflow 一轮，2 步采样的误差降 11.1 倍。 用第 1 轮训好的模型把噪声推到数据端、拿得到的配对重训一轮，配对插值上的归一化回归残差 $S_{pair}$ 从 $0.6229$ 掉到 $0.0027$（约 231 倍；它不是几何曲率），NFE=2 的终点偏差从 $0.04469$ 降到 $0.00402$，已经低于方差保持路径 2 步的 $0.0105$。但高 NFE 下的终点偏差没有因此继续下降：生成质量仍受教师分布、重训误差和有限样本评估影响，06 节展开。

所以这篇文章要回答三个问题：流匹配到底在训练什么、它和 DDPM 是不是两个东西、以及"直线路径"这个卖点真实兑现了多少。

---

02. 最小可用理解

三句话讲完：

1. 流匹配训练的是一个速度场，不是噪声。 采样是解一条 ODE：从噪声端 $t=1$ 出发，跟着 $v_\theta$ 走到数据端 $t=0$。训练只是回归：给定当前点 $z_t$，预测条件速度 $u_t$；在线性 Rectified Flow 路径上，它就是这一对端点的相对位移。最优解是条件期望 $E[u_t \mid z_t = z]$，而由连续性方程（03 节证），这个期望场恰好把 $p_1$ 运到 $p_0$——所以不需要知道任何密度，样本对就够。

2. 高斯路径提供了统一描述，但不同路径不等于同一个模型。 把路径统一写成 $z_t=\alpha_t x+\sigma_t\epsilon$：DDPM 那一路取 $\alpha_t=\sqrt{\bar\alpha_t}$、$\sigma_t=\sqrt{1-\bar\alpha_t}$；Rectified Flow 取 $\alpha_t=1-t$、$\sigma_t=t$。两者可以使用相同的网络架构。三种预测目标（干净图 $x$、噪声 $\epsilon$、速度 $v$）在同一条已知路径及非退化时间内可以代数换算输出（无需重训网络）；不同路径的模型不能仅换公式就变成彼此，04 节把换算残差验到 $10^{-15}$。

3. "直"的是训练插值线，不是采样轨迹。 在本文独立配对的平滑数据分布上，线性路径边缘场满足 $v(z,0)=-z$、$v(z,1)=z-\mu$，其中 $\mu=E[x]$。采样从 1 积分到 0，时间步为负，因此噪声端先朝数据均值走，数据端局部向外走；中间还有模式分流。实际积分轨迹可以弯曲，reflow 用模型生成的端点配对重训，是减小这种弯曲的一种办法。

这张图要看什么：在这 14 个共同起点上，左图（直线路径）出现明显回转，右图（VP 余弦路径）较平缓。灰点是 8 个数据模式，蓝圆是噪声起点，绿方是生成终点。两图使用相同的精确场计算方式和 300 步 RK4，只改变路径；其形状是本分布上的实验结果。VP 数据端速度为零，但由于数据均值不为零，噪声端速度并不为零，见 3.4 节。

---

03. 数学推导

3.1 高斯路径与条件速度

把常见高斯插值类扩散/流方法统一成一条路径。取数据样本 $x\sim p_{\text{data}}$、独立噪声 $\epsilon\sim N(0,I)$，令

$$z_t=\alpha_t x+\sigma_t\epsilon,\qquad t\in[0,1]$$

约定 $t=0$ 是数据端、$t=1$ 是噪声端，即 $\alpha_0=1,\sigma_0=0$、$\alpha_1=0,\sigma_1=1$。各符号的含义：$\alpha_t$ 是数据分量的幅度，$\sigma_t$ 是噪声分量的幅度，两者是标量函数，选不同的曲线就得到不同的方法：

方法	$\alpha_t$	$\sigma_t$	备注
VP / DDPM	$\sqrt{\bar\alpha_t}$	$\sqrt{1-\bar\alpha_t}$	满足 $\alpha_t^2+\sigma_t^2=1$，当数据协方差也为单位阵时保持单位方差；一般数据协方差仍随 t 改变
Rectified Flow	$1-t$	$t$	线性插值，中间分布方差会缩水（06 节）

对固定的一对 $(x,\epsilon)$，$z_t$ 是一条确定的曲线，它对时间的导数是

$$u_t=\dot\alpha_t x+\dot\sigma_t\epsilon$$

每个符号：$\dot\alpha_t$、$\dot\sigma_t$ 是两条幅度曲线的导数；$u_t$ 叫条件速度——它说的是"这一对端点对应的粒子此刻在往哪走"。直线路径下 $\dot\alpha_t=-1$、$\dot\sigma_t=1$，所以 $u_t=\epsilon-x$，与 $t$ 无关：整条插值线是匀速直线。这就是"直线"的全部含义，它只说了这条以 $(x,\epsilon)$ 为参数的曲线，没说采样时走的那条。

3.2 边缘速度场：为什么样本对就够训练

采样时我们手里没有那对 $(x,\epsilon)$，只有 $z_t$。所以要用的是边缘速度场：

$$v^*(z,t)=E\left[u_t \mid z_t=z\right]$$

读法：在所有"此刻恰好经过 $z$ 的粒子对"里，平均下来往哪走。要证它是"对的场"，也就是用它解 ODE 确实把 $p_1$ 运到 $p_0$。对任意光滑测试函数 $\phi$，沿条件路径求导再取期望：

$$\frac{d}{dt}E\left[\phi(z_t)\right]=E\left[\nabla\phi(z_t)\cdot u_t\right]=E\left[\nabla\phi(z_t)\cdot E[u_t\mid z_t]\right]=\int \nabla\phi(z)\cdot v^*(z,t)\,p_t(z)\,dz$$

第二步是把里层的条件期望提出来（塔性质）。最后那个积分正是连续性方程 $\partial_t p_t+\nabla\cdot(v^*p_t)=0$ 的弱形式，在速度场和密度满足适当正则性、相应 ODE 与连续性方程有唯一解的条件下，$v^*$ 生成的流在每个时刻匹配既定的 $p_t$，从 $p_1$ 走到 $p_0$。这里 $p_t$ 本身随时间变化，并不是保持某个固定分布不变。这一步就是"marginalization trick"：条件场逐对可得，边缘场取个条件期望就行。

于是在线性路径下，训练目标可以完全绕开密度（一般路径把目标换成 $u_t$）：

$$\min_\theta\;E_{t,x,\epsilon}\left\|v_\theta(z_t,t)-(\epsilon-x)\right\|^2$$

固定 $t$ 后，这个回归的逐点最优解恰是 $v^*(z,t)$。没有任何一项需要 $p_t$ 的表达式——这是流匹配在工程上能起飞的根本原因：它把"学分布"变成了"学一个回归"。

3.3 同一路径的三种预测输出如何换算

设 $m(z,t)=E[x\mid z_t=z]$（干净图估计）、$e(z,t)=E[\epsilon\mid z_t=z]$（噪声估计）。对 $z_t=\alpha_t x+\sigma_t\epsilon$ 两边取条件期望（条件期望是线性的，$z_t$ 在条件下是常数）：

$$z=\alpha_t m+\sigma_t e$$

这是贯穿全文的恒等式，$\alpha m+\sigma e=z$。再配合速度的定义 $v=\dot\alpha_t m+\dot\sigma_t e$，两个方程、两个未知数，解出

$$m=\frac{\sigma_t v-\dot\sigma_t z}{\dot\alpha_t\sigma_t-\dot\sigma_t\alpha_t},\qquad e=\frac{\dot\alpha_t z-\alpha_t v}{\dot\alpha_t\sigma_t-\dot\sigma_t\alpha_t}$$

分母 $\dot\alpha_t\sigma_t-\dot\sigma_t\alpha_t$ 是个行列式。直线路径下它恒等于 $-1$（$(-1)\cdot t-1\cdot(1-t)=-1$），于是换算干净得没有任何除法：

已知	求 $m$（干净图）	求 $e$（噪声）
速度 $v$	$m=z-t\,v$	$e=z+(1-t)\,v$
噪声 $\epsilon$	$m=(z-t\,e)/(1-t)$	—
干净图 $m$	—	$e=(z-(1-t)m)/t$

第一行值得盯一眼：从速度换算到干净图和噪声都不需要除以 $\alpha$ 或 $\sigma$，而第二、三行各有一个会趋奇的除法。这就是速度参数化在数值上更稳的那一半理由。

那为什么不能说"三种参数化完全等价"？ 真值层面等价（上表），loss 层面不等价。假设模型在 $m$ 上有误差 $\delta$，把它代入三种目标的误差：

$$\text{x0 目标：}\ \|\delta\|,\qquad \text{噪声目标：}\ \frac{\alpha_t}{\sigma_t}\|\delta\|,\qquad \text{速度目标：}\ \left|\dot\alpha_t-\frac{\dot\sigma_t\alpha_t}{\sigma_t}\right|\|\delta\|$$

来历：$e=(z-\alpha_t m)/\sigma_t$ 对 $m$ 求导得 $-\alpha_t/\sigma_t$；$v=\dot\alpha_t m+\dot\sigma_t e$ 对 $m$ 求导得 $\dot\alpha_t-\dot\sigma_t\alpha_t/\sigma_t$。直线路径代入，得到三条曲线

$$\text{x0：}1,\qquad \text{噪声：}\frac{1-t}{t},\qquad \text{速度：}\frac{1}{t}$$

这张图要看什么：横轴是时间 $t$，纵轴是对数刻度的放大倍数。$t=0.99$（接近纯噪声端）噪声参数化的曲线是 $0.0101$——在高噪声端，噪声目标对"你把干净图猜错了"几乎无感，因为 $z$ 里本来就几乎没有 $x$ 的信息；而速度目标的系数是 $1.0101$；真正取 $t\to1$ 的极限时，二者分别趋于 0 和 1。反过来在 $t\to 0$（数据端）二者对固定干净图误差的换算系数都增大，速度系数比噪声系数加法上大 1；这不等于直接 velocity 训练的目标或梯度必然发散。这里能直接推出的是高噪声端对干净图估计误差的相对损失权重不同，不能单凭它保证真实网络的梯度大小或学习效率。

3.4 直的是插值线，不是轨迹

先看端点行为。本文数据是平滑的高斯混合，数据与噪声独立；因此在数据端，$m(z,0)=z$、$e(z,0)=E[\epsilon\mid x=z]=0$；在噪声端，$m(z,1)=\mu=E[x]$、$e(z,1)=z$。这些关系来自端点的条件独立性，不能用 $(z-\alpha m)/\sigma$ 中的“分子趋零”来推断一个 $0/0$ 极限。

make_data 的 8 个分量权重与 $1+0.35\cos(\cdot)$ 成正比，并不均匀。精确加权均值为

$$\mu=(0.294464,\,-0.248024).$$

线性路径 $\dot\alpha=-1,\dot\sigma=1$ 的端点速度是

$$v(z,0)=-z,\qquad v(z,1)=z-\mu.$$

采样沿负时间方向积分，所以噪声端局部朝 $\mu$ 走，数据端局部沿 $z$ 向外走。中间的模式分流也会影响轨迹，端点公式本身不能决定全程弯曲程度；01 节的弧长比是实际积分测得的结果，不是由端点方向推出的普遍定理。

VP 余弦路径取 $\alpha=\cos(\pi t/2)$、$\sigma=\sin(\pi t/2)$，因此

$$v_{\rm VP}(z,0)=0,\qquad v_{\rm VP}(z,1)=-\frac{\pi}{2}\mu=(-0.462543,\,0.389595).$$

只有数据均值为零时，这条 VP 路径的两个端点速度才都为零。本实验不满足该条件。

需要区分两个诊断量。本文脚本 straightness_of 实际计算的是原始端点配对插值上的归一化回归残差：

$$S_{pair}=\frac{\int_0^1 E\|v((1-t)z_0+t z_1,t)-(z_1-z_0)\|^2dt}{E\|z_1-z_0\|^2}.$$

当 $v=v^*$ 时，分子是条件回归的不可约误差。下式的期望同时包含均匀时间与端点配对，并且 $S_{pair}$ 用 $v^*$ 计算：

$$E\|v_\theta-u\|^2=E\|v_\theta-v^*\|^2+S_{pair}E\|z_1-z_0\|^2.$$

Rectified Flow 论文沿实际 ODE 轨迹 $Z_t$ 定义的 straightness 则比较 $v(Z_t,t)$ 与 $Z_1-Z_0$。这不是在独立原始配对插值线上评估，不能直接把本文 $S_{pair}$ 当成论文的同名量。本文另用实际积分轨迹的弧长/弦长检测几何变直，两个指标应分别报告。弧长的离散计算必须是 $\sum_j\|Z_{t_{j+1}}-Z_{t_j}\|_2$，不能先分别累加坐标上的绝对位移再取范数。

path_lab.py 另取 1500 个起点、400 步 RK4（1600 NFE）：线性路径的平均弧长/弦长为 1.6232，VP 为 1.2915。沿实际轨迹的速度偏差再除以平均弦长平方，得到归一化 $S_{traj}$ 分别为 1.2293、0.5317。这与原始配对插值上的 $S_{pair}$ 不同。01 节的 1.6499 来自 reflow_lab.py 的另一批 10 万个起点和 64 步 RK2；样本与积分精度不同，诊断数值也会不同。

04. 代码实现

全部代码在文末附录，五个脚本，数值实验依赖 numpy、画图另需 matplotlib：fm_oracle.py（实验台与精确场）、param_lab.py（三种参数化）、path_lab.py（路径对比与 NFE）、reflow_lab.py（reflow 两轮）、make_figures.py（配图）。实验台是一个二维八模高斯混合，噪声是 $N(0,I)$。选它的原因是：高斯混合经过任何高斯路径之后仍是高斯混合，于是 $m$、$e$、$v^*$ 全都有闭式解，不用训练也不用采样近似。

4.1 精确场：四条闭式

分量 $k$ 的中间方差是 $V_k(t)=\alpha_t^2 s_k^2+\sigma_t^2$（$s_k$ 是该分量的标准差），后验责任 $r_k$ 由贝叶斯公式给出。三个量都是"责任加权的条件均值"：

def m(self, z, t):                      # E[x | z_t = z]
    a = self.path.alpha(t)
    r, V, diff = self._post(z, t)
    ex = self.data.mu[None,:,:] + (a * self.data.s**2 / V)[None,:,None] * diff
    return (r[:,:,None] * ex).sum(1)

def e(self, z, t):                      # E[eps | z_t = z]
    sg = self.path.sigma(t)
    r, V, diff = self._post(z, t)
    ee = (sg / V)[None,:,None] * diff
    return (r[:,:,None] * ee).sum(1)

def v(self, z, t):                      # 边缘速度场
    return self.path.dalpha(t) * self.m(z, t) + self.path.dsigma(t) * self.e(z, t)

先验两个恒等式。第一个是 3.3 节的 $\alpha m+\sigma e=z$，第二个是把 $v^*$ 用另一种方式算一遍——直线路径下 3.3 节的表给出 $v^*=(z-m)/t$：

=== 恒等式自检：alpha*m + sigma*e 应等于 z（残差 ~1e-15）===
path        t     |a*m+s*e-z|    |v-(da*m+ds*e)|
linear      0.05        2.66e-15          0.00e+00
linear      0.95        1.78e-15          0.00e+00
cosine_vp   0.05        1.78e-15          0.00e+00
cosine_vp   0.95        1.78e-15          0.00e+00

=== 直线路径下 v* 的两种算法是否一致：(z-m)/t 与 alpha' m + sigma' e ===
  t=0.1  最大绝对差 = 1.332e-14
  t=0.9  最大绝对差 = 1.776e-15

符号的物理含义都压到了 $10^{-15}$ 量级，说明推导和实现是同一件事。

4.2 换算表与放大倍数

把 3.3 节的换算写成代码并逐点验真值（注意残差随 $t\to 1$ 变大——那是除以 $(1-t)$ 的数值放大，不是推导错）：

def v_to_me(z, v, path, t):
    a, sg = path.alpha(t), path.sigma(t)
    da, ds = path.dalpha(t), path.dsigma(t)
    det = da * sg - ds * a               # 直线路径恒等于 -1
    m = (sg * v - ds * z) / det
    e = (da * z - a * v) / det
    return m, e

path        t     v->m        v->e        eps->m      eps->v
linear      0.10    1.78e-15    1.33e-15    1.78e-15    1.78e-15
linear      0.90    1.22e-15    8.88e-16    1.30e-14    1.29e-14
linear      0.99    1.78e-15    1.78e-15    2.04e-13    2.03e-13

放大倍数不靠公式背，用数值扰动直接量：给 $m$ 加一个长度固定为 $0.01$ 的随机扰动，看三种预测误差范数各被放大多少；平方损失对应这些系数的平方。实测与解析式在小数点后四位完全一致（$t=0.1$ 时噪声目标 $9.0000$ 对 $9.0000$、速度目标 $10.0000$ 对 $10.0000$），整条曲线就是图 1。

4.3 路径对比：NFE-误差曲线

误差尺子值得单独说一句。常用做法是"拿一条很高步数的参考解当真值"，但那条参考解自己也有截断误差。这里利用 $p_t$ 仍是高斯混合这一点，取一组 RBF 特征 $\phi_c(z)=\exp(-\|z-c\|^2/(2h^2))$，它在高斯下的期望有闭式，于是参照侧没有任何采样噪声。生成样本一侧仍有有限样本误差：一批 4000 个真实样本对精确 $p_{\text{data}}$ 的偏差是 $0.00211$，这里只把它画成参考线，既不是误差下限，也不是显著性阈值。更换随机种子后，这个参考偏差也会变化。要判断小差异是否稳定，应对方法使用共同起点并重复多个种子。有限组 RBF 特征只是分布诊断，特征偏差小不能证明两个分布相同。

NFE 统计实际速度场评估次数：Euler 每步 1 次，RK2 每步 2 次。所以 NFE=8 时分别运行 8 步 Euler 或 4 步 RK2，表中按这一预算公平比较。

=== NFE-误差：直线路径 vs VP 余弦路径（同一批起点、同一批随机数）===
  NFE   linear/euler  linear/rk2   vp/euler     vp/rk2
    2       0.04074      0.04113     0.01047    0.04603
    4       0.01491      0.00736     0.00562    0.00318
    8       0.00802      0.00202     0.00329    0.00236
   16       0.00466      0.00183     0.00234    0.00194
   32       0.00307      0.00184     0.00199    0.00186
  256       0.00194      0.00183     0.00184    0.00183

这张图要看什么：高 NFE 时，四条线都接近同一量级的小偏差，与有限样本误差并存；这不能单独证明场正确。低 NFE 的差异依赖路径和积分器：在 Euler、NFE=8 这一点，线性路径误差是 VP 的约 2.4 倍（$0.00802$ 对 $0.00329$），不是“需要 2.4 倍步数”。要比较所需步数，应先规定相同误差阈值。RK2 在 NFE=2 时只有一个积分步，误差甚至大于 Euler；较高阶也不保证每个极低预算下都更好。

数值积分误差受速度场的空间与时间变化、时间网格和积分器共同影响，不能仅用“曲线弯”或错误的“两端速度为零”解释。时间 shift 也可能改变弯曲轨迹的积分精度；这里只能报告本次测试：Euler、NFE=8 时，shift=1、3、6 的误差分别为 $0.00802$、$0.00864$、$0.01843$，这两个非均匀网格没有带来改善。

4.4 Reflow：把弯的轨迹拉直

reflow 的实现出奇地短：用第 1 轮模型从噪声端积分到数据端，得到新配对，再训一轮。

def round1_pairs(n, seed=0):
    rng = base_rng(seed)
    data = make_data()
    z0 = data.sample(n, rng)              # 数据端（t=0）
    z1 = rng.standard_normal((n, D))      # 噪声端（t=1），独立配对
    return z0, z1

def generate_pairs(model, z1, n_steps=100):
    sched = make_schedule(n_steps)        # RK2：100 步 = 200 次场评估
    z0, _ = integrate(lambda z, t: model(z, t), z1, sched, "rk2")
    return z0                             # 新配对就是 (z0, z1)

模型是 numpy 手搓的两层 128 宽 tanh MLP。第 1 轮在 10 万条独立配对上训练（每步从数据样本池抽样并重采独立噪声），训练 loss 收敛到 $2.1853$——按样本平方范数是 $4.3706$，而估计的不可约回归误差 $S_{pair}\cdot E\|z_1-z_0\|^2\approx4.3702$（未舍入计算；两个因子分别约为 $0.6229$、$7.0162$）。loss 接近这一估计，但有限采样下略低或略高都可能发生，不能证明模型已达精确最优；在分布内查询点上它与精确场的相对均方误差是 $1.28\%$。

然后是关键的一步——量两轮的归一化配对残差和 NFE：

=== 归一化配对回归残差（另用弧长/弦长量轨迹）===
  第 1 轮配对 + 精确场   S_pair = 0.6229
  第 1 轮配对 + 第1轮模型 S_pair = 0.6289
  第 2 轮配对 + 第2轮模型 S_pair = 0.0027
  弧长/弦长：精确场 1.6499 | 第1轮 1.6463 | 第2轮 1.0027

=== NFE-误差（直线路径；精确场 / 第1轮模型 / 第2轮模型）===
  NFE    精确场    第1轮    第2轮
    2   0.04012   0.04469   0.00402
    4   0.01469   0.01789   0.00401
    8   0.00785   0.00927   0.00402
   16   0.00427   0.00563   0.00402
   64   0.00182   0.00389   0.00403

这张图要看什么：左图是三条 NFE-误差曲线，绿色（reflow 后）从 NFE=2 起就是一条平线——步数不再是瓶颈；右图分别比较配对回归残差与实际轨迹弧长/弦长：配对残差第 1 轮用精确场、第 2 轮用重训模型；弧长柱均用对应训练模型。第 2 轮的弧长/弦长 $1.0027$，轨迹已经基本就是两点连线。另外注意第 2 轮配对的弦长平方从 $7.0162$ 掉到 $1.2410$：reflow 之后噪声端和数据端被强相关地配对了，这种配对改变也体现在实际轨迹变直上；不能把较小的配对残差本身称为曲率。

---

05. 工业级实现对照

以 diffusers 的 FlowMatchEulerDiscreteScheduler 为例（以 2026-09 的实现为准，scheduling_flow_match_euler_discrete.py），最小实现和生产实现的差别集中在四处。

一处一行更新。 step() 的默认分支就是速度空间里的显式 Euler：

dt = sigma_next - sigma
prev_sample = sample + dt * model_output

model_output 被直接当作速度用，不做任何转换。和 04 节最小实现的 z = z - h * v 是同一行——diffusers 用 $\sigma$（就是本文的 $t$）当时间轴，$\sigma$ 从 1 递减到 0，所以步长是负的，方向自动正确。

换算真的写在生产代码里。 随机采样分支里有这一行：

x0 = sample - current_sigma * model_output

这正是 3.3 节换算表的第一格 $m=z-t\,v$。也就是说"速度预测的模型可以零成本拿到干净图估计"这件事，在 diffusers 里是一行乘法。

模型评估点与最终样本落点不同。 默认噪声表的最后非零模型评估点约为 0.001，set_timesteps 仍在末尾追加 0。最后执行 sample + (0 - sigma) * model_output，样本实际到达 0，只是不再在 0 调用网络；这也等于直线路径的 $x_0$ 估计。二维 toy 的 linspace(1,0,...) 在循环中同样先评估当前非零时刻，再积分到目标，不是高维模型禁止的做法。

时间 shift 调整有效时间分配，生产模型常根据分辨率设置它。 静态 shift 的公式是

sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas)

而动态 shift 按序列长度插值：base_shift=0.5 对应 256 个 token，max_shift=1.15 对应 4096 个 token，中间用指数插值 $\exp(\mu)/(\exp(\mu)+(1/t-1))$ 过渡。动机是：更高分辨率往往增加空间相关信号的冗余，使同等加噪后信息仍较易恢复，实践中据此调整训练/采样的有效噪声分配；具体 shift 应遵循模型配置。04 节只验证了这个 toy 上的几个 shift 设置：它们没有降低对应预算下的误差，不能推广成 shift 对弯曲轨迹无效。

另外 SD3 论文在训练侧还做了一件事：时间 $t$ 不用均匀分布采，用 logit-normal（众数在中间）。对照图 1 就能理解为什么——不同时间采样相当于重新分配训练权重；SD3 直接回归 velocity，不能拿噪声参数化在高噪声端的弱信号直接解释其 logit-normal 选择。

---

06. 代价与边界

直线路径不是免费的午餐。 在这组实验的 Euler、NFE=8 处，线性路径的误差约为 VP 余弦路径的 2.4 倍；这个误差比不等于达到同一质量所需的步数比。需要诚实标注边界：这是一个二维、数据尺度（标准差 1.56）与噪声尺度（1）不匹配的 toy，不能据此断言"VP 路径普遍更省步数"；但至少说明直线插值不保证少步积分准确。实际轨迹接近直线有助于理解 reflow 的效果，积分难度仍取决于沿轨迹的速度变化与数值格式。真实模型里路径选择还和训练分布、模型容量、蒸馏方案纠缠在一起，单独归因很难。

reflow 仍受教师分布与学习误差限制。 本次第 2 轮的高 NFE 偏差约为 $0.00403$；第 1 轮模型用于训练的 10 万个生成端点偏差为 $0.00328$。二者评估样本量不同，不能据此定量断言“误差几乎全部继承自教师”。在理论条件满足、回归和积分精确的理想 reflow 中，端点边缘分布保持为教师的分布；有限模型还会增加重训和积分误差。图 4 显示的是少步曲线被压平，不能当作重训自动改善真实数据拟合的证据。

reflow 的账单。 它需要用第 1 轮模型做一次全量生成，生成量要够训第 2 轮。其成本由配对生成量、教师步数和重训步数共同决定，不能固定说成翻倍，所以工程上要么只在最后的精调阶段做，要么直接换成少步蒸馏（知识树上的 step_distillation，待写）。

$\sigma\to0$ 的数值边界。 需要避免在端点进行会除以零的参数化转换；速度场 Euler 更新自身不含该除法。最后从 $\sigma>0$ 到 0 的 Euler 步正是 $x_0=x-\sigma v$，两种说法在直线路径下是同一操作。

---

07. 经典论文脉络

• Lipman et al., 2022（arXiv:2210.02747）Flow Matching for Generative Modeling：用 simulation-free 条件回归训练速度场，避免训练期间反复数值求解 ODE；不是从 O(1) 复杂度降到 1，提出条件路径 + marginalization trick，是"用回归训速度场"的源头。
• Liu et al., 2022（arXiv:2209.03003）Rectified Flow：从线性插值回归出发构造流，并提出 reflow。在论文的理想化条件下，reflow 具有直线度和凸运输成本的相关保证；本文 4.4 节的有限样本、有限模型实验用于观察趋势，不能替代理论条件。
• Albergo & Vanden-Eijnden, 2022，Building Normalizing Flows with Stochastic Interpolants：通过随机端点的插值和回归目标构造确定性概率流。
• Albergo, Boffi & Vanden-Eijnden, 2023，Stochastic Interpolants: A Unifying Framework for Flows and Diffusions：进一步引入可调噪声项，建立流与扩散过程的统一框架；应与上一项区分引用。
• Karras et al., 2022（arXiv:2206.00364）EDM：不谈"流"，把预条件（参数化）和时间调度当成独立的自由度来调。它把输出预条件、损失权重与采样调度分开设计；本文 3.3 节涉及其中的输出换算与相对权重，不能把换算恒等式当作训练行为完全等价。
• Esser et al., 2024（arXiv:2403.03206）Stable Diffusion 3：把 rectified flow 推到大规模文生图，给出 logit-normal 时间采样与分辨率相关 shift——工业界从 DDPM 切换到流匹配的标志点。

---

08. 常见误解

误解一："流匹配不是扩散模型，是另一套东西。" 同一族高斯路径，DDPM 是 $\alpha=\sqrt{\bar\alpha},\sigma=\sqrt{1-\bar\alpha}$ 的一支，Rectified Flow 是 $\alpha=1-t,\sigma=t$ 的一支。可以使用相同的网络架构；同一路径上的三种预测输出可以代数互转（端点需处理退化）（04 节残差 $10^{-15}$）。真正不同的是中间分布 $p_t$ 的形状和 loss 的加权，"两个流派"的说法遮住了这些可调的自由度。

误解二："路径是直的，所以一两步就能出图。" 直的是训练时的插值线；采样走的是边缘场的积分曲线。实测 1-rectified 的弧长/弦长是 $1.65$，NFE=2 误差 $0.04469$。要两步出图，靠的是 reflow 之后的第 2 轮（$0.00402$），不是第 1 轮的"直线路径"。

误解三："reflow 之后质量和速度都变好了。" 只对了一半。步数-精度曲线确实被压平（NFE=2 就到位），但精度仍受教师样本分布与重训误差限制：第 2 轮在本次实验中停在约 $0.00403$，教师生成端点的参考偏差为 $0.00328$。理想 reflow 保持教师端点边缘，有限训练下则需另外评估分布偏差，不能保证质量自动提高。

误解四："velocity 预测只是把回归目标换了个写法。" 真值层面是，loss 层面不是。$t=0.99$ 处噪声目标对干净图误差的放大倍数是 $0.0101$，速度目标是 $1.0101$，幅度系数差 100 倍、平方权重差 10000 倍；这只描述固定干净图误差的相对权重，不代表高噪声端一半训练样本都无用。

误解五："流匹配需要知道边缘密度或 score。" 3.2 节的推导全程只用了条件期望和塔性质，训练目标里没有任何一项含 $p_t$。需要密度的是评估（比如算 NLL），不是训练。

---

09. 动手验证

三个可以自己跑的小实验，按顺序：

1. 恒等式链：python fm_oracle.py。应该看到 $\alpha m+\sigma e=z$ 的残差在 $10^{-15}$ 量级，$v^*$ 的两种算法差 $10^{-14}$ 以内。如果你的实现里这两个数在 $10^{-6}$ 量级，大概率是后验责任没做 log-sum-exp。
2. 放大倍数：python param_lab.py。看解析式与数值扰动两列是否完全一致，再看 $t=0.99$ 那一行——噪声目标 $0.0101$ 对速度目标 $1.0101$。
3. reflow 的残差与采样曲线：python reflow_lab.py（约两分钟，纯 numpy）。确认三件事：第 1 轮逐坐标平均 loss 乘以维数 2 后，接近估计的 $S_{pair}\cdot E\|z_1-z_0\|^2$；第 2 轮弧长/弦长 $\approx 1.00$；第 2 轮 NFE 曲线近乎平坦，但仍存在相对真实数据的偏差。

想改实验设置的话：make_data 里的 radius 控制数据尺度（默认 2.2）。改变它后重新运行 path_lab.py，比较不同路径在相同实际 NFE 下的误差。半径改变会同时影响数据协方差、模式间距和路径难度；不要把某一次误差比当作固定的步数收益。

---

10. 延伸阅读

• 前置：ddpm（VP 路径与噪声预测）、diffusion_math（前向反向推导与 score）、vae_elbo（另一条"回归式"生成目标）。
• 相邻：ddim_samplers（同一套 ODE 视角下的高阶格式与步长调度）、cfg（流匹配模型上的引导，公式形式几乎不变）。
• 后继：step_distillation（少步蒸馏，reflow 的工程替代品，待写）；video_rl（已发布，Flow-GRPO 在流匹配模型上做在线强化学习，直接建立在这些速度场上）。

附录：完整代码

09 节用到的脚本全文如下（reflow_lab.py、path_lab.py、fm_oracle.py、param_lab.py、make_figures.py）。复制到本地存成同名文件，按各脚本开头的依赖说明准备环境后即可运行。

reflow_lab.py

# -*- coding: utf-8 -*-
"""flow_matching 实验台（四）：Reflow —— 用模型自己造的配对再训一轮。

Rectified Flow 最容易被转述错的一句话是「它的路径是直的」。
真的是直的，但直的是**训练时那条插值线** (1-t)x + t*eps；
ODE 真正走出来的轨迹是边缘速度场 v*(z,t) = E[eps - x | z_t = z] 的积分曲线，
本例的积分轨迹会弯曲，弧长须把每一小段的欧式长度相加后再与弦长比较。

Reflow 做的事是：用第 1 轮训好的模型把噪声推到数据端，得到一批**配对**
(z_1, z_0) = (起点噪声, 模型生成的终点)，再拿这批配对重训一轮。
理想化 reflow 的性质有理论条件；本实验另外测量有限模型的轨迹长度比和采样误差。

这一份用 numpy 手搓一个小 MLP 当"模型"，分别报告配对插值上的归一化回归残差 S_pair、实际 ODE 轨迹弧长/弦长和 NFE-误差。
"""

import numpy as np

from fm_oracle import (
    D, base_rng, make_data, LinearPath, OracleFlow,
    make_rbf_centers, discrepancy, make_schedule, integrate,
)

C, BW = make_rbf_centers()
PATH = LinearPath()


# ─────────────────── 一个小 MLP（numpy 手写） ───────────────────
class MLP:
    def __init__(self, sizes, rng):
        self.W, self.b = [], []
        for i in range(len(sizes) - 1):
            lim = np.sqrt(6.0 / (sizes[i] + sizes[i + 1]))
            self.W.append(rng.uniform(-lim, lim, (sizes[i], sizes[i + 1])))
            self.b.append(np.zeros(sizes[i + 1]))
        self.m = [np.zeros_like(w) for w in self.W]
        self.v = [np.zeros_like(w) for w in self.W]
        self.mb = [np.zeros_like(x) for x in self.b]
        self.vb = [np.zeros_like(x) for x in self.b]

    def forward(self, X):
        A = [X]
        H = X
        for i in range(len(self.W) - 1):
            H = np.tanh(H @ self.W[i] + self.b[i])
            A.append(H)
        A.append(H @ self.W[-1] + self.b[-1])
        return A[-1], A

    def step(self, X, Y, lr, tstep, beta1=0.9, beta2=0.999, eps=1e-8):
        pred, A = self.forward(X)
        g = 2.0 * (pred - Y) / pred.size
        for i in range(len(self.W) - 1, -1, -1):
            gW = A[i].T @ g
            gb = g.sum(0)
            self.m[i] = beta1 * self.m[i] + (1 - beta1) * gW
            self.v[i] = beta2 * self.v[i] + (1 - beta2) * (gW * gW)
            self.mb[i] = beta1 * self.mb[i] + (1 - beta1) * gb
            self.vb[i] = beta2 * self.vb[i] + (1 - beta2) * (gb * gb)
            mh = self.m[i] / (1 - beta1 ** tstep)
            vh = self.v[i] / (1 - beta2 ** tstep)
            # 先用本次前向的旧权重传播梯度，再更新参数。
            g_prev = (g @ self.W[i].T) * (1 - A[i] ** 2) if i > 0 else None
            self.W[i] -= lr * mh / (np.sqrt(vh) + eps)
            bmh = self.mb[i] / (1 - beta1 ** tstep)
            bvh = self.vb[i] / (1 - beta2 ** tstep)
            self.b[i] -= lr * bmh / (np.sqrt(bvh) + eps)
            if i > 0:
                g = g_prev
        return float(((pred - Y) ** 2).mean())

    def __call__(self, z, t):
        tt = np.full((z.shape[0], 1), float(t))
        return self.forward(np.concatenate([z, tt], axis=1))[0]


def train(pairs_z0, pairs_z1, n_step=4000, batch=512, lr=3e-3, seed=0, fresh=False):
    """在给定配对上训练速度场：目标 u = z_1 - z_0，z_t = (1-t) z_0 + t z_1。

    fresh=True 时每一步现采一批新配对（第 1 轮的配对可以无限造），
    避免模型把 2 万条配对背下来——背下来会让 loss 掉到"条件方差"以下，
    看起来很美，其实场是有偏的。
    """
    rng = base_rng(seed)
    net = MLP([D + 1, 128, 128, D], np.random.default_rng(seed + 1))
    n = pairs_z0.shape[0]
    losses = []
    for s in range(1, n_step + 1):
        if fresh:
            idx = rng.integers(0, n, batch)
            z1 = rng.standard_normal((batch, D))
            z0 = pairs_z0[idx]          # 数据端样本池足够大，随机抽即视为新样本
        else:
            idx = rng.integers(0, n, batch)
            z0 = pairs_z0[idx]
            z1 = pairs_z1[idx]
        t = rng.uniform(0.0, 1.0, (batch, 1))
        zt = (1 - t) * z0 + t * z1
        u = z1 - z0
        X = np.concatenate([zt, t], axis=1)
        cur_lr = lr * (0.3 ** (s / n_step))     # 余弦退火换成简单指数退火
        l = net.step(X, u, cur_lr, s)
        losses.append(l)
    return net, float(np.mean(losses[-200:]))


# ─────────────────── 两轮的配对 ───────────────────
def round1_pairs(n, seed=0):
    """第 1 轮：数据与噪声独立配对（这就是 Rectified Flow 原始的训练配对）。"""
    rng = base_rng(seed)
    data = make_data()
    z0 = data.sample(n, rng)              # 数据端（t=0）
    z1 = rng.standard_normal((n, D))      # 噪声端（t=1）
    return z0, z1


def generate_pairs(model, z1, n_steps=100):
    """用 RK2 的 n_steps 个积分步生成配对，实际场评估次数为 2*n_steps。"""
    sched = make_schedule(n_steps)
    z0, _ = integrate(lambda z, t: model(z, t), z1, sched, "rk2")
    return z0


# ─────────────────── 评价指标 ───────────────────
def straightness_of(pairs_z0, pairs_z1, vfun, nstep=32):
    """配对插值上的归一化回归残差；不是论文沿实际 ODE 轨迹定义的 straightness。"""
    ts = np.linspace(1.0, 0.0, nstep + 1)[:-1]
    dev = 0.0
    for t in ts:
        zt = (1 - t) * pairs_z0 + t * pairs_z1
        dev += float(np.mean(np.sum((vfun(zt, t) - (pairs_z1 - pairs_z0)) ** 2, axis=1)))
    dev /= len(ts)
    return dev / float(np.mean(np.sum((pairs_z1 - pairs_z0) ** 2, axis=1)))


def arc_over_chord(pairs_z1, vfun, nstep=64):
    """真走一遍：从 z_1 出发积分到 t=0，量轨迹弧长与端点弦长之比。"""
    sched = make_schedule(nstep)
    z0, traj = integrate(vfun, pairs_z1, sched, "rk2")
    seg = np.linalg.norm(np.diff(traj, axis=0), axis=2).sum(0)
    chord = np.linalg.norm(z0 - pairs_z1, axis=1)
    return float(np.mean(seg / np.maximum(chord, 1e-12))), z0


def nfe_curve(pairs_z1, vfun, nfe_list, data):
    flow = OracleFlow(data, PATH)
    out = []
    for nfe in nfe_list:
        sched = make_schedule(nfe)
        z0, _ = integrate(vfun, pairs_z1, sched, "euler")
        out.append(discrepancy(z0, 0.0, flow, C, BW))
    return np.array(out)


def field_error(model, flow, n=4000, seed=3):
    """训练出来的场与精确场差多少（相对均方）。

    查询点取自真实边缘 p_t（与第 1 轮训练分布一致），并在五个时刻汇总平方误差。
    """
    rng = base_rng(seed)
    data = make_data()
    num = 0.0
    den = 0.0
    for t in (0.1, 0.3, 0.5, 0.7, 0.9):
        x = data.sample(n, rng)
        eps = rng.standard_normal((n, D))
        zt = (1 - t) * x + t * eps
        v = flow.v(zt, t)
        num += float(np.mean(np.sum((model(zt, t) - v) ** 2, axis=1)))
        den += float(np.mean(np.sum(v ** 2, axis=1)))
    return num / den


def chord_scale(pairs_z0, pairs_z1):
    """E||z_1 - z_0||^2，配对残差 S_pair 的分母。"""
    return float(np.mean(np.sum((pairs_z1 - pairs_z0) ** 2, axis=1)))


def state():
    data = make_data()
    flow = OracleFlow(data, PATH)
    z0_1, z1 = round1_pairs(100000, seed=0)
    net1, loss1 = train(z0_1, z1, n_step=6000, seed=0, fresh=True)
    z0_2 = generate_pairs(net1, z1, n_steps=100)
    net2, loss2 = train(z0_2, z1, n_step=6000, seed=0)

    nfe_list = [2, 4, 8, 16, 32, 64]
    eval_z1 = base_rng(999).standard_normal((4000, D))
    out = {
        "nfe": np.array(nfe_list),
        "loss1": loss1,
        "loss2": loss2,
        "ferr1": field_error(net1, flow),
        "ferr2": field_error(net2, flow),
        "S1_oracle": straightness_of(z0_1, z1, lambda z, t: flow.v(z, t)),
        "S1_learned": straightness_of(z0_1, z1, lambda z, t: net1(z, t)),
        "S2_learned": straightness_of(z0_2, z1, lambda z, t: net2(z, t)),
        "arc1_oracle": arc_over_chord(z1, lambda z, t: flow.v(z, t))[0],
        "arc1_learned": arc_over_chord(z1, lambda z, t: net1(z, t))[0],
        "arc2_learned": arc_over_chord(z1, lambda z, t: net2(z, t))[0],
        "nfe_oracle": nfe_curve(eval_z1, lambda z, t: flow.v(z, t), nfe_list, data),
        "nfe_learned1": nfe_curve(eval_z1, lambda z, t: net1(z, t), nfe_list, data),
        "nfe_learned2": nfe_curve(eval_z1, lambda z, t: net2(z, t), nfe_list, data),
    }
    return out


if __name__ == "__main__":
    data = make_data()
    flow = OracleFlow(data, PATH)

    print("=== 第 1 轮：独立配对上训练 ===")
    z0_1, z1 = round1_pairs(100000, seed=0)
    net1, loss1 = train(z0_1, z1, n_step=6000, seed=0, fresh=True)
    s1 = chord_scale(z0_1, z1)
    print(f"  训练 loss（末 200 步均值）= {loss1:.4f}（按样本平方范数是 {2*loss1:.4f}）")
    s_pair_oracle = straightness_of(z0_1, z1, lambda z, t: flow.v(z, t))
    print(f"  估计的不可约回归误差 = S_pair * E||z1-z0||^2 = "
          f"{s_pair_oracle:.4f} * {s1:.4f} = {s_pair_oracle*s1:.4f}")
    print(f"  学出来的场 vs 精确场，相对均方误差（on-distribution）= {field_error(net1, flow):.4f}")

    print()
    print("=== 用第 1 轮模型生成新配对 ===")
    z0_2 = generate_pairs(net1, z1, n_steps=100)
    print(f"  生成终点与精确 p_data 的特征偏差 = {discrepancy(z0_2, 0.0, flow, C, BW):.5f}"
          f"（另有一组 4000 个真实样本的参考偏差约 0.00211，非误差下限）")
    net2, loss2 = train(z0_2, z1, n_step=6000, seed=0)
    s2 = chord_scale(z0_2, z1)
    print(f"  第 2 轮训练 loss = {loss2:.4f}（按样本平方范数是 {2*loss2:.4f}）")
    print(f"  第 2 轮弦长平方 E||z1-z0||^2 = {s2:.4f}（第 1 轮是 {s1:.4f}，配对更紧了）")

    print()
    print("=== 归一化配对残差（另用弧长/弦长量轨迹）===")
    print(f"  第 1 轮配对 + 精确场   S_pair = {straightness_of(z0_1, z1, lambda z, t: flow.v(z, t)):.4f}")
    print(f"  第 1 轮配对 + 第1轮模型 S_pair = {straightness_of(z0_1, z1, lambda z, t: net1(z, t)):.4f}")
    print(f"  第 2 轮配对 + 第2轮模型 S_pair = {straightness_of(z0_2, z1, lambda z, t: net2(z, t)):.4f}")
    print(f"  弧长/弦长：精确场 {arc_over_chord(z1, lambda z, t: flow.v(z, t))[0]:.4f} | "
          f"第1轮 {arc_over_chord(z1, lambda z, t: net1(z, t))[0]:.4f} | "
          f"第2轮 {arc_over_chord(z1, lambda z, t: net2(z, t))[0]:.4f}")

    print()
    nfe_list = [2, 4, 8, 16, 32, 64]
    eval_z1 = base_rng(999).standard_normal((4000, D))
    o = nfe_curve(eval_z1, lambda z, t: flow.v(z, t), nfe_list, data)
    a = nfe_curve(eval_z1, lambda z, t: net1(z, t), nfe_list, data)
    b = nfe_curve(eval_z1, lambda z, t: net2(z, t), nfe_list, data)
    print("=== NFE-误差（直线路径；精确场 / 第1轮模型 / 第2轮模型）===")
    print("  NFE    精确场    第1轮    第2轮")
    for i, n in enumerate(nfe_list):
        print(f"{n:>5}{o[i]:>10.5f}{a[i]:>10.5f}{b[i]:>10.5f}")

path_lab.py

# -*- coding: utf-8 -*-
"""flow_matching 实验台（三）：路径的选择到底值多少 NFE。

直线路径（Rectified Flow）和 VP 余弦路径（DDPM 那一路）通向**同一个**目标分布，
中间分布不同，ODE 的轨迹与时间参数化也会不同。这一份把这件事量成三条曲线：

1. NFE-误差曲线：同一批起点，比较路径和积分器；Euler/RK2/RK4 每步分别评估 1/2/4 次场。
2. 欧式弧长 / 弦长，以及沿实际 ODE 轨迹的归一化速度偏差 S_traj；
   后者对 Rectified Flow 的轨迹 straightness 再除以 E||Z_1-Z_0||^2，
   不等于 reflow_lab 在原始配对插值上计算的 S_pair。
3. 时间 shift（SD3/FLUX 的做法）在直线路径上到底省不省步数。

误差尺子是 fm_oracle.discrepancy：拿精确边缘 p_t 的 RBF 特征矩当参照，
参照侧没有采样噪声，生成粒子一侧仍有有限样本误差。共享起点可减轻比较方差，
小差异是否稳定仍应更换随机种子验证；单次真实样本偏差不是硬下限。
"""

import numpy as np

from fm_oracle import (
    D, base_rng, make_data, LinearPath, CosineVPPath, OracleFlow,
    make_rbf_centers, discrepancy, make_schedule, integrate,
)

NPART = 4000
C, BW = make_rbf_centers()


def start_particles(n, rng):
    """t=1 端的粒子：精确就是 N(0, I)。"""
    return rng.standard_normal((n, D))


def sampling_reference(data, n=NPART):
    """固定一次真实样本与精确 p_data 的偏差，仅作有限采样参考，不是误差下限。"""
    rng = base_rng(12345)
    x = data.sample(n, rng)
    flow = OracleFlow(data, LinearPath())
    return discrepancy(x, 0.0, flow, C, BW)


def run_sweep(path, method, nfe_list, npart=NPART):
    rng = base_rng()
    z1 = start_particles(npart, rng)
    flow = OracleFlow(make_data(), path)
    out = []
    calls_per_step = {"euler": 1, "rk2": 2, "rk4": 4}[method]
    for nfe in nfe_list:
        if nfe < calls_per_step or nfe % calls_per_step:
            raise ValueError(f"{method}: NFE must be a positive multiple of {calls_per_step}")
        sched = make_schedule(nfe // calls_per_step)
        z0, _ = integrate(lambda z, t: flow.v(z, t), z1, sched, method)
        out.append(discrepancy(z0, 0.0, flow, C, BW))
    return np.array(out)


def straightness(path, npart=1500, nstep=400):
    """细步长积分出真实轨迹，量它离弦有多远。"""
    rng = base_rng(777)
    z1 = start_particles(npart, rng)
    flow = OracleFlow(make_data(), path)
    sched = make_schedule(nstep)
    z0, traj = integrate(lambda z, t: flow.v(z, t), z1, sched, "rk4")
    # traj: [nstep+1, npart, 2]，index 0 是 t=1
    arc = np.linalg.norm(np.diff(traj, axis=0), axis=2).sum(axis=0)  # [npart]，各段欧式长度相加
    chord = np.linalg.norm(z0 - z1, axis=1)
    ratio = float(np.mean(arc / np.maximum(chord, 1e-12)))

    # S_traj = mean_t ||v(z_t,t) - (z_1 - z_0)||^2 / mean ||z_1 - z_0||^2
    chord_dir = (z1 - z0)[None, :, :]                    # [1,npart,2]
    ts = sched[:-1]
    dev = 0.0
    for i, t in enumerate(ts):
        v = flow.v(traj[i], t)
        dev += float(np.mean(np.sum((v - chord_dir[0]) ** 2, axis=1)))
    dev /= len(ts)
    scale = float(np.mean(np.sum((z1 - z0) ** 2, axis=1)))
    return ratio, dev / scale


def shift_sweep(path, nfe_list, shifts, npart=NPART):
    rng = base_rng()
    z1 = start_particles(npart, rng)
    flow = OracleFlow(make_data(), path)
    res = {}
    for sh in shifts:
        errs = []
        for nfe in nfe_list:
            sched = make_schedule(nfe, shift=sh)
            z0, _ = integrate(lambda z, t: flow.v(z, t), z1, sched, "euler")
            errs.append(discrepancy(z0, 0.0, flow, C, BW))
        res[sh] = np.array(errs)
    return res


def state():
    """给画图脚本复用：返回与上面打印完全一致的数字。"""
    data = make_data()
    nfe_list = [2, 4, 8, 16, 32, 64, 128, 256]
    out = {
        "nfe": np.array(nfe_list),
        "sampling_reference": sampling_reference(data),
        "sweep": {},
        "straight": {},
        "shift": {},
    }
    for path in (LinearPath(), CosineVPPath()):
        for method in ("euler", "rk2"):
            out["sweep"][(path.name, method)] = run_sweep(path, method, nfe_list)
        out["straight"][path.name] = straightness(path)
    out["shift"] = shift_sweep(LinearPath(), [4, 8, 16, 32, 64], [1.0, 3.0, 6.0])
    out["shift_nfe"] = np.array([4, 8, 16, 32, 64])
    return out


if __name__ == "__main__":
    data = make_data()
    print(f"=== 单次采样参考：{NPART} 个真实样本 vs 精确 p_data 的偏差 = "
          f"{sampling_reference(data):.5f} ===")
    print("（这不是硬下限或显著性阈值；评估小差异需重复采样）")
    print()

    nfe_list = [2, 4, 8, 16, 32, 64, 128, 256]
    print("=== NFE-误差：直线路径 vs VP 余弦路径 ===")
    print("  NFE   linear/euler  linear/rk2   vp/euler     vp/rk2")
    lin_e = run_sweep(LinearPath(), "euler", nfe_list)
    lin_r = run_sweep(LinearPath(), "rk2", nfe_list)
    vp_e = run_sweep(CosineVPPath(), "euler", nfe_list)
    vp_r = run_sweep(CosineVPPath(), "rk2", nfe_list)
    for i, n in enumerate(nfe_list):
        print(f"{n:>5}{lin_e[i]:>14.5f}{lin_r[i]:>13.5f}{vp_e[i]:>12.5f}{vp_r[i]:>11.5f}")

    print()
    print("=== 轨迹指标（400 步 RK4，即 1600 NFE；数值近似 ODE 轨迹）===")
    for path in (LinearPath(), CosineVPPath()):
        ratio, s = straightness(path)
        print(f"  {path.label:<28} 弧长/弦长 = {ratio:.4f}   归一化 S_traj = {s:.4f}")

    print()
    print("=== 时间 shift（仅直线路径，Euler）===")
    res = shift_sweep(LinearPath(), [4, 8, 16, 32, 64], [1.0, 3.0, 6.0])
    print("  NFE   shift=1     shift=3     shift=6")
    for i, n in enumerate([4, 8, 16, 32, 64]):
        print(f"{n:>5}{res[1.0][i]:>11.5f}{res[3.0][i]:>11.5f}{res[6.0][i]:>11.5f}")

fm_oracle.py

# -*- coding: utf-8 -*-
"""flow_matching 实验台（一）：任意高斯路径下的精确速度场。

这一份是整个文章所有数字的来源。思想是：

    数据分布取二维高斯混合 p_data（K 个各向同性分量），噪声取 N(0, I)。
    对任意「高斯路径」 z_t = alpha_t * x + sigma_t * eps（t=1 是纯噪声，t=0 是数据），
    p_t 仍然是高斯混合（分量均值 alpha*mu_k，方差 alpha^2 s_k^2 + sigma^2），
    于是下面四样东西全都有闭式解，不需要训练、不需要采样近似：

        m(z,t)  = E[x | z_t = z]      去噪均值（x0 预测的真值）
        e(z,t)  = E[eps | z_t = z]    噪声均值（epsilon 预测的真值）
        v(z,t)  = alpha' m + sigma' e 边缘速度场（速度预测的真值）
        score   = grad log p_t(z)

    有了这些，三种参数化（x0 / eps / v）可以逐点互相换算并验到浮点误差，
    ODE 的 NFE-误差曲线可以用「精确边缘 p_t」当尺子，不需要跑一条高精参考解。

    误差尺子：取一组 RBF 特征 phi_c(z) = exp(-||z-c||^2 / (2h^2))，
    它在高斯分布下的期望有闭式，于是
        disc(粒子云, t) = mean_c | (1/N) sum_i phi_c(z_i)  -  E_{p_t}[phi_c] |
    经验侧仍有有限样本误差，解析侧不额外引入采样噪声。单次真实样本参考值既不是硬地板，
    也不是统计显著性阈值；有限组 RBF 特征矩一致也不能证明两个分布完全一致。
"""

import numpy as np

D = 2  # 二维，画图方便；公式与代码对任意维都成立


# ─────────────────────────── 随机流 ───────────────────────────
def base_rng(seed=20260928):
    """所有脚本与画图共用同一条随机流，保证正文数字与图上的数字一致。"""
    return np.random.default_rng(seed)


# ─────────────────────────── 数据分布 ───────────────────────────
class GMM2D:
    """二维各向同性高斯混合。"""

    def __init__(self, means, stds, weights=None):
        self.mu = np.asarray(means, dtype=float)          # [K, 2]
        self.s = np.asarray(stds, dtype=float)            # [K]
        K = self.mu.shape[0]
        if weights is None:
            self.w = np.full(K, 1.0 / K)
        else:
            self.w = np.asarray(weights, dtype=float)
            self.w = self.w / self.w.sum()

    @property
    def K(self):
        return self.mu.shape[0]

    def sample(self, n, rng):
        k = rng.choice(self.K, size=n, p=self.w)
        return self.mu[k] + self.s[k][:, None] * rng.standard_normal((n, D))

    def mean_cov(self):
        m = (self.w[:, None] * self.mu).sum(0)
        c = (self.w[:, None, None] * (
            (self.s ** 2)[:, None, None] * np.eye(D)[None]
            + (self.mu - m)[:, :, None] * (self.mu - m)[:, None, :]
        )).sum(0)
        return m, c


def make_data(K=8, radius=2.2, std=0.28):
    """8 个分量摆在半径 2.2 的圆上，每个分量标准差 0.28。"""
    ang = np.arange(K) * 2 * np.pi / K
    mu = np.stack([radius * np.cos(ang), radius * np.sin(ang)], axis=1)
    s = np.full(K, std)
    # 权重做成确定性的非均匀（1 + 0.35 cos），让混合更不像"一圈一样的点"
    w = 1.0 + 0.35 * np.cos(ang + 0.7)
    return GMM2D(mu, s, w / w.sum())


# ─────────────────────────── 路径 ───────────────────────────
class LinearPath:
    """Rectified Flow 的直线插值路径：z_t = (1-t) x + t eps。"""

    name = "linear"
    label = "直线路径（Rectified Flow）"

    def alpha(self, t):
        return 1.0 - t

    def sigma(self, t):
        return t

    def dalpha(self, t):
        return -1.0

    def dsigma(self, t):
        return 1.0


class CosineVPPath:
    """方差保持（VP）路径：z_t = cos(pi t/2) x + sin(pi t/2) eps。

    alpha^2 + sigma^2 = 1，也就是 DDPM 那一路 cosine schedule 的连续化版本。
    """

    name = "cosine_vp"
    label = "VP 余弦路径（DDPM 那一路）"

    def alpha(self, t):
        return np.cos(0.5 * np.pi * t)

    def sigma(self, t):
        return np.sin(0.5 * np.pi * t)

    def dalpha(self, t):
        return -0.5 * np.pi * np.sin(0.5 * np.pi * t)

    def dsigma(self, t):
        return 0.5 * np.pi * np.cos(0.5 * np.pi * t)


# ─────────────────────────── 精确场 ───────────────────────────
class OracleFlow:
    """给定数据与路径后的精确边缘速度场。"""

    def __init__(self, data: GMM2D, path):
        self.data = data
        self.path = path

    # 后验responsibility r_k(z,t) 与分量方差 V_k(t)
    def _post(self, z, t):
        a = self.path.alpha(t)
        sg = self.path.sigma(t)
        V = a * a * self.data.s ** 2 + sg * sg          # [K]
        diff = z[:, None, :] - a * self.data.mu[None, :, :]   # [N,K,2]
        d2 = (diff ** 2).sum(-1)                        # [N,K]
        logp = -0.5 * d2 / V[None, :] - np.log(V)[None, :] + np.log(self.data.w)[None, :]
        logp = logp - logp.max(1, keepdims=True)
        r = np.exp(logp)
        r = r / r.sum(1, keepdims=True)
        return r, V, diff

    def m(self, z, t):
        """E[x | z_t = z]，也就是 x0 预测的真值。"""
        a = self.path.alpha(t)
        r, V, diff = self._post(z, t)
        ex = self.data.mu[None, :, :] + (a * self.data.s ** 2 / V)[None, :, None] * diff
        return (r[:, :, None] * ex).sum(1)

    def e(self, z, t):
        """E[eps | z_t = z]，也就是 epsilon 预测的真值。"""
        sg = self.path.sigma(t)
        r, V, diff = self._post(z, t)
        ee = (sg / V)[None, :, None] * diff
        return (r[:, :, None] * ee).sum(1)

    def v(self, z, t):
        """边缘速度场 v*(z,t) = alpha' m + sigma' e。"""
        da = self.path.dalpha(t)
        ds = self.path.dsigma(t)
        return da * self.m(z, t) + ds * self.e(z, t)

    def score(self, z, t):
        """grad log p_t(z)。"""
        r, V, diff = self._post(z, t)
        return -(r[:, :, None] * diff / V[None, :, None]).sum(1)

    def marginal(self, t):
        """p_t 的分量参数（仍是 GMM）：均值 [K,2]、标准差 [K]、权重 [K]。"""
        a = self.path.alpha(t)
        sg = self.path.sigma(t)
        return a * self.data.mu, np.sqrt(a * a * self.data.s ** 2 + sg * sg), self.data.w


# ─────────────────────────── 误差尺子 ───────────────────────────
def make_rbf_centers(n=9, extent=3.6, bw=0.9):
    g = np.linspace(-extent, extent, n)
    C = np.stack(np.meshgrid(g, g, indexing="ij"), axis=-1).reshape(-1, D)
    return C, bw


def rbf_expectation(means, sds, weights, C, bw):
    """E_{p_t}[phi_c]，phi_c(z)=exp(-||z-c||^2/(2 bw^2))，p_t 为各向同性 GMM。

    单个高斯分量 N(m, v I) 下：E[phi_c] = (bw^2/(bw^2+v))^{d/2} * exp(-||m-c||^2/(2(bw^2+v)))
    """
    v = sds ** 2                                   # [K]
    coef = (bw ** 2 / (bw ** 2 + v)) ** (D / 2.0)  # [K]
    d2 = ((means[:, None, :] - C[None, :, :]) ** 2).sum(-1)     # [K, M]
    val = coef[:, None] * np.exp(-0.5 * d2 / (bw ** 2 + v)[:, None])
    return (weights[:, None] * val).sum(0)          # [M]


def discrepancy(z, t, flow: OracleFlow, C, bw):
    """粒子云 z 与精确边缘 p_t 的 RBF 特征矩偏差（越小越好）。"""
    mu_t, sd_t, w_t = flow.marginal(t)
    exact = rbf_expectation(mu_t, sd_t, w_t, C, bw)
    d2 = ((z[:, None, :] - C[None, :, :]) ** 2).sum(-1)
    emp = np.exp(-0.5 * d2 / bw ** 2).mean(0)
    return float(np.abs(emp - exact).mean())


# ─────────────────────────── 积分器 ───────────────────────────
def make_schedule(n_steps, shift=1.0):
    """含 n_steps 个积分区间（不是统一 NFE 预算）的时刻表。shift>1 是 SD3/FLUX 那套把时间往高噪声端推的做法。"""
    t = np.linspace(1.0, 0.0, n_steps + 1)
    if shift != 1.0:
        t = shift * t / (1.0 + (shift - 1.0) * t)
    return t


def integrate(vfun, z, schedule, method="euler"):
    """从 t=1 走到 t=0。vfun(z, t) 返回对递增时间定义的 dz/dt；负步长负责反向积分。"""
    z = z.copy()
    traj = [z.copy()]
    for i in range(len(schedule) - 1):
        t = schedule[i]
        h = schedule[i] - schedule[i + 1]     # >0
        if method == "euler":
            dz = vfun(z, t)
        elif method == "rk2":
            zm = z - 0.5 * h * vfun(z, t)
            dz = vfun(zm, t - 0.5 * h)
        elif method == "rk4":
            k1 = vfun(z, t)
            k2 = vfun(z - 0.5 * h * k1, t - 0.5 * h)
            k3 = vfun(z - 0.5 * h * k2, t - 0.5 * h)
            k4 = vfun(z - h * k3, t - h)
            dz = (k1 + 2 * k2 + 2 * k3 + k4) / 6.0
        else:
            raise ValueError(method)
        z = z - h * dz
        traj.append(z.copy())
    return z, np.stack(traj)   # traj: [S+1, N, 2]


# ─────────────────────────── 自测 ───────────────────────────
def selfcheck():
    rng = base_rng()
    data = make_data()
    out = []
    for path in (LinearPath(), CosineVPPath()):
        flow = OracleFlow(data, path)
        z = rng.standard_normal((4000, D)) * 1.6
        for t in (0.05, 0.25, 0.5, 0.75, 0.95):
            a = path.alpha(t)
            sg = path.sigma(t)
            m = flow.m(z, t)
            e = flow.e(z, t)
            v = flow.v(z, t)
            # 恒等式 1：alpha*m + sigma*e == z
            r1 = float(np.abs(a * m + sg * e - z).max())
            # 恒等式 2：v == alpha' m + sigma' e（定义，顺手确认没写反导数）
            r2 = float(np.abs(v - (path.dalpha(t) * m + path.dsigma(t) * e)).max())
            out.append((path.name, t, r1, r2))
    return out


if __name__ == "__main__":
    print("=== 恒等式自检：alpha*m + sigma*e 应等于 z（残差 ~1e-15）===")
    hdr = "path        t     |a*m+s*e-z|    |v-(da*m+ds*e)|"
    print(hdr)
    for name, t, r1, r2 in selfcheck():
        print(f"{name:<12}{t:>6.2f}{r1:>16.2e}{r2:>18.2e}")

    rng = base_rng()
    data = make_data()
    flow = OracleFlow(data, LinearPath())
    print()
    print("=== 直线路径下 v* 的两种算法是否一致：(z-m)/t 与 alpha' m + sigma' e ===")
    for t in (0.1, 0.3, 0.5, 0.7, 0.9):
        z = rng.standard_normal((2000, D)) * 1.6
        lhs = (z - flow.m(z, t)) / t
        rhs = flow.v(z, t)
        print(f"  t={t:.1f}  最大绝对差 = {np.abs(lhs - rhs).max():.3e}")

param_lab.py

# -*- coding: utf-8 -*-
"""flow_matching 实验台（二）：三种参数化（x0 / eps / v）的换算与 loss 权重。

要回答两个问题：

1. 固定同一条高斯路径和非退化时间，eps / x0 / v 输出能否代数换算？
   可以换算同一路径上的输出，但不能据此把 VP 模型直接变成另一条 RF 路径的模型。
   三种参数化的真值来自同一个去噪均值 m：
       z = alpha*m + sigma*e            （路径定义，恒等）
       v = alpha'*m + sigma'*e          （速度定义）
   两式联立解出 m、e，就得到「速度 -> 噪声/干净图」的换算；
   反过来也成立。下面把这些换算逐点验到 1e-14。

2. 输出可换算，为什么不同训练目标仍会产生差异？
   因为「等价」指的是**真值**等价，**loss 不等价**：同一个 m 上的误差 delta，
   在三种 loss 里被放大的倍数不同：
       x0  : 1
       eps : alpha / sigma
       v   : |alpha' - sigma' * alpha / sigma|
   这里列的是误差幅度系数，平方损失权重还需平方。直线路径 t→1 时，
   eps 系数趋于 0，v 系数趋于 1；这比较固定 m 误差的相对权重，
   不等于真实网络的所有梯度或质量。t=1 的 eps->m 换算本身退化。
"""

import numpy as np

from fm_oracle import (
    D, base_rng, make_data, LinearPath, CosineVPPath, OracleFlow,
)


# ─────────────────── 换算表的通用形式 ───────────────────
def v_to_me(z, v, path, t):
    """由 (z, v) 反解 (m, e)。

    解 [[alpha', sigma'], [alpha, sigma]] @ [m, e] = [v, z]
    行列式 det = alpha'*sigma - sigma'*alpha
    """
    a, sg = path.alpha(t), path.sigma(t)
    da, ds = path.dalpha(t), path.dsigma(t)
    det = da * sg - ds * a
    m = (sg * v - ds * z) / det
    e = (da * z - a * v) / det
    return m, e


def eps_to_me(z, eps_hat, path, t):
    """由 eps 预测反解 (m, e)。"""
    a, sg = path.alpha(t), path.sigma(t)
    m = (z - sg * eps_hat) / a
    return m, np.broadcast_to(eps_hat, m.shape).copy()


def me_to_v(m, e, path, t):
    return path.dalpha(t) * m + path.dsigma(t) * e


def check_conversions():
    """真值层面三种参数化互转，误差应到浮点量级。"""
    rng = base_rng()
    data = make_data()
    rows = []
    for path in (LinearPath(), CosineVPPath()):
        flow = OracleFlow(data, path)
        z = rng.standard_normal((3000, D)) * 1.6
        for t in (0.1, 0.3, 0.5, 0.7, 0.9, 0.99):
            m = flow.m(z, t)
            e = flow.e(z, t)
            v = flow.v(z, t)
            # 真值 v -> (m,e)
            m2, e2 = v_to_me(z, v, path, t)
            # 真值 eps -> m -> v
            m3, _ = eps_to_me(z, e, path, t)
            v3 = me_to_v(m3, e, path, t)
            rows.append((
                path.name, t,
                float(np.abs(m2 - m).max()),
                float(np.abs(e2 - e).max()),
                float(np.abs(m3 - m).max()),
                float(np.abs(v3 - v).max()),
            ))
    return rows


# ─────────────────── 放大倍数 ───────────────────
def weights(path, ts):
    """返回三种参数化下「m 上的单位误差」被放大的倍数（振幅，非平方）。"""
    wx = np.ones_like(ts)
    we, wv = [], []
    for t in ts:
        a, sg = path.alpha(t), path.sigma(t)
        da, ds = path.dalpha(t), path.dsigma(t)
        we.append(a / sg)
        wv.append(abs(da - ds * a / sg))
    return wx, np.array(we), np.array(wv)


def numeric_weight_check():
    """不靠公式，直接用数值扰动验证放大倍数：给 m 加一个固定扰动，看 loss 变化。"""
    rng = base_rng()
    data = make_data()
    flow = OracleFlow(data, LinearPath())
    path = LinearPath()
    z = rng.standard_normal((4000, D)) * 1.6
    out = []
    for t in (0.1, 0.3, 0.5, 0.7, 0.9, 0.99):
        m = flow.m(z, t)
        e = flow.e(z, t)
        v = flow.v(z, t)
        delta = rng.standard_normal((4000, D))
        delta = delta / np.linalg.norm(delta) * 0.01   # 固定长度 0.01 的扰动
        mp = m + delta
        ep = (z - path.alpha(t) * mp) / path.sigma(t)
        vp = me_to_v(mp, ep, path, t)
        r_eps = np.linalg.norm(ep - e) / np.linalg.norm(delta)
        r_v = np.linalg.norm(vp - v) / np.linalg.norm(delta)
        _, we, wv = weights(path, np.array([t]))
        out.append((t, r_eps, float(we[0]), r_v, float(wv[0])))
    return out


if __name__ == "__main__":
    print("=== 换算残差（真值层面互转，应为 1e-14 量级）===")
    print("path        t     v->m        v->e        eps->m      eps->v")
    for name, t, r1, r2, r3, r4 in check_conversions():
        print(f"{name:<11}{t:>5.2f}{r1:>12.2e}{r2:>12.2e}{r3:>12.2e}{r4:>12.2e}")

    print()
    print("=== 放大倍数：解析式 vs 数值扰动（直线路径）===")
    print("   t     eps实测   eps解析    v实测     v解析")
    for t, re_, we_, rv_, wv_ in numeric_weight_check():
        print(f"{t:>5.2f}{re_:>10.4f}{we_:>10.4f}{rv_:>10.4f}{wv_:>10.4f}")

    print()
    print("=== 放大倍数随 t 的变化（直线路径 vs VP 余弦路径）===")
    ts = np.array([0.01, 0.05, 0.1, 0.2, 0.4, 0.6, 0.8, 0.9, 0.95, 0.99])
    for path in (LinearPath(), CosineVPPath()):
        wx, we, wv = weights(path, ts)
        print(f"-- {path.label}")
        print("     t      x0      eps        v")
        for i, t in enumerate(ts):
            print(f"{t:>7.2f}{wx[i]:>8.3f}{we[i]:>10.3f}{wv[i]:>10.3f}")

make_figures.py

# -*- coding: utf-8 -*-
"""画配图。数字全部来自同目录的实验脚本，不另算一遍。

运行： python make_figures.py [--only 图名]
"""

import os
import sys

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from fm_oracle import (                                  # noqa: E402
    D, base_rng, make_data, LinearPath, CosineVPPath, OracleFlow,
    make_schedule, integrate,
)
from param_lab import weights                            # noqa: E402
import path_lab                                          # noqa: E402
import reflow_lab                                        # noqa: E402

HERE = os.path.dirname(os.path.abspath(__file__))
FIGDIR = os.path.join(os.path.dirname(HERE), "figures")
os.makedirs(FIGDIR, exist_ok=True)

C0 = "#2f4b7c"
C1 = "#d45087"
C2 = "#f0a35e"
C3 = "#4c9f70"
CGREY = "#8a8a8a"

plt.rcParams["font.sans-serif"] = ["Heiti TC", "Arial Unicode MS", "DejaVu Sans"]
plt.rcParams["axes.unicode_minus"] = False


def _save(fig, name):
    p = os.path.join(FIGDIR, name)
    fig.savefig(p, dpi=150, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"  [ok] {name}  ({os.path.getsize(p)} bytes)")


# ─────────────── 图 1：三种参数化的放大倍数 ───────────────
def fig_weights():
    ts = np.linspace(0.01, 0.999, 400)
    wx, we, wv = weights(LinearPath(), ts)
    fig, ax = plt.subplots(figsize=(7.4, 4.2))
    ax.plot(ts, wv, color=C1, lw=2.4, label="velocity 参数化  v")
    ax.plot(ts, we, color=C0, lw=2.4, label="噪声参数化  eps")
    ax.plot(ts, wx, color=CGREY, lw=1.8, ls="--", label="干净图参数化  x0")
    ax.set_yscale("log")
    ax.set_xlabel("时间 t（0 = 数据端，1 = 纯噪声端）")
    ax.set_ylabel("同一份误差被放大的倍数（对数轴）")
    ax.set_title("图 1：预测误差幅度系数（平方损失权重为其平方）")
    ax.axvline(0.99, color="#cccccc", lw=1, ls=":")
    ax.annotate("t=0.99 处 eps 系数为 0.010，\nvelocity 还有 1.010",
                xy=(0.99, 1.0), xytext=(0.55, 0.35), fontsize=10, color="#333333",
                arrowprops=dict(arrowstyle="->", color="#999999", lw=1))
    ax.legend(loc="upper right", fontsize=10)
    ax.grid(alpha=0.25)
    _save(fig, "fig1_param_weights.png")


# ─────────────── 图 2：两条路径的真实 ODE 轨迹 ───────────────
def fig_trajectories():
    data = make_data()
    rng = base_rng(4242)
    z1 = rng.standard_normal((600, D))
    xs = data.sample(1200, base_rng(11))

    fig, axes = plt.subplots(1, 2, figsize=(12.0, 5.4), sharex=True, sharey=True)
    sched = make_schedule(300)
    for ax, path in zip(axes, (LinearPath(), CosineVPPath())):
        flow = OracleFlow(data, path)
        z0, traj = integrate(lambda z, t: flow.v(z, t), z1, sched, "rk4")
        ax.scatter(xs[:, 0], xs[:, 1], s=8, color="#dddddd", label="数据样本（8 个模式）")
        for i in range(14):
            ax.plot(traj[:, i, 0], traj[:, i, 1], color=C1, lw=1.3, alpha=0.9)
        ax.scatter(z1[:14, 0], z1[:14, 1], s=34, color=C0, zorder=5, label="起点（纯噪声）")
        ax.scatter(z0[:14, 0], z0[:14, 1], s=34, color=C3, marker="s", zorder=5, label="终点（生成）")
        ax.set_title(path.label)
        ax.set_xlabel("dim 1")
        ax.grid(alpha=0.2)
    axes[0].set_ylabel("dim 2")
    axes[0].legend(fontsize=9, loc="upper left")
    fig.suptitle("图 2：同一批起点，两条路径走出来的 ODE 轨迹（300 步 RK4）", fontsize=13)
    fig.tight_layout()
    _save(fig, "fig2_trajectories.png")


# ─────────────── 图 3：NFE-误差 ───────────────
def fig_nfe():
    nfe_list = [2, 4, 8, 16, 32, 64, 128, 256]
    floor = path_lab.sampling_reference(make_data())
    cur = {
        "linear/euler": path_lab.run_sweep(LinearPath(), "euler", nfe_list),
        "linear/rk2": path_lab.run_sweep(LinearPath(), "rk2", nfe_list),
        "vp/euler": path_lab.run_sweep(CosineVPPath(), "euler", nfe_list),
        "vp/rk2": path_lab.run_sweep(CosineVPPath(), "rk2", nfe_list),
    }
    fig, ax = plt.subplots(figsize=(7.6, 4.6))
    for (k, v), col, mk in zip(cur.items(), [C1, C0, C2, C3], ["o", "o", "s", "s"]):
        ax.loglog(nfe_list, v, marker=mk, color=col, lw=2, label=k)
    ax.loglog(nfe_list, [floor] * len(nfe_list), color=CGREY, ls="--", lw=1.6,
              label="单次采样参考（4000 个真实样本）")
    ax.set_xlabel("NFE（模型评估次数）")
    ax.set_ylabel("终点分布与精确 p_data 的偏差")
    ax.set_title("图 3：1-rectified 的直线路径 vs VP 余弦路径")
    ax.grid(alpha=0.25, which="both")
    ax.legend(fontsize=10)
    _save(fig, "fig3_nfe_paths.png")
    return {"nfe": nfe_list, "curves": {k: v.tolist() for k, v in cur.items()}}


# ─────────────── 图 4：Reflow ───────────────
def fig_reflow():
    data = make_data()
    flow = OracleFlow(data, LinearPath())
    z0_1, z1 = reflow_lab.round1_pairs(100000, seed=0)
    net1, _ = reflow_lab.train(z0_1, z1, n_step=6000, seed=0, fresh=True)
    z0_2 = reflow_lab.generate_pairs(net1, z1, n_steps=100)
    net2, _ = reflow_lab.train(z0_2, z1, n_step=6000, seed=0)

    nfe_list = [2, 4, 8, 16, 32, 64]
    eval_z1 = base_rng(999).standard_normal((4000, D))
    o = reflow_lab.nfe_curve(eval_z1, lambda z, t: flow.v(z, t), nfe_list, data)
    a = reflow_lab.nfe_curve(eval_z1, lambda z, t: net1(z, t), nfe_list, data)
    b = reflow_lab.nfe_curve(eval_z1, lambda z, t: net2(z, t), nfe_list, data)
    floor = path_lab.sampling_reference(data)

    arc1 = reflow_lab.arc_over_chord(z1, lambda z, t: net1(z, t))[0]
    arc2 = reflow_lab.arc_over_chord(z1, lambda z, t: net2(z, t))[0]
    s1 = reflow_lab.straightness_of(z0_1, z1, lambda z, t: flow.v(z, t))
    s2 = reflow_lab.straightness_of(z0_2, z1, lambda z, t: net2(z, t))

    fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.8))
    ax = axes[0]
    ax.loglog(nfe_list, o, marker="o", color=CGREY, lw=2, label="精确场（数值积分对照）")
    ax.loglog(nfe_list, a, marker="o", color=C1, lw=2, label="第 1 轮模型")
    ax.loglog(nfe_list, b, marker="s", color=C3, lw=2, label="第 2 轮模型（reflow 后）")
    ax.loglog(nfe_list, [floor] * len(nfe_list), color="#bbbbbb", ls="--", lw=1.5,
              label="单次采样参考（非下限）")
    ax.set_xlabel("NFE")
    ax.set_ylabel("终点偏差")
    ax.set_title("本实验 reflow 后：NFE≥2 的偏差变化很小")
    ax.grid(alpha=0.25, which="both")
    ax.legend(fontsize=9)

    ax = axes[1]
    names = ["第 1 轮", "第 2 轮"]
    x = np.arange(2)
    ax.bar(x - 0.18, [arc1, arc2], width=0.34, color=C0, label="弧长/弦长（1 = 完全直线）")
    ax.bar(x + 0.18, [s1, s2], width=0.34, color=C2, label="归一化配对残差 S_pair")
    ax.set_xticks(x)
    ax.set_xticklabels(names)
    ax.set_yscale("log")
    ax.set_ylabel("配对残差与轨迹长度比（对数轴）")
    ax.set_title("两轮之间：轨迹从弯的变成直的")
    ax.grid(alpha=0.25, axis="y")
    ax.legend(fontsize=9)
    for i, (av, sv) in enumerate([(arc1, s1), (arc2, s2)]):
        ax.text(i - 0.18, av * 1.08, f"{av:.4f}", ha="center", fontsize=9)
        ax.text(i + 0.18, sv * 1.08, f"{sv:.4f}", ha="center", fontsize=9)
    fig.suptitle("图 4：Reflow 的两轮对比（同一个直线路径、同一批起点）", fontsize=13)
    fig.tight_layout()
    _save(fig, "fig4_reflow.png")
    return {"nfe1": a.tolist(), "nfe2": b.tolist(), "S1": s1, "S2": s2,
            "arc1": arc1, "arc2": arc2}


if __name__ == "__main__":
    only = sys.argv[2] if len(sys.argv) > 2 and sys.argv[1] == "--only" else None
    jobs = {
        "weights": fig_weights,
        "trajectories": fig_trajectories,
        "nfe": fig_nfe,
        "reflow": fig_reflow,
    }
    for k, fn in jobs.items():
        if only and k != only:
            continue
        print(f"[draw] {k}")
        fn()
    print("done")