Repository: adrn/gala
Issue: gd.fit_isochrone(...) has trouble converging for certain near-circular orbits in MilkyWayPotential

Context: calculating actions in MilkyWayPotential following [procedure in gala docs](http://gala.adrian.pw/en/latest/dynamics/actionangle.html). Tested on python 2.7.9, gala 0.3, and two versions of numpy (1.14.2, 1.15.4).

If you perturb a planar circular orbit by rotating its velocity slightly out of plane (for orbit at `$\\vec{r} = (x,0,0)$`, transfer some `$v_y$` to `$v_z$` while preserving `$|v|$ = $|v_\\mt{circ}|$`), and then try to fit an isochrone using the perturbed orbit, the fit often breaks.  Failure modes include:
* `fit_isochrone(...)` throws `ValueError: Failed to fit toy potential to orbit.`
* fitted potential mass is too low compared to MW potential; toy action
  calculations fail because orbit that was bound in MW potential is unbound in
  toy potential

The issue can be side-stepped by lowering starting velocity with respect to circular, e.g. setting `$|v| = 0.8 |v_\\mt{circ}|$`.

Some observations, demonstrated by example code below.
* Fit is sensitive to numerical noise (e.g., 1e-8 km/s change to initial velocity can break fit)
* Fit can break at multiple galactocentric radii: 0.3, 3, 100 kpc
* Increasing resolution + integration time doesn't always help the fit converge
* Fitting NFWPotential works better than fitting MilkyWayPotential, but at least one example below still breaks in NFWPotential

So it seems like the fit is snagging on a weird part of `$\\chi^2$` space. It's not yet clear to me whether there's a good physical explanation, or if it's easy/possible to tune the fit for these cases.

    from __future__ import division, print_function
    import numpy as np
    import astropy.units as u
    import gala.dynamics as gd
    import gala.potential as gp
    from gala.units import galactic
    
    # configuration
    example = 2
    use_mw = False # Toggle to make some fits work
    lower_v = False # Toggle to make all fits work (and quickly too!)
    
    if use_mw:
        pot = gp.MilkyWayPotential(units=galactic)
    else:
        pot = gp.NFWPotential.from_circular_velocity(
                    v_c=220*u.km/u.s, r_s=4.24672*u.kpc, units=galactic)
    
    if example == 0:  # 3 kpc example - mass is too low, orbit unbound in toy pot
        dt = 0.1
        n_steps = 5000
        N_max = 8
        rvec = [3,0,0]*u.kpc
        vinit = pot.circular_velocity(rvec)[0].to(u.km/u.s).value
        vvec = [0, vinit*np.cos(0.01), vinit*np.sin(0.01)]*u.km/u.s  # fails
        #vvec = [0, vinit*np.cos(0.2), vinit*np.sin(0.2)]*u.km/u.s  # also fails
        #vvec = [0, vinit*np.cos(1), vinit*np.sin(1)]*u.km/u.s  # also fails
    
    elif example == 1:  # 0.3 kpc example - "Failed to fit toy potential"
        dt = 0.01
        n_steps = 50000
        N_max = 8
        rvec = [0.3,0,0]*u.kpc
        vinit = pot.circular_velocity(rvec)[0].to(u.km/u.s).value
        vvec = [0, vinit*np.cos(0.01), vinit*np.sin(0.01)]*u.km/u.s
        vvec = 0.999*vvec
    
    elif example == 2:  # 100 kpc example
        dt = 0.2
        n_steps = 500000
        N_max = 8
        rvec = [100,0,0]*u.kpc
        vinit = pot.circular_velocity(rvec)[0].to(u.km/u.s).value
        vvec = [0, vinit*np.cos(0.1), vinit*np.sin(0.1)]*u.km/u.s
    
    elif example == 3:  # high resolution, but fit still fails
        dt = 0.001
        n_steps = 500000
        N_max = 8
        rvec = [3,0,0]*u.kpc
        vinit = pot.circular_velocity(rvec)[0].to(u.km/u.s).value
        vvec = [0, vinit*np.cos(0.04), vinit*np.sin(0.04)]*u.km/u.s
        #vvec = [0, vinit*np.cos(0.2), vinit*np.sin(0.2)]*u.km/u.s  # also fails
        #vvec = [0, vinit*np.cos(1), vinit*np.sin(1)]*u.km/u.s  # also fails
    
    if lower_v:
        vvec = 0.8 * vvec
    
    print("Example#",example)
    print("Got rvec=", rvec)
    print("Got vvec=", vvec)
    
    ics = gd.PhaseSpacePosition(pos=rvec, vel=vvec)
    orbit = gp.Hamiltonian(pot).integrate_orbit(ics, dt=dt, n_steps=n_steps)    
    toy_potential = gd.fit_isochrone(orbit)
    
    if use_mw:
        m_tot = sum([pot[k].parameters['m'] for k in pot])
        print("MW potential mass:", m_tot)
    else:
        print("NFW potential:", pot.parameters)
    print("Fitted isochrone:", toy_potential.parameters)
    
    toy_actions, toy_angles, toy_freqs = toy_potential.action_angle(orbit)

Changed files:
gala/dynamics/tests/test_actionangle.py, gala/dynamics/actionangle.py

diff --git a/gala/dynamics/actionangle.py b/gala/dynamics/actionangle.py
index 41c181ff..734b320b 100644
--- a/gala/dynamics/actionangle.py
+++ b/gala/dynamics/actionangle.py
@@ -11,7 +11,7 @@ import warnings
 import numpy as np
 from astropy import log as logger
 from scipy.linalg import solve
-from scipy.optimize import leastsq
+from scipy.optimize import minimize, leastsq
 
 # Project
 from ..potential import HarmonicOscillatorPotential, IsochronePotential
@@ -60,17 +60,22 @@ def generate_n_vectors(N_max, dx=1, dy=1, dz=1, half_lattice=True):
     vecs = np.meshgrid(np.arange(-N_max, N_max+1, dx),
                        np.arange(-N_max, N_max+1, dy),
                        np.arange(-N_max, N_max+1, dz))
-    vecs = np.vstack(map(np.ravel,vecs)).T
-    vecs = vecs[np.linalg.norm(vecs,axis=1) <= N_max]
+    vecs = np.vstack(map(np.ravel, vecs)).T
+    vecs = vecs[np.linalg.norm(vecs, axis=1) <= N_max]
 
     if half_lattice:
-        ix = ((vecs[:,2] > 0) | ((vecs[:,2] == 0) & (vecs[:,1] > 0)) | ((vecs[:,2] == 0) & (vecs[:,1] == 0) & (vecs[:,0] > 0)))
+        ix = ((vecs[:, 2] > 0) |
+              ((vecs[:, 2] == 0) &
+               (vecs[:, 1] > 0)) |
+              ((vecs[:, 2] == 0) &
+               (vecs[:, 1] == 0) &
+               (vecs[:, 0] > 0)))
         vecs = vecs[ix]
 
-    vecs = np.array(sorted(vecs, key=lambda x: (x[0],x[1],x[2])))
+    vecs = np.array(sorted(vecs, key=lambda x: (x[0], x[1], x[2])))
     return vecs
 
-def fit_isochrone(orbit, m0=2E11, b0=1.):
+def fit_isochrone(orbit, m0=2E11, b0=1., minimize_kwargs=None):
     r"""
     Fit the toy Isochrone potential to the sum of the energy residuals relative
     to the mean energy by minimizing the function
@@ -89,6 +94,8 @@ def fit_isochrone(orbit, m0=2E11, b0=1.):
         Initial mass guess.
     b0 : numeric (optional)
         Initial b guess.
+    minimize_kwargs : dict (optional)
+        Keyword arguments to pass through to `scipy.optimize.minimize`.
 
     Returns
     -------
@@ -106,24 +113,33 @@ def fit_isochrone(orbit, m0=2E11, b0=1.):
     if w.ndim > 2:
         raise ValueError("Input orbit object must be a single orbit.")
 
-    def f(p,w):
-        logm,b = p
-        potential = IsochronePotential(m=np.exp(logm), b=b, units=pot.units)
-        H = potential.value(w[:3]).decompose(pot.units).value + 0.5*np.sum(w[3:]**2, axis=0)
-        return np.squeeze(H - np.mean(H))
+    def f(p, w):
+        logm, logb = p
+        potential = IsochronePotential(m=np.exp(logm), b=np.exp(logb),
+                                       units=pot.units)
+        H = (potential.value(w[:3]).decompose(pot.units).value +
+             0.5*np.sum(w[3:]**2, axis=0))
+        return np.sum(np.squeeze(H - np.mean(H))**2)
 
     logm0 = np.log(m0)
-    p,ier = leastsq(f, np.array([logm0, b0]), args=(w,))
+    logb0 = np.log(b0)
 
-    if ier < 1 or ier > 4:
+    if minimize_kwargs is None:
+        minimize_kwargs = dict()
+    minimize_kwargs['x0'] = np.array([logm0, logb0])
+    minimize_kwargs['method'] = minimize_kwargs.get('method', 'Nelder-Mead')
+    res = minimize(f, args=(w,), **minimize_kwargs)
+
+    if not res.success:
         raise ValueError("Failed to fit toy potential to orbit.")
 
-    logm,b = np.abs(p)
+    logm, logb = np.abs(res.x)
     m = np.exp(logm)
+    b = np.exp(logb)
 
     return IsochronePotential(m=m, b=b, units=pot.units)
 
-def fit_harmonic_oscillator(orbit, omega0=[1.,1.,1.]):
+def fit_harmonic_oscillator(orbit, omega0=[1., 1, 1], minimize_kwargs=None):
     r"""
     Fit the toy harmonic oscillator potential to the sum of the energy
     residuals relative to the mean energy by minimizing the function
@@ -140,6 +156,8 @@ def fit_harmonic_oscillator(orbit, omega0=[1.,1.,1.]):
     orbit : `~gala.dynamics.Orbit`
     omega0 : array_like (optional)
         Initial frequency guess.
+    minimize_kwargs : dict (optional)
+        Keyword arguments to pass through to `scipy.optimize.minimize`.
 
     Returns
     -------
@@ -159,14 +177,20 @@ def fit_harmonic_oscillator(orbit, omega0=[1.,1.,1.]):
 
     def f(omega, w):
         potential = HarmonicOscillatorPotential(omega=omega, units=pot.units)
-        H = potential.value(w[:3]).decompose(pot.units).value + 0.5*np.sum(w[3:]**2, axis=0)
-        return np.squeeze(H - np.mean(H))
+        H = (potential.value(w[:3]).decompose(pot.units).value +
+             0.5*np.sum(w[3:]**2, axis=0))
+        return np.sum(np.squeeze(H - np.mean(H))**2)
+
+    if minimize_kwargs is None:
+        minimize_kwargs = dict()
+    minimize_kwargs['x0'] = omega0
+    minimize_kwargs['method'] = minimize_kwargs.get('method', 'Nelder-Mead')
+    res = minimize(f, args=(w,), **minimize_kwargs)
 
-    p,ier = leastsq(f, omega0, args=(w,))
-    if ier < 1 or ier > 4:
+    if not res.success:
         raise ValueError("Failed to fit toy potential to orbit.")
 
-    best_omega = np.abs(p)
+    best_omega = np.abs(res.x)
     return HarmonicOscillatorPotential(omega=best_omega, units=pot.units)
 
 def fit_toy_potential(orbit, force_harmonic_oscillator=False):
@@ -205,7 +229,8 @@ def fit_toy_potential(orbit, force_harmonic_oscillator=False):
         logger.debug("Using triaxial harmonic oscillator toy potential")
 
         toy_potential = fit_harmonic_oscillator(orbit)
-        logger.debug("Best omegas ({})".format(toy_potential.parameters['omega']))
+        logger.debug("Best omegas ({})"
+                     .format(toy_potential.parameters['omega']))
 
     return toy_potential
 
@@ -235,10 +260,10 @@ def check_angle_sampling(nvecs, angles):
     failed_nvecs = []
     failures = []
 
-    for i,vec in enumerate(nvecs):
+    for i, vec in enumerate(nvecs):
         # N = np.linalg.norm(vec)
         # X = np.dot(angles,vec)
-        X = (angles*vec[:,None]).sum(axis=0)
+        X = (angles*vec[:, None]).sum(axis=0)
         diff = float(np.abs(X.max() - X.min()))
 
         if diff < (2.*np.pi):
@@ -293,7 +318,7 @@ def _action_prepare(aa, N_max, dx, dy, dz, sign=1., throw_out_modes=False):
     nvecs = generate_n_vectors(N_max, dx, dy, dz)
 
     # make sure we have enough angle coverage
-    modes,P = check_angle_sampling(nvecs, angles)
+    modes, P = check_angle_sampling(nvecs, angles)
 
     # throw out modes?
     # if throw_out_modes:
@@ -301,30 +326,31 @@ def _action_prepare(aa, N_max, dx, dy, dz, sign=1., throw_out_modes=False):
 
     n = len(nvecs) + 3
     b = np.zeros(shape=(n, ))
-    A = np.zeros(shape=(n,n))
+    A = np.zeros(shape=(n, n))
 
     # top left block matrix: identity matrix summed over timesteps
-    A[:3,:3] = aa.shape[1]*np.identity(3)
+    A[:3, :3] = aa.shape[1]*np.identity(3)
 
     actions = aa[:3]
     angles = aa[3:]
 
     # top right block matrix: transpose of C_nk matrix (Eq. 12)
-    C_T = 2.*nvecs.T * np.sum(np.cos(np.dot(nvecs,angles)), axis=-1)
+    C_T = 2.*nvecs.T * np.sum(np.cos(np.dot(nvecs, angles)), axis=-1)
     A[:3,3:] = C_T
-    A[3:,:3] = C_T.T
+    A[3:, :3] = C_T.T
 
     # lower right block matrix: C_nk dotted with C_nk^T
-    cosv = np.cos(np.dot(nvecs,angles))
-    A[3:,3:] = 4.*np.dot(nvecs,nvecs.T)*np.einsum('it,jt->ij', cosv, cosv)
+    cosv = np.cos(np.dot(nvecs, angles))
+    A[3:,3:] = 4.*np.dot(nvecs, nvecs.T)*np.einsum('it,jt->ij', cosv, cosv)
 
     # b vector first three is just sum of toy actions
     b[:3] = np.sum(actions, axis=1)
 
     # rest of the vector is C dotted with actions
-    b[3:] = 2*np.sum(np.dot(nvecs,actions)*np.cos(np.dot(nvecs,angles)), axis=1)
+    b[3:] = 2*np.sum(np.dot(nvecs, actions)*np.cos(np.dot(nvecs, angles)),
+                     axis=1)
 
-    return A,b,nvecs
+    return A, b, nvecs
 
 def _angle_prepare(aa, t, N_max, dx, dy, dz, sign=1.):
     """
@@ -365,7 +391,7 @@ def _angle_prepare(aa, t, N_max, dx, dy, dz, sign=1.):
     nvecs = generate_n_vectors(N_max, dx, dy, dz)
 
     # make sure we have enough angle coverage
-    modes,P = check_angle_sampling(nvecs, angles)
+    modes, P = check_angle_sampling(nvecs, angles)
 
     # TODO: throw out modes?
     # if(throw_out_modes):
@@ -375,41 +401,44 @@ def _angle_prepare(aa, t, N_max, dx, dy, dz, sign=1.):
     n = 3 + 3 + 3*nv # angle(0)'s, freqs, 3 derivatives of Sn
 
     b = np.zeros(shape=(n,))
-    A = np.zeros(shape=(n,n))
+    A = np.zeros(shape=(n, n))
 
     # top left block matrix: identity matrix summed over timesteps
-    A[:3,:3] = aa.shape[1]*np.identity(3)
+    A[:3, :3] = aa.shape[1]*np.identity(3)
 
     # identity matrices summed over times
-    A[:3,3:6] = A[3:6,:3] = np.sum(t)*np.identity(3)
-    A[3:6,3:6] = np.sum(t*t)*np.identity(3)
+    A[:3, 3:6] = A[3:6, :3] = np.sum(t)*np.identity(3)
+    A[3:6, 3:6] = np.sum(t*t)*np.identity(3)
 
     # S1,2,3
-    A[6:6+nv,0] = -2.*np.sum(np.sin(np.dot(nvecs,angles)),axis=1)
-    A[6+nv:6+2*nv,1] = A[6:6+nv,0]
-    A[6+2*nv:6+3*nv,2] = A[6:6+nv,0]
+    A[6:6+nv, 0] = -2.*np.sum(np.sin(np.dot(nvecs, angles)), axis=1)
+    A[6+nv:6+2*nv, 1] = A[6:6+nv, 0]
+    A[6+2*nv:6+3*nv, 2] = A[6:6+nv, 0]
 
     # t*S1,2,3
-    A[6:6+nv,3] = -2.*np.sum(t[None,:]*np.sin(np.dot(nvecs,angles)),axis=1)
-    A[6+nv:6+2*nv,4] = A[6:6+nv,3]
-    A[6+2*nv:6+3*nv,5] = A[6:6+nv,3]
+    A[6:6+nv, 3] = -2.*np.sum(t[None, :]*np.sin(np.dot(nvecs, angles)),
+                              axis=1)
+    A[6+nv:6+2*nv, 4] = A[6:6+nv, 3]
+    A[6+2*nv:6+3*nv, 5] = A[6:6+nv, 3]
 
     # lower right block structure: S dot S^T
-    sinv = np.sin(np.dot(nvecs,angles))
+    sinv = np.sin(np.dot(nvecs, angles))
     SdotST = np.einsum('it,jt->ij', sinv, sinv)
-    A[6:6+nv,6:6+nv] = A[6+nv:6+2*nv,6+nv:6+2*nv] = \\
-        A[6+2*nv:6+3*nv,6+2*nv:6+3*nv] = 4*SdotST
+    A[6:6+nv, 6:6+nv] = A[6+nv:6+2*nv, 6+nv:6+2*nv] = \\
+        A[6+2*nv:6+3*nv, 6+2*nv:6+3*nv] = 4*SdotST
 
     # top rectangle
-    A[:6,:] = A[:,:6].T
+    A[:6, :] = A[:, :6].T
 
     b[:3] = np.sum(angles.T, axis=0)
-    b[3:6] = np.sum(t[:,None]*angles.T, axis=0)
-    b[6:6+nv] = -2.*np.sum(angles[0]*np.sin(np.dot(nvecs,angles)), axis=1)
-    b[6+nv:6+2*nv] = -2.*np.sum(angles[1]*np.sin(np.dot(nvecs,angles)), axis=1)
-    b[6+2*nv:6+3*nv] = -2.*np.sum(angles[2]*np.sin(np.dot(nvecs,angles)), axis=1)
+    b[3:6] = np.sum(t[:, None]*angles.T, axis=0)
+    b[6:6+nv] = -2.*np.sum(angles[0]*np.sin(np.dot(nvecs, angles)), axis=1)
+    b[6+nv:6+2*nv] = -2.*np.sum(angles[1]*np.sin(np.dot(nvecs, angles)),
+                                axis=1)
+    b[6+2*nv:6+3*nv] = -2.*np.sum(angles[2]*np.sin(np.dot(nvecs, angles)),
+                                  axis=1)
 
-    return A,b,nvecs
+    return A, b, nvecs
 
 def _single_orbit_find_actions(orbit, N_max, toy_potential=None,
                                force_harmonic_oscillator=False):
@@ -440,21 +469,23 @@ def _single_orbit_find_actions(orbit, N_max, toy_potential=None,
         raise ValueError("must be a single orbit")
 
     if toy_potential is None:
-        toy_potential = fit_toy_potential(orbit, force_harmonic_oscillator=force_harmonic_oscillator)
+        toy_potential = fit_toy_potential(
+            orbit, force_harmonic_oscillator=force_harmonic_oscillator)
 
     else:
-        logger.debug("Using *fixed* toy potential: {}".format(toy_potential.parameters))
+        logger.debug("Using *fixed* toy potential: {}"
+                     .format(toy_potential.parameters))
 
     if isinstance(toy_potential, IsochronePotential):
         orbit_align = orbit.align_circulation_with_z()
         w = orbit_align.w()
 
-        dxyz = (1,2,2)
-        circ = np.sign(w[0,0]*w[4,0]-w[1,0]*w[3,0])
-        sign = np.array([1.,circ,1.])
+        dxyz = (1, 2, 2)
+        circ = np.sign(w[0, 0]*w[4, 0]-w[1, 0]*w[3, 0])
+        sign = np.array([1., circ, 1.])
         orbit = orbit_align
     elif isinstance(toy_potential, HarmonicOscillatorPotential):
-        dxyz = (2,2,2)
+        dxyz = (2, 2, 2)
         sign = 1.
         w = orbit.w()
     else:
@@ -466,31 +497,32 @@ def _single_orbit_find_actions(orbit, N_max, toy_potential=None,
     aaf = toy_potential.action_angle(orbit)
 
     if aaf[0].ndim > 2:
-        aa = np.vstack((aaf[0].value[...,0], aaf[1].value[...,0]))
+        aa = np.vstack((aaf[0].value[..., 0], aaf[1].value[..., 0]))
     else:
         aa = np.vstack((aaf[0].value, aaf[1].value))
 
     if np.any(np.isnan(aa)):
-        ix = ~np.any(np.isnan(aa),axis=0)
-        aa = aa[:,ix]
+        ix = ~np.any(np.isnan(aa), axis=0)
+        aa = aa[:, ix]
         t = t[ix]
         warnings.warn("NaN value in toy actions or angles!")
         if sum(ix) > 1:
             raise ValueError("Too many NaN value in toy actions or angles!")
 
     t1 = time.time()
-    A,b,nvecs = _action_prepare(aa, N_max, dx=dxyz[0], dy=dxyz[1], dz=dxyz[2])
+    A, b, nvecs = _action_prepare(aa, N_max, dx=dxyz[0], dy=dxyz[1], dz=dxyz[2])
     actions = np.array(solve(A,b))
     logger.debug("Action solution found for N_max={}, size {} symmetric"
                  " matrix in {} seconds"
-                 .format(N_max,len(actions),time.time()-t1))
+                 .format(N_max, len(actions), time.time()-t1))
 
     t1 = time.time()
-    A,b,nvecs = _angle_prepare(aa, t, N_max, dx=dxyz[0], dy=dxyz[1], dz=dxyz[2], sign=sign)
-    angles = np.array(solve(A,b))
+    A, b, nvecs = _angle_prepare(aa, t, N_max, dx=dxyz[0],
+                                 dy=dxyz[1], dz=dxyz[2], sign=sign)
+    angles = np.array(solve(A, b))
     logger.debug("Angle solution found for N_max={}, size {} symmetric"
                  " matrix in {} seconds"
-                 .format(N_max,len(angles),time.time()-t1))
+                 .format(N_max, len(angles), time.time()-t1))
 
     # Just some checks
     if len(angles) > len(aa):
@@ -500,7 +532,8 @@ 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,
+    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, toy_potential=None):
@@ -533,24 +566,27 @@ def find_actions(orbit, N_max, force_harmonic_oscillator=False, toy_potential=No
     """
 
     if orbit.norbits == 1:
-        return _single_orbit_find_actions(orbit, N_max,
-                                          force_harmonic_oscillator=force_harmonic_oscillator,
-                                          toy_potential=toy_potential)
+        return _single_orbit_find_actions(
+            orbit, N_max,
+            force_harmonic_oscillator=force_harmonic_oscillator,
+            toy_potential=toy_potential)
 
     else:
         norbits = orbit.norbits
-        actions = np.zeros((3,norbits))
-        angles = np.zeros((3,norbits))
-        freqs = np.zeros((3,norbits))
+        actions = np.zeros((3, norbits))
+        angles = np.zeros((3, norbits))
+        freqs = np.zeros((3, norbits))
         for n in range(norbits):
-            aaf = _single_orbit_find_actions(orbit[:,n], N_max,
-                                             force_harmonic_oscillator=force_harmonic_oscillator,
-                                             toy_potential=toy_potential)
+            aaf = _single_orbit_find_actions(
+                orbit[:, n], N_max,
+                force_harmonic_oscillator=force_harmonic_oscillator,
+                toy_potential=toy_potential)
             actions[n] = aaf['actions'].value
             angles[n] = aaf['angles'].value
             freqs[n] = aaf['freqs'].value
 
-    return dict(actions=actions*aaf['actions'].unit, angles=angles*aaf['angles'].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'])
 
diff --git a/gala/dynamics/tests/test_actionangle.py b/gala/dynamics/tests/test_actionangle.py
index 35b205a8..02c0dd14 100644
--- a/gala/dynamics/tests/test_actionangle.py
+++ b/gala/dynamics/tests/test_actionangle.py
@@ -40,7 +40,8 @@ def test_fit_isochrone():
     orbit = potential.integrate_orbit([15.,0,0,0,0.2,0], dt=2., n_steps=10000)
 
     fit_potential = fit_isochrone(orbit)
-    m,b = fit_potential.parameters['m'].value, fit_potential.parameters['b'].value
+    m, b = (fit_potential.parameters['m'].value,
+            fit_potential.parameters['b'].value)
     assert np.allclose(m, true_m, rtol=1E-2)
     assert np.allclose(b, true_b, rtol=1E-2)
 
@@ -232,3 +233,39 @@ def test_compare_action_prepare():
     # assert np.allclose(ang_apw, ang_san)
 
     # TODO: this could be critical -- why don't our angles agree?
+
+
+def test_regression_113():
+    """Test that fit_isochrone succeeds for a variety of orbits. See issue:
+    https://github.com/adrn/gala/issues/113
+    """
+    from ...potential import MilkyWayPotential, Hamiltonian
+    from ...dynamics import PhaseSpacePosition
+
+    pot = MilkyWayPotential()
+
+    dt = 0.01
+    n_steps = 50000
+
+    rvec = [0.3, 0, 0]*u.kpc
+    vinit = pot.circular_velocity(rvec)[0].to(u.km/u.s).value
+    vvec = [0, vinit*np.cos(0.01), vinit*np.sin(0.01)]*u.km/u.s
+    vvec = 0.999*vvec
+
+    ics = PhaseSpacePosition(pos=rvec, vel=vvec)
+    orbit = Hamiltonian(pot).integrate_orbit(ics, dt=dt, n_steps=n_steps)
+    toy_potential = fit_isochrone(orbit)
+
+    m = toy_potential.parameters['m'].value
+    b = toy_potential.parameters['b'].value
+    assert np.log10(m) > 11 and np.log10(m) < 12
+    assert np.log10(b) > 0 and np.log10(b) < 1
+
+    # try again with Nelder-Mead
+    toy_potential = fit_isochrone(orbit,
+                                  minimize_kwargs=dict(method='Nelder-Mead'))
+
+    m = toy_potential.parameters['m'].value
+    b = toy_potential.parameters['b'].value
+    assert np.log10(m) > 11 and np.log10(m) < 12
+    assert np.log10(b) > 0 and np.log10(b) < 1