Perform Python leapfrog integration in cylindrical coordinates in all…

… cases; also combines subsequent drifts, thus fixes #365

galpy/orbit/FullOrbit.py

@@ -589,28 +589,17 @@ def _integrateFullOrbit(vxvv,pot,t,method,dt):
             method= 'odeint'
         warnings.warn("Cannot use symplectic integration because some of the included forces are dissipative (using non-symplectic integrator %s instead)" % (method), galpyWarning)
     if method.lower() == 'leapfrog':
-        #go to the rectangular frame
-        this_vxvv= nu.array([vxvv[0]*nu.cos(vxvv[5]),
-                             vxvv[0]*nu.sin(vxvv[5]),
-                             vxvv[3],
-                             vxvv[1]*nu.cos(vxvv[5])-vxvv[2]*nu.sin(vxvv[5]),
-                             vxvv[2]*nu.cos(vxvv[5])+vxvv[1]*nu.sin(vxvv[5]),
-                             vxvv[4]])
         #integrate
-        out= symplecticode.leapfrog(_rectForce,this_vxvv,
-                                    t,args=(pot,),rtol=10.**-8)
-        #go back to the cylindrical frame
-        R= nu.sqrt(out[:,0]**2.+out[:,1]**2.)
-        phi= nu.arccos(out[:,0]/R)
-        phi[(out[:,1] < 0.)]= 2.*nu.pi-phi[(out[:,1] < 0.)]
-        vR= out[:,3]*nu.cos(phi)+out[:,4]*nu.sin(phi)
-        vT= out[:,4]*nu.cos(phi)-out[:,3]*nu.sin(phi)
-        out[:,3]= out[:,2]
-        out[:,4]= out[:,5]
-        out[:,0]= R
-        out[:,1]= vR
-        out[:,2]= vT
-        out[:,5]= phi
+        out= symplecticode.leapfrog_cyl(\
+            lambda x: nu.array([x[1],0.,0.,x[4],0.,0.]),
+            lambda x,t: nu.array([0.,
+                                  _evaluateRforces(pot,x[0],x[3],phi=x[5],t=t)
+                                  +x[2]**2./x[0]**3.,
+                                  _evaluatephiforces(pot,x[0],x[3],phi=x[5],t=t),
+                                  0.,
+                                  _evaluatezforces(pot,x[0],x[3],phi=x[5],t=t),
+                                  x[2]/x[0]**2]),
+            vxvv,t,rtol=10.**-8)
     elif ext_loaded and \
             (method.lower() == 'leapfrog_c' or method.lower() == 'rk4_c' \
             or method.lower() == 'rk6_c' or method.lower() == 'symplec4_c' \
@@ -690,34 +679,6 @@ def _FullEOM(y,t,pot):
             _evaluatezforces(pot,y[0],y[4],phi=y[2],t=t,
                              v=[y[1],y[0]*y[3],y[5]])]
 
-def _rectForce(x,pot,t=0.):
-    """
-    NAME:
-       _rectForce
-    PURPOSE:
-       returns the force in the rectangular frame
-    INPUT:
-       x - current position
-       t - current time
-       pot - (list of) Potential instance(s)
-    OUTPUT:
-       force
-    HISTORY:
-       2011-02-02 - Written - Bovy (NYU)
-    """
-    #x is rectangular so calculate R and phi
-    R= nu.sqrt(x[0]**2.+x[1]**2.)
-    phi= nu.arccos(x[0]/R)
-    sinphi= x[1]/R
-    cosphi= x[0]/R
-    if x[1] < 0.: phi= 2.*nu.pi-phi
-    #calculate forces
-    Rforce= _evaluateRforces(pot,R,x[2],phi=phi,t=t)
-    phiforce= _evaluatephiforces(pot,R,x[2],phi=phi,t=t)
-    return nu.array([cosphi*Rforce-1./R*sinphi*phiforce,
-                     sinphi*Rforce+1./R*cosphi*phiforce,
-                     _evaluatezforces(pot,R,x[2],phi=phi,t=t)])
-
 def _fit_orbit(orb,vxvv,vxvv_err,pot,radec=False,lb=False,
                customsky=False,lb_to_customsky=None,
                pmllpmbb_to_customsky=None,

galpy/orbit/RZOrbit.py

@@ -11,6 +11,7 @@
 from .integrateFullOrbit import _ext_loaded as ext_loaded
 from galpy.util.bovy_conversion import physical_conversion
 from galpy.util.leung_dop853 import dop853
+import galpy.util.bovy_symplecticode as symplecticode
 from .OrbitTop import OrbitTop
 class RZOrbit(OrbitTop):
     """Class that holds and integrates orbits in axisymetric potentials 
@@ -481,8 +482,21 @@ def _integrateRZOrbit(vxvv,pot,t,method,dt):
                 warnings.warn("Cannot use C integration because C extension not loaded (using %s instead)" % (method), galpyWarning)
             else:
                 warnings.warn("Cannot use C integration because some of the potentials are not implemented in C (using %s instead)" % (method), galpyWarning)
-    if method.lower() == 'leapfrog' \
-            or method.lower() == 'leapfrog_c' or method.lower() == 'rk4_c' \
+    if method.lower() == 'leapfrog':
+        #integrate
+        Lz2= (vxvv[0]*vxvv[2])**2.
+        tmp_out= symplecticode.leapfrog_cyl(\
+            lambda x: nu.array([x[1],0.,x[3],0.]),
+            lambda x,t: nu.array([0.,
+                                  _evaluateRforces(pot,x[0],x[2],t=t)
+                                  +Lz2/x[0]**3.,
+                                  0.,
+                                  _evaluatezforces(pot,x[0],x[2],t=t)]),
+            nu.array(vxvv)[[0,1,3,4]],t,rtol=10.**-8,meridional=Lz2)
+        out= nu.empty((len(tmp_out),5))
+        out[:,[0,1,3,4]]= tmp_out
+        out[:,2]= vxvv[0]*vxvv[2]/tmp_out[:,0]
+    elif method.lower() == 'leapfrog_c' or method.lower() == 'rk4_c' \
             or method.lower() == 'rk6_c' or method.lower() == 'symplec4_c' \
             or method.lower() == 'symplec6_c' or method.lower() == 'dopr54_c' or method.lower() == 'dop853_c':
         #We hack this by upgrading to a FullOrbit

galpy/orbit/linearOrbit.py

@@ -159,10 +159,10 @@ def _integrateLinearOrbit(vxvv,pot,t,method,dt):
             else:
                 warnings.warn("Cannot use C integration because some of the potentials are not implemented in C (using %s instead)" % (method), galpyWarning)
     if method.lower() == 'leapfrog':
-        return symplecticode.leapfrog(lambda x,t=t: _evaluatelinearForces(pot,x,
-                                                                         t=t),
-                                      nu.array(vxvv),
-                                      t,rtol=10.**-8)
+        return symplecticode.leapfrog_cyl(\
+            lambda x: nu.array([x[1],0.]),
+            lambda x,t=t: nu.array([0.,_evaluatelinearForces(pot,x[0],t=t)]),
+            vxvv,t,rtol=10.**-8)
     elif method.lower() == 'dop853':
         return dop853(func=_linearEOM, x=vxvv, t=t, args=(pot,))
     elif ext_loaded and \

galpy/orbit/planarOrbit.py

@@ -508,12 +508,17 @@ def _integrateROrbit(vxvv,pot,t,method,dt):
             else:
                 warnings.warn("Cannot use C integration because some of the potentials are not implemented in C (using %s instead)" % (method), galpyWarning)
     if method.lower() == 'leapfrog':
-        #We hack this by putting in a dummy phi
-        this_vxvv= nu.zeros(len(vxvv)+1)
-        this_vxvv[0:len(vxvv)]= vxvv
-        tmp_out, msg= _integrateOrbit(this_vxvv,pot,t,method,dt)
-        #tmp_out is (nt,4)
-        out= tmp_out[:,0:3]
+        Lz2= (vxvv[0]*vxvv[2])**2.
+        tmp_out= symplecticode.leapfrog_cyl(\
+            lambda x: nu.array([x[1],0.]),
+            lambda x,t: nu.array([0.,
+                                  _evaluateplanarRforces(pot,x[0],t=t)
+                                  +Lz2/x[0]**3.]),
+            nu.array(vxvv)[:2],t,rtol=10.**-8,meridional=Lz2)
+        out= nu.empty((len(tmp_out),3))
+        out[:,:2]= tmp_out
+        out[:,2]= vxvv[0]*vxvv[2]/tmp_out[:,0]
+        msg= 0
     elif ext_loaded and \
             (method.lower() == 'leapfrog_c' or method.lower() == 'rk4_c' \
             or method.lower() == 'rk6_c' or method.lower() == 'symplec4_c' \
@@ -594,25 +599,15 @@ def _integrateOrbit(vxvv,pot,t,method,dt):
             else:
                 warnings.warn("Cannot use C integration because some of the potentials are not implemented in C (using %s instead)" % (method), galpyWarning)
     if method.lower() == 'leapfrog':
-        #go to the rectangular frame
-        this_vxvv= nu.array([vxvv[0]*nu.cos(vxvv[3]),
-                             vxvv[0]*nu.sin(vxvv[3]),
-                             vxvv[1]*nu.cos(vxvv[3])-vxvv[2]*nu.sin(vxvv[3]),
-                             vxvv[2]*nu.cos(vxvv[3])+vxvv[1]*nu.sin(vxvv[3])])
         #integrate
-        tmp_out= symplecticode.leapfrog(_rectForce,this_vxvv,
-                                        t,args=(pot,),rtol=10.**-8)
-        #go back to the cylindrical frame
-        R= nu.sqrt(tmp_out[:,0]**2.+tmp_out[:,1]**2.)
-        phi= nu.arccos(tmp_out[:,0]/R)
-        phi[(tmp_out[:,1] < 0.)]= 2.*nu.pi-phi[(tmp_out[:,1] < 0.)]
-        vR= tmp_out[:,2]*nu.cos(phi)+tmp_out[:,3]*nu.sin(phi)
-        vT= tmp_out[:,3]*nu.cos(phi)-tmp_out[:,2]*nu.sin(phi)
-        out= nu.zeros((len(t),4))
-        out[:,0]= R
-        out[:,1]= vR
-        out[:,2]= vT
-        out[:,3]= phi
+        out= symplecticode.leapfrog_cyl(\
+            lambda x: nu.array([x[1],0.,0.,0.]),
+            lambda x,t: nu.array([0.,
+                                  _evaluateplanarRforces(pot,x[0],phi=x[3],t=t)
+                                  +x[2]**2./x[0]**3.,
+                                  _evaluateplanarphiforces(pot,x[0],phi=x[3],t=t),
+                                  x[2]/x[0]**2]),
+            vxvv,t,rtol=10.**-8)
         msg= 0
     elif ext_loaded and \
             (method.lower() == 'leapfrog_c' or method.lower() == 'rk4_c' \
@@ -840,33 +835,6 @@ def _EOM(y,t,pot):
             1./y[0]**2.*(_evaluateplanarphiforces(pot,y[0],phi=y[2],t=t)-
                          2.*y[0]*y[1]*y[3])]
 
-def _rectForce(x,pot,t=0.):
-    """
-    NAME:
-       _rectForce
-    PURPOSE:
-       returns the force in the rectangular frame
-    INPUT:
-       x - current position
-       t - current time
-       pot - (list of) Potential instance(s)
-    OUTPUT:
-       force
-    HISTORY:
-       2011-02-02 - Written - Bovy (NYU)
-    """
-    #x is rectangular so calculate R and phi
-    R= nu.sqrt(x[0]**2.+x[1]**2.)
-    phi= nu.arccos(x[0]/R)
-    sinphi= x[1]/R
-    cosphi= x[0]/R
-    if x[1] < 0.: phi= 2.*nu.pi-phi
-    #calculate forces
-    Rforce= _evaluateplanarRforces(pot,R,phi=phi,t=t)
-    phiforce= _evaluateplanarphiforces(pot,R,phi=phi,t=t)
-    return nu.array([cosphi*Rforce-1./R*sinphi*phiforce,
-                     sinphi*Rforce+1./R*cosphi*phiforce])
-
 def _parse_warnmessage(msg):
     if msg == 1: #pragma: no cover
         warnings.warn("During numerical integration, steps smaller than the smallest step were requested; integration might not be accurate",galpyWarning)

galpy/util/bovy_symplecticode.py

@@ -32,7 +32,7 @@
 #############################################################################
 import numpy as nu
 _MAX_DT_REDUCE= 10000.
-def leapfrog(func,yo,t,args=(),rtol=1.49012e-12,atol=1.49012e-12):
+def _leapfrog(func,yo,t,args=(),rtol=1.49012e-12,atol=1.49012e-12):
     """
     NAME:
        leapfrog
@@ -112,3 +112,103 @@ def _leapfrog_estimate_step(func,qo,po,dt,to,args,rtol,atol):
         dt/= 2.
     return dt
 
+def leapfrog_cyl(drift,kick,yo,t,args=(),rtol=1.49012e-12,atol=1.49012e-12,
+                 meridional=False):
+    """
+    NAME:
+       leapfrog_cyl
+    PURPOSE:
+       leapfrog integrate a Hamiltonian ODE in cylindrical coordinates in 2,3,4,5,6 phase-space dimensions
+    INPUT:
+       drift - (+,-)dH/d(p,q) for drift step, function of y
+       kick -  (+,-)dH/d(p,q) for kick step, function of (y,*args,t=)
+       yo - initial condition [R,vR,vT,z,vz,phi] or galpy subsets
+       t - set of times at which one wants the result
+       rtol, atol
+       meridional= (False) if True, we are integrating in the meridional plane (overloaded as Lz**2)
+    OUTPUT:
+       y : array, shape (len(y0), len(t))
+       Array containing the value of y for each desired time in t, \
+       with the initial value y0 in the first row.
+    HISTORY:
+       2018-12-21 - Written - Bovy (UofT)
+    """
+    #Initialize
+    y= nu.copy(yo)
+    if not meridional and len(yo) != 2:
+        # vT --> Lz
+        y[2]*= y[0]
+    out= nu.zeros((len(t),len(yo)))
+    out[0,:]= y
+    #Estimate necessary step size
+    dt= t[1]-t[0] #assumes that the steps are equally spaced
+    init_dt= dt
+    dt= _leapfrog_cyl_estimate_step(drift,kick,yo,dt,t[0],args,rtol,atol,
+                                    meridional)
+    dt2= dt/2.
+    ndt= int(init_dt/dt)
+    #Integrate
+    to= t[0]
+    for ii in range(1,len(t)):
+        #drift half
+        y12= y+dt2*drift(y)
+        for jj in range(ndt): #loop over number of sub-intervals
+            #kick
+            y12+= dt*kick(y12,*args,t=to+dt/2)
+            #drift full
+            y12+= dt*drift(y12)
+            #Get ready for next
+            to+= dt
+        #drift half back to correct overshoot
+        y= y12-dt2*drift(y12)
+        out[ii,:]= y
+    if not meridional and len(yo) != 2:
+        # Lz --> vT
+        out[:,2]/= out[:,0]
+    return out
+
+def _leapfrog_cyl_estimate_step(drift,kick,yo,dt,to,args,rtol,atol,meridional):
+    init_dt= dt
+    if len(yo) == 6:
+        xscale= nu.sqrt(yo[0]**2.+yo[3]**2.)
+        vscale= nu.sqrt((yo[2]/yo[0])**2.+yo[1]**2.+yo[4]**2.)
+        scale= atol+rtol*nu.array([xscale,vscale,vscale*xscale,
+                                   xscale,vscale,xscale])
+    elif meridional and len(yo) == 4:
+        xscale= nu.sqrt(yo[0]**2.+yo[2]**2.)
+        vscale= nu.sqrt(meridional/yo[0]**2.+yo[1]**2.+yo[3]**2.)
+        scale= atol+rtol*nu.array([xscale,vscale,xscale,vscale])
+    elif len(yo) == 4:
+        xscale= yo[0]
+        vscale= nu.sqrt((yo[2]/yo[0])**2.+yo[1]**2.)
+        scale= atol+rtol*nu.array([xscale,vscale,vscale*xscale,xscale])
+    elif meridional and len(yo) == 2:
+        xscale= nu.fabs(yo[0])
+        vscale= nu.sqrt(meridional/yo[0]**2.+yo[1]**2.)
+        scale= atol+rtol*nu.array([xscale,vscale])
+    else: # len(yo) == 2:
+        xscale= nu.fabs(yo[0])
+        vscale= nu.fabs(yo[1])
+        scale= atol+rtol*nu.array([xscale,vscale])
+    err= 2.
+    dt*= 2.
+    while err > 1. and init_dt/dt < _MAX_DT_REDUCE:
+        #Do one leapfrog step with step dt and one with dt/2.
+        #dt
+        y12= yo+dt/2.*drift(yo)
+        y12+= dt*kick(y12,*args,t=to+dt/2)
+        y11= y12+dt/2.*drift(y12)
+        #dt/2.
+        y12= yo+dt/4.*drift(yo)
+        y12+= dt/2.*kick(y12,*args,t=to+dt/4.)
+        y12+= dt/2.*drift(y12)#Take full step combining two half
+        y12+= dt/2.*kick(y12,*args,t=to+3.*dt/4.)
+        y12+= dt/4.*drift(y12)
+        #Norm
+        delta= nu.fabs(y11-y12)
+        if len(yo) == 6 or (not meridional and len(yo) == 4):
+            delta[-1]= nu.fabs(yo[0]*(y11[-1]-y12[-1]))
+        err= nu.sqrt(nu.mean((delta/scale)**2.))
+        dt/= 2.
+    return dt
+