ODE solver class
| 81 | |
| 82 | |
| 83 | class ode: |
| 84 | """ODE solver class""" |
| 85 | |
| 86 | def __init__( |
| 87 | self, |
| 88 | drift, |
| 89 | *, |
| 90 | t0, |
| 91 | t1, |
| 92 | sampler_type, |
| 93 | num_steps, |
| 94 | atol, |
| 95 | rtol, |
| 96 | ): |
| 97 | # assert t0 < t1, "ODE sampler has to be in forward time", , comment it out, to make x2z ODE works |
| 98 | |
| 99 | self.drift = drift |
| 100 | self.t = th.linspace(t0, t1, num_steps) |
| 101 | self.atol = atol |
| 102 | self.rtol = rtol |
| 103 | self.sampler_type = sampler_type |
| 104 | |
| 105 | def sample(self, x, model, **model_kwargs): |
| 106 | |
| 107 | device = x[0].device if isinstance(x, tuple) else x.device |
| 108 | |
| 109 | def _fn(t, x): |
| 110 | t = ( |
| 111 | th.ones(x[0].size(0)).to(device) * t |
| 112 | if isinstance(x, tuple) |
| 113 | else th.ones(x.size(0)).to(device) * t |
| 114 | ) |
| 115 | model_output = self.drift(x, t, model, **model_kwargs) |
| 116 | return model_output |
| 117 | |
| 118 | t = self.t.to(device) |
| 119 | |
| 120 | atol = [self.atol] * len(x) if isinstance(x, tuple) else [self.atol] |
| 121 | rtol = [self.rtol] * len(x) if isinstance(x, tuple) else [self.rtol] |
| 122 | samples = odeint(_fn, x, t, method=self.sampler_type, atol=atol, rtol=rtol) |
| 123 | return samples |
no outgoing calls
no test coverage detected