Compute the diffusion term of the SDE Args: x: [batch_dim, ...], data point t: [batch_dim,], time vector form: str, form of the diffusion term norm: float, norm of the diffusion term
(self, x, t, form="constant", norm=1.0)
| 43 | return -drift, diffusion |
| 44 | |
| 45 | def compute_diffusion(self, x, t, form="constant", norm=1.0): #compute w_t in Eq 4. |
| 46 | """Compute the diffusion term of the SDE |
| 47 | Args: |
| 48 | x: [batch_dim, ...], data point |
| 49 | t: [batch_dim,], time vector |
| 50 | form: str, form of the diffusion term |
| 51 | norm: float, norm of the diffusion term |
| 52 | """ |
| 53 | t = expand_t_like_x(t, x) |
| 54 | choices = { |
| 55 | "constant": norm, |
| 56 | "SBDM": norm * self.compute_drift(x, t)[1], #follow the calculation of w_t for SBDM |
| 57 | "sigma": norm * self.compute_sigma_t(t)[0], # This suggests the choice wt = σt in (4) to cancel this singularity (see Appendix A.3) |
| 58 | "linear": norm * (1 - t), #Table 2, seems never used |
| 59 | "decreasing": 0.25 * (norm * th.cos(np.pi * t) + 1) ** 2, #Table2, seems never used |
| 60 | "inccreasing-decreasing": norm * th.sin(np.pi * t) ** 2, #seems never used |
| 61 | } |
| 62 | |
| 63 | try: |
| 64 | diffusion = choices[form] |
| 65 | except KeyError: |
| 66 | raise NotImplementedError(f"Diffusion form {form} not implemented") |
| 67 | |
| 68 | return diffusion |
| 69 | |
| 70 | def get_score_from_velocity(self, velocity, x, t): #Eq 9 in SiT paper by simple algebra |
| 71 | """Wrapper function: transfrom velocity prediction model to score |
no test coverage detected