RFdiffusion sample_vec函数解读

sample_vec 是IGSO3类一个高层次的采样方法，用于从 IGSO(3) 分布中生成旋转向量。它的实现依赖于多个其他函数，调用关系如下：

1. sample_vec 调用 sample:

   • sample 从 IGSO(3) 分布的角度分布中采样旋转角度 ω。
   • sample_vec 利用这些角度 ω 和随机方向生成旋转向量 v。

2. sample 调用 t_to_idx:

   • t_to_idx 将给定的时间步 tt 映射到对应的离散化 σ 索引。
   • 这是因为不同时间步 tt 对应不同的噪声尺度 σ(t)。

3. t_to_idx 调用 sigma_idx 和 sigma:

   • sigma_idx 将连续 σ(t) 映射到预离散化的 σ 索引。
   • sigma 根据时间步 tt 和噪声日程（例如指数或线性）计算对应的噪声尺度 σ(t)。

sample_vec → sample → t_to_idx → sigma_idx + sigma。

sample_vec函数通过调用其他函数实现从 IGSO(3) 分布中采样的方法。它包括计算离散化的 σ 值索引、将时间步 t 映射为对应的 σ、以及使用逆 CDF 方法生成旋转角度和旋转向量。

1. discrete_sigma 属性源代码：

   @propertydef discrete_sigma(self):return self.igso3_vals["discrete_sigma"]

• 功能: 返回预计算的离散化 σ 值，这些值是从 IGSO(3) 分布的离散化过程中得到的。
• 用途: 用于索引和采样，与分布的噪声尺度有关。

2.sigma_idx 方法源代码：

    def sigma_idx(self, sigma: np.ndarray):"""Calculates the index for discretized sigma during IGSO(3) initialization."""return np.digitize(sigma, self.discrete_sigma) - 1

功能: 将输入的 σ 值映射到 discrete_sigma 中的索引。

• 使用 np.digitize 来查找 σ 所属的离散化区间。
• 减去 1，使索引从 0 开始。

3. t_to_idx 方法源代码：

    def t_to_idx(self, t: np.ndarray):"""Helper function to go from discrete time index t to corresponding sigma_idx.Args:t: time index (integer between 1 and 200)"""continuous_t = t / self.Treturn self.sigma_idx(self.sigma(continuous_t))

• 功能: 将时间步 t 映射为对应的 σ 索引。