(
self,
drift,
diffusion,
*,
t0,
t1,
num_steps,
sampler_type,
)
| 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) |
nothing calls this directly
no outgoing calls
no test coverage detected