from gala.dynamics.actionangle import get_staeckel_fudge_delta
from galpy.actionAngle import actionAngleStaeckel

galpy_potential = pot.to_galpy_potential()
J = np.zeros((3, orbits.norbits))
Omega = np.zeros((3, orbits.norbits))
for n, orbit in enumerate(orbits.orbit_gen()):
    o = orbit.to_galpy_orbit()
    delta = get_staeckel_fudge_delta(pot, orbit)
    staeckel = actionAngleStaeckel(pot=galpy_potential, delta=delta)
    af = staeckel.actionsFreqs(o)
    af = np.mean(np.stack(af), axis=1)

    J[:3, n] = af[:3]
    Omega[:3, n] = af[3:]

fig, ax = plt.subplots(figsize=(6, 6), constrained_layout=True)
ax.plot(w0.v_z, J[2])
ax.set_xlabel(f"$v_z$ [{w0.v_z.unit:latex_inline}]")
ax.set_ylabel(rf"$J_z$")