| 185 | class StepSDE: |
| 186 | """SDE solver class""" |
| 187 | def __init__(self, dt, drift, diffusion, sampler_type): |
| 188 | self.dt = dt |
| 189 | self.drift = drift |
| 190 | self.diffusion = diffusion |
| 191 | self.sampler_type = sampler_type |
| 192 | self.sampler_dict = { |
| 193 | "euler": self.__Euler_Maruyama_step, |
| 194 | "heun": self.__Heun_step, |
| 195 | } |
| 196 | |
| 197 | try: self.sampler = self.sampler_dict[sampler_type] |
| 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) |