import astropy.units as u
import numpy as np
import gala.dynamics as gd
import gala.potential as gp
from gala.units import galactic

pot = gp.HernquistPotential(m=1E10*u.Msun, c=1.*u.kpc,
                            units=galactic)
w0 = gd.PhaseSpacePosition(pos=[5.,0,0]*u.kpc,
                           vel=[0,0,50.]*u.km/u.s)

rotation_axis = np.array([8.2, -1.44, 3.25])
rotation_axis /= np.linalg.norm(rotation_axis) # make a unit vector
frame_freq = 42. * u.km/u.s/u.kpc
rot_frame = gp.ConstantRotatingFrame(Omega=frame_freq * rotation_axis,
                                     units=galactic)

H_rot = gp.Hamiltonian(potential=pot, frame=rot_frame)
rot_orbit = H_rot.integrate_orbit(w0, dt=0.5, n_steps=1000)
_ = rot_orbit.plot(marker='') # doctest: +SKIP