calculate_igso3
函数旨在对 IGSO(3) 分布的概率密度函数 (PDF)、累积分布函数 (CDF)、以及相关统计量进行数值近似计算,特别用于预计算以加速后续操作(例如采样、反向扩散等)。
calculate_igso3函数源代码:
def calculate_igso3(*, num_sigma, num_omega, min_sigma, max_sigma):
"""calculate_igso3 pre-computes numerical approximations to the IGSO3 cdfs
and score norms and expected squared score norms.
Args:
num_sigma: number of different sigmas for which to compute igso3
quantities.
num_omega: number of point in the discretization in the angle of
rotation.
min_sigma, max_sigma: the upper and lower ranges for the angle of
rotation on which to consider the IGSO3 distribution. This cannot
be too low or it will create numerical instability.
"""
# Discretize omegas for calculating CDFs. Skip omega=0.
discrete_omega = np.linspace(0, np.pi, num_omega+1)[1:]
# Exponential noise schedule. This choice is closely tied to the
# scalings used when simulating the reverse time SDE. For each step n,
# discrete_sigma[n] = min_eps^(1-n/num_eps) * max_eps^(n/num_eps)
discrete_sigma = 10 ** np.linspace(np.log10(min_sigma), np.log10(max_sigma), num_sigma + 1)[1:]
# Compute the pdf and cdf values for the marginal distribution of the angle
# of rotation (which is needed for sampling)
pdf_vals = np.asarray(
[igso3_density_angle(discrete_omega, sigma**2) for sigma in discrete_sigma])
cdf_vals = np.asarray(
[pdf.cumsum() / num_omega * np.pi for pdf in pdf_vals])
# Compute the norms of the scores. This are used to scale the rotation axis when
# computing the score as a vector.
score_norm = np.asarray(
[d_logf_d_omega(discrete_omega, sigma**2)