import numpy as np import torch as th import torch.nn as nn from torchdiffeq import odeint from functools import partial from tqdm import tqdm class sde: """SDE solver class""" def __init__( self, drift, diffusion, *, t0, t1, num_steps, sampler_type, ): assert t0 < t1, "SDE sampler has to be in forward time" self.num_timesteps = num_steps self.t = th.linspace(t0, t1, num_steps) self.dt = self.t[1] - self.t[0] self.drift = drift self.diffusion = diffusion self.sampler_type = sampler_type def __Euler_Maruyama_step(self, x, mean_x, t, model, **model_kwargs): w_cur = th.randn(x.size()).to(x) t = th.ones(x.size(0)).to(x) * t dw = w_cur * th.sqrt(self.dt) drift = self.drift(x, t, model, **model_kwargs) diffusion = self.diffusion(x, t) mean_x = x + drift * self.dt x = mean_x + th.sqrt(2 * diffusion) * dw return x, mean_x def __Heun_step(self, x, _, t, model, **model_kwargs): w_cur = th.randn(x.size()).to(x) dw = w_cur * th.sqrt(self.dt) t_cur = th.ones(x.size(0)).to(x) * t diffusion = self.diffusion(x, t_cur) xhat = x + th.sqrt(2 * diffusion) * dw K1 = self.drift(xhat, t_cur, model, **model_kwargs) xp = xhat + self.dt * K1 K2 = self.drift(xp, t_cur + self.dt, model, **model_kwargs) return xhat + 0.5 * self.dt * (K1 + K2), xhat # at last time point we do not perform the heun step def __forward_fn(self): """TODO: generalize here by adding all private functions ending with steps to it""" sampler_dict = { "Euler": self.__Euler_Maruyama_step, "Heun": self.__Heun_step, } try: sampler = sampler_dict[self.sampler_type] except: raise NotImplementedError("Smapler type not implemented.") return sampler def sample(self, init, model, **model_kwargs): """forward loop of sde""" x = init mean_x = init samples = [] sampler = self.__forward_fn() for ti in self.t[:-1]: with th.no_grad(): x, mean_x = sampler(x, mean_x, ti, model, **model_kwargs) samples.append(x) return samples class ode: """ODE solver class""" def __init__( self, drift, *, t0, t1, sampler_type, num_steps, atol, rtol, ): assert t0 < t1, "ODE sampler has to be in forward time" self.drift = drift self.t = th.linspace(t0, t1, num_steps) self.atol = atol self.rtol = rtol self.sampler_type = sampler_type def sample(self, x, model, **model_kwargs): device = x[0].device if isinstance(x, tuple) else x.device def _fn(t, x): t = th.ones(x[0].size(0)).to(device) * t if isinstance(x, tuple) else th.ones(x.size(0)).to(device) * t model_output = self.drift(x, t, model, **model_kwargs) return model_output t = self.t.to(device) atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol] rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol] samples = odeint( _fn, x, t, method=self.sampler_type, atol=atol, rtol=rtol ) return samples