| 198 | except: raise NotImplementedError(f"Sampler type '{sampler_type}' not implemented.") |
| 199 | |
| 200 | def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs): |
| 201 | w_cur = torch.randn(x.size()).to(x) |
| 202 | t = torch.ones(x.size(0)).to(x) * t |
| 203 | dw = w_cur * torch.sqrt(self.dt) |
| 204 | drift = self.drift(x, t, model, **model_kwargs) |
| 205 | diffusion = self.diffusion(x, t) |
| 206 | mean_x = x + drift * self.dt |
| 207 | x = mean_x + torch.sqrt(2 * diffusion) * dw |
| 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) |