class for VP path flow matching
| 137 | |
| 138 | |
| 139 | class VPCPlan(ICPlan): |
| 140 | """class for VP path flow matching""" |
| 141 | |
| 142 | def __init__(self, sigma_min=0.1, sigma_max=20.0): |
| 143 | self.sigma_min = sigma_min |
| 144 | self.sigma_max = sigma_max |
| 145 | self.log_mean_coeff = lambda t: -0.25 * ((1 - t) ** 2) * (self.sigma_max - self.sigma_min) - 0.5 * (1 - t) * self.sigma_min |
| 146 | self.d_log_mean_coeff = lambda t: 0.5 * (1 - t) * (self.sigma_max - self.sigma_min) + 0.5 * self.sigma_min |
| 147 | |
| 148 | |
| 149 | def compute_alpha_t(self, t): |
| 150 | """Compute coefficient of x1""" |
| 151 | alpha_t = self.log_mean_coeff(t) |
| 152 | alpha_t = th.exp(alpha_t) |
| 153 | d_alpha_t = alpha_t * self.d_log_mean_coeff(t) |
| 154 | return alpha_t, d_alpha_t |
| 155 | |
| 156 | def compute_sigma_t(self, t): |
| 157 | """Compute coefficient of x0""" |
| 158 | p_sigma_t = 2 * self.log_mean_coeff(t) |
| 159 | sigma_t = th.sqrt(1 - th.exp(p_sigma_t)) |
| 160 | d_sigma_t = th.exp(p_sigma_t) * (2 * self.d_log_mean_coeff(t)) / (-2 * sigma_t) |
| 161 | return sigma_t, d_sigma_t |
| 162 | |
| 163 | def compute_d_alpha_alpha_ratio_t(self, t): |
| 164 | """Special purposed function for computing numerical stabled d_alpha_t / alpha_t""" |
| 165 | return self.d_log_mean_coeff(t) |
| 166 | |
| 167 | def compute_drift(self, x, t): |
| 168 | """Compute the drift term of the SDE""" |
| 169 | t = expand_t_like_x(t, x) |
| 170 | beta_t = self.sigma_min + (1 - t) * (self.sigma_max - self.sigma_min) |
| 171 | return -0.5 * beta_t * x, beta_t / 2 |
| 172 | |
| 173 | |
| 174 | class GVPCPlan(ICPlan): |
nothing calls this directly
no outgoing calls
no test coverage detected