fixes #231

adrn · Jun 21, 2021 · 79c5126 · 79c5126
1 parent d1b2e1e
commit 79c5126
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 9 deletions.
diff --git a/gala/dynamics/actionangle.py b/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):

diff --git a/gala/dynamics/analyticactionangle.py b/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