from scipy.optimize import root

def func_ydot(val, x, H, E_J):
   ydot = val[0]
   Om_cross_x = np.cross(H.frame.parameters['Omega'].value, x)
   eff_pot = H.potential.energy(x).value[0] - 0.5*Om_cross_x.dot(Om_cross_x)
   return E_J - 0.5*ydot**2 - eff_pot

def xxdot_to_qp(x, xdot, Omega):
    q = x
    p = np.array(xdot) + np.cross(Omega, x)
    return q, p

x0 = [0.5, 0., 0.]
orbits = []
for level in E_J_levels:
    res = root(func_ydot, x0=0.3, args=(x0, H, level))
    xdot0 = [0, res.x[0], 0.]
    w0 = np.concatenate(xxdot_to_qp(x0, xdot0, Omega))
    orbit = H.integrate_orbit(w0, dt=1E-2, n_steps=100000,
                              Integrator=gi.DOPRI853Integrator)
    orbits.append(orbit)

fig,axes = plt.subplots(2, 2, figsize=(8,8), sharex=True, sharey=True)

for ax, level, orbit in zip(axes.flat, E_J_levels, orbits):
    ax.contourf(x_grid, y_grid, E_J.reshape(128,128),
                levels=[level,0], colors='#aaaaaa')
    ax.scatter(-mu, 0, c='r')
    ax.scatter(1-mu, 0, c='r')
    ax.set_title(r'$E_{{\rm J}} = {:.2f}$'.format(level))

    ax.plot(orbit.x, orbit.y, marker='None', linewidth=1.)

ax.set_xlim(-1.6, 1.6)
ax.set_ylim(-1.6, 1.6)

fig.tight_layout()