1import numpy as np
 2
 3def deriv(z, omega, mu, R, drive=None):
 4    n = z.shape[-1] // 2
 5    x, y = z[..., :n], z[..., n:]
 6    r2 = x*x + y*y
 7    out = np.concatenate((omega*y, -omega*x + mu*(1-r2/(R*R))*y), axis=-1)
 8    return out if drive is None else out + drive
 9
10def rk4_step(z, h, omega, mu, R, drive=None):
11    f = lambda q: deriv(q, omega, mu, R, drive)
12    k1=f(z); k2=f(z+h*k1/2); k3=f(z+h*k2/2); k4=f(z+h*k3)
13    return z + h*(k1+2*k2+2*k3+k4)/6
14
15def simulate(z0, steps, h, omega, mu, R):
16    out=np.empty((steps+1,len(z0))); out[0]=z0
17    for t in range(steps): out[t+1]=rk4_step(out[t],h,omega,mu,R)
18    return out
19
20def radius(z):
21    n=z.shape[-1]//2
22    return np.sqrt(np.sum(z[...,:n]**2+z[...,n:]**2,axis=-1))