SDE solver class
| 7 | |
| 8 | |
| 9 | class sde: |
| 10 | """SDE solver class""" |
| 11 | |
| 12 | def __init__( |
| 13 | self, |
| 14 | drift, |
| 15 | diffusion, |
| 16 | *, |
| 17 | t0, |
| 18 | t1, |
| 19 | num_steps, |
| 20 | sampler_type, |
| 21 | ): |
| 22 | assert t0 < t1, "SDE sampler has to be in forward time" |
| 23 | |
| 24 | self.num_timesteps = num_steps |
| 25 | self.t = th.linspace(t0, t1, num_steps) |
| 26 | self.dt = self.t[1] - self.t[0] |
| 27 | self.drift = drift |
| 28 | self.diffusion = diffusion |
| 29 | self.sampler_type = sampler_type |
| 30 | |
| 31 | def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs): |
| 32 | w_cur = th.randn(x.size()).to(x) |
| 33 | t = th.ones(x.size(0)).to(x) * t |
| 34 | dw = w_cur * th.sqrt(self.dt) |
| 35 | drift = self.drift(x, t, model, **model_kwargs) |
| 36 | diffusion = self.diffusion(x, t) |
| 37 | mean_x = x + drift * self.dt |
| 38 | x = mean_x + th.sqrt(2 * diffusion) * dw |
| 39 | return x, mean_x |
| 40 | |
| 41 | def __Heun_step(self, x, _, t, model, **model_kwargs): |
| 42 | w_cur = th.randn(x.size()).to(x) |
| 43 | dw = w_cur * th.sqrt(self.dt) |
| 44 | t_cur = th.ones(x.size(0)).to(x) * t |
| 45 | diffusion = self.diffusion(x, t_cur) |
| 46 | xhat = x + th.sqrt(2 * diffusion) * dw |
| 47 | K1 = self.drift(xhat, t_cur, model, **model_kwargs) |
| 48 | xp = xhat + self.dt * K1 |
| 49 | K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs) |
| 50 | return ( |
| 51 | xhat + 0.5 * self.dt * (K1 + K2), |
| 52 | xhat, |
| 53 | ) # at last time point we do not perform the heun step |
| 54 | |
| 55 | def __forward_fn(self): |
| 56 | """TODO: generalize here by adding all private functions ending with steps to it""" |
| 57 | sampler_dict = { |
| 58 | "Euler": self.__Euler_Maruyama_step, |
| 59 | "Heun": self.__Heun_step, |
| 60 | } |
| 61 | |
| 62 | try: |
| 63 | sampler = sampler_dict[self.sampler_type] |
| 64 | except: |
| 65 | raise NotImplementedError("Smapler type not implemented.") |
| 66 |