Merge pull request #232 from adrn/fix-aa-freq-unit

Fix action-angle frequency unit

CHANGES.rst

@@ -27,6 +27,8 @@ Bug fixes
 - Implemented a density function for ``LogarithmicPotential``, which was
   missing previously.
 
+- The analytic action-angle and ``find_actions()`` utilities now correctly
+  return frequencies with angular frequency units rather than frequency.
 
 API changes
 -----------

gala/dynamics/actionangle.py

@@ -541,9 +541,12 @@ def _single_orbit_find_actions(orbit, N_max, toy_potential=None,
     theta = angles[:3]
     freqs = angles[3:6]  # * sign
 
-    return dict(actions=J*aaf[0].unit, angles=theta*aaf[1].unit,
-                freqs=freqs*aaf[2].unit,
-                Sn=actions[3:], dSn_dJ=angles[6:], nvecs=nvecs)
+    return dict(actions=J * aaf[0].unit,
+                angles=theta * aaf[1].unit,
+                freqs=freqs * aaf[2].unit,
+                Sn=actions[3:],
+                dSn_dJ=angles[6:],
+                nvecs=nvecs)
 
 
 def find_actions(orbit, N_max, force_harmonic_oscillator=False,
@@ -596,9 +599,9 @@ def find_actions(orbit, N_max, force_harmonic_oscillator=False,
             angles[n] = aaf['angles'].value
             freqs[n] = aaf['freqs'].value
 
-    return dict(actions=actions*aaf['actions'].unit,
-                angles=angles*aaf['angles'].unit,
-                freqs=freqs*aaf['freqs'].unit,
+    return dict(actions=actions * aaf['actions'].unit,
+                angles=angles * aaf['angles'].unit,
+                freqs=freqs * aaf['freqs'].unit,
                 Sn=actions[3:], dSn=angles[6:], nvecs=aaf['nvecs'])
 
 # def solve_hessian(relative_actions, relative_freqs):

gala/dynamics/actionangle_staeckel.py

@@ -52,6 +52,6 @@ def get_staeckel_fudge_delta(potential, w, median=True):
 
     # Median over time if the inputs were orbits
     if len(delta.shape) > 1 and median:
-        delta = np.median(delta, axis=1)
+        delta = np.nanmedian(delta, axis=1)
 
     return delta * potential.units['length']

gala/dynamics/analyticactionangle.py

@@ -148,8 +148,8 @@ def F(x, y):
     freqs[1] = np.sign(actions[1]) * omega_th * omega_r
     freqs[2] = omega_th * omega_r
 
-    a_unit = (1*usys['angular momentum']/usys['mass']).decompose(usys).unit
-    f_unit = (1*usys['frequency']).decompose(usys).unit
+    a_unit = (1 * usys['angular momentum'] / usys['mass']).decompose(usys).unit
+    f_unit = (1 * usys['angular speed']).decompose(usys).unit
     return actions*a_unit, angles*u.radian, freqs*f_unit
 
 
@@ -349,7 +349,7 @@ def harmonic_oscillator_to_aa(w, potential):
 
     if usys is not None and usys:
         a_unit = (1*usys['angular momentum']/usys['mass']).decompose(usys).unit
-        f_unit = (1*usys['frequency']).decompose(usys).unit
+        f_unit = (1*usys['angular speed']).decompose(usys).unit
         return action*a_unit, (angle % (2.*np.pi))*u.radian, freq*f_unit
     else:
         return action*u.one, (angle % (2.*np.pi))*u.one, freq*u.one