<jupyter_start><jupyter_text>Test method vs `torchdiffeq`  Runge-Kutta methods<jupyter_code>import collections
ButcherTableau = collections.namedtuple('ButcherTableau', 'a, b, c, c_err')
class GenericExplicitButcher(nn.Module):
    def __init__(self, tableau):
        super().__init__()
        self.tab = tableau
    
    def forward(self, s, z, f, eps):
        stages = []
        for i, a in enumerate(self.tab.a): 
            stval = torch.zeros_like(z)
            for j, stage in enumerate(stages):
                # determine stage value
                stval += a[j]*stages[j]
            # evaluate f and store stage
            stages += [f(s + tableau.c[i], z + eps*stval)]
        return z + eps*sum([tableau.b[i]*stages[i] for i in range(len(tableau.b))])<jupyter_output><empty_output><jupyter_text>Test Euler<jupyter_code>tableau = ButcherTableau([[0]], [1], [0], [])
solver = GenericExplicitButcher(tableau)

def f(s, z): return -0.99*z
z = torch.Tensor([0.9])[None]
s_span = torch.linspace(0, 1, 100)

def odeint(f, z, solver, s_span):
    eps = s_span[1] - s_span[0]
    traj = []
    for i in s_span:
        traj += [z]
        z = solver(s, z, f, eps)
    return torch.cat(traj)

traj = odeint(f, z, solver, s_span)
import matplotlib.pyplot as plt; plt.plot(traj)
z = torch.Tensor([0.9])
traj_torchdiffeq = torchdiffeq.odeint(f, z, method='euler', t=torch.linspace(0, 1, 100))[:, 0]
assert((traj - traj_torchdiffeq <= 1e-2).all())
%timeit z = torch.Tensor([0.9]); traj_torchdiffeq = odeint(f, z, solver, s_span)
%timeit z = torch.Tensor([0.9]); traj = torchdiffeq.odeint(f, z, method='euler', t=torch.linspace(0, 1, 100))[:, 0]<jupyter_output>3.98 ms ± 49 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)<jupyter_text>Test Explicit Midpoint<jupyter_code>tableau = ButcherTableau([[0, 0], [0.5, 0]], [0, 1], [0, 0.5], [])
solver = GenericExplicitButcher(tableau)

def f(s, z): return -0.99*z
z = torch.Tensor([0.9])
s_span = torch.linspace(0, 1, 100)

traj = odeint(f, z, solver, s_span)
plt.plot(traj)
z = torch.Tensor([0.9])
traj_torchdiffeq = torchdiffeq.odeint(f, z, method='midpoint', t=torch.linspace(0, 1, 100))[:, 0]
assert((traj - traj_torchdiffeq <= 1e-3).all())
%timeit z = torch.Tensor([0.9]); traj_torchdiffeq = odeint(f, z, solver, s_span)
%timeit z = torch.Tensor([0.9]); traj = torchdiffeq.odeint(f, z, method='midpoint', t=torch.linspace(0, 1, 100))[:, 0]<jupyter_output>5.32 ms ± 89.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)<jupyter_text>Test Heun<jupyter_code>tableau = ButcherTableau([[0, 0], [1, 0]], [0, 1], [0.5, 0.5], [])
solver = GenericExplicitButcher(tableau)

def f(s, z): return -0.99*z
z = torch.Tensor([0.9])
s_span = torch.linspace(0, 1, 100)

traj = odeint(f, z, solver, s_span)
plt.plot(traj)<jupyter_output><empty_output><jupyter_text>Test Ralston<jupyter_code>tableau = ButcherTableau([[0, 0], [2/3, 0]], [0, 2/3], [1/4, 3/4], [])
solver = GenericExplicitButcher(tableau)

def f(s, z): return -0.99*z
z = torch.Tensor([0.9])
s_span = torch.linspace(0, 1, 100)

traj = odeint(f, z, solver, s_span)
plt.plot(traj)<jupyter_output><empty_output><jupyter_text>Test Generic 2nd order (\alpha)<jupyter_code>alpha = 1

tableau = ButcherTableau([[0, 0], [alpha, 0]], [0, alpha], [1-1/(2*alpha), 1/(2*alpha)], [])
solver = GenericExplicitButcher(tableau)

def f(s, z): return -0.99*z
z = torch.Tensor([0.9])
s_span = torch.linspace(0, 1, 100)

traj = odeint(f, z, solver, s_span)
plt.plot(traj)<jupyter_output><empty_output><jupyter_text>Test "Classic" RK4<jupyter_code>tableau = ButcherTableau([[0, 0, 0, 0], [0.5, 0, 0, 0], [0, 0.5, 0, 0], [0, 0, 1, 0]], [0, 0.5, 0.5, 1], [1/6, 1/3, 1/3, 1/6], [])
solver = GenericExplicitButcher(tableau)

def f(s, z): return -0.99*z
z = torch.Tensor([0.9])
s_span = torch.linspace(0, 1, 100)

traj = odeint(f, z, solver, s_span)
plt.plot(traj)
z = torch.Tensor([0.9])
traj_torchdiffeq = torchdiffeq.odeint(f, z, method='rk4', t=torch.linspace(0, 1, 100))[:, 0]
assert((traj - traj_torchdiffeq <= 1e-3).all())
%timeit z = torch.Tensor([0.9]); traj = odeint(f, z, solver, s_span)
%timeit z = torch.Tensor([0.9]); traj_torchdiffeq = torchdiffeq.odeint(f, z, method='rk4', t=torch.linspace(0, 1, 100))[:, 0]<jupyter_output>11.1 ms ± 62.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
StarcoderdataJupyterStructuredCleanDedup