import numpy as np

def deriv(z, omega, mu, R, drive=None):
    n = z.shape[-1] // 2
    x, y = z[..., :n], z[..., n:]
    r2 = x*x + y*y
    out = np.concatenate((omega*y, -omega*x + mu*(1-r2/(R*R))*y), axis=-1)
    return out if drive is None else out + drive

def rk4_step(z, h, omega, mu, R, drive=None):
    f = lambda q: deriv(q, omega, mu, R, drive)
    k1=f(z); k2=f(z+h*k1/2); k3=f(z+h*k2/2); k4=f(z+h*k3)
    return z + h*(k1+2*k2+2*k3+k4)/6

def simulate(z0, steps, h, omega, mu, R):
    out=np.empty((steps+1,len(z0))); out[0]=z0
    for t in range(steps): out[t+1]=rk4_step(out[t],h,omega,mu,R)
    return out

def radius(z):
    n=z.shape[-1]//2
    return np.sqrt(np.sum(z[...,:n]**2+z[...,n:]**2,axis=-1))