SDE solver class
| 183 | |
| 184 | |
| 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) |
| 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) |
| 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) |
| 223 | |
| 224 | |
| 225 | class FlowSDE: |