| 208 | return x, mean_x |
| 209 | |
| 210 | def __Heun_step(self, x, mean_x, t, model, **model_kwargs): |
| 211 | w_cur = torch.randn(x.size()).to(x) |
| 212 | dw = w_cur * torch.sqrt(self.dt) |
| 213 | t_cur = torch.ones(x.size(0)).to(x) * t |
| 214 | diffusion = self.diffusion(x, t_cur) |
| 215 | xhat = x + torch.sqrt(2 * diffusion) * dw |
| 216 | K1 = self.drift(xhat, t_cur, model, **model_kwargs) |
| 217 | xp = xhat + self.dt * K1 |
| 218 | K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs) |
| 219 | return xhat + 0.5 * self.dt * (K1 + K2), xhat # at last time point we do not perform the heun step |
| 220 | |
| 221 | def __call__(self, x, mean_x, t, model, **model_kwargs): |
| 222 | return self.sampler(x, mean_x, t, model, **model_kwargs) |