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)
| 94 | return -drift, diffusion |
| 95 | |
| 96 | def compute_diffusion(self, x, t, form="constant", norm=1.0): |
| 97 | """Compute the diffusion term of the SDE |
| 98 | Args: |
| 99 | x: [batch_dim, ...], data point |
| 100 | t: [batch_dim,], time vector |
| 101 | form: str, form of the diffusion term |
| 102 | norm: float, norm of the diffusion term |
| 103 | """ |
| 104 | t = pad_v_like_x(t, x) |
| 105 | choices = { |
| 106 | "constant": norm, |
| 107 | "SBDM": norm * self.compute_drift(x, t)[1], |
| 108 | "sigma": norm * self.compute_sigma_t(t)[0], |
| 109 | "linear": norm * (1 - t), |
| 110 | "decreasing": 0.25 * (norm * torch.cos(np.pi * t) + 1) ** 2, |
| 111 | "increasing-decreasing": norm * torch.sin(np.pi * t) ** 2, |
| 112 | } |
| 113 | |
| 114 | try: diffusion = choices[form] |
| 115 | except KeyError: raise NotImplementedError(f"Diffusion form {form} not implemented") |
| 116 | |
| 117 | return diffusion |
| 118 | |
| 119 | def get_score_from_velocity(self, velocity, x, t): |
| 120 | """Wrapper function: transfrom velocity prediction model to score |
no test coverage detected