test_closed_orbit_finder.py

# ----------------------------------------------------------------------------
#
# TITLE - test_closed_orbit_finder.py
# AUTHOR - James Lane
# PROJECT - AST1501
# CONTENTS:
#
# ----------------------------------------------------------------------------

### Docstrings and metadata:
'''Test the new closed orbit finding routine
'''

### Imports

## Basic
import numpy as np
import pdb

## Plotting
from matplotlib import pyplot as plt

## galpy
from galpy import orbit
from galpy import potential

# ----------------------------------------------------------------------------

def find_closed_orbit(pot,Lz,rtol=0.01,R0=1.0,vR0=0.0,plot_loop=False):
    '''find_closed_orbit:
    
    Calculate a closed orbit for a given angular momentum
    
    Args:
        pot (galpy Potential object) - Potential for which to determine 
            the closed orbit
        Lz (float) - Angular momentum for the closed orbit
        rtol (float) - change in radius marking the end of the search [0.01]
        R0 (float) - Starting radius [1.0]
        vR0 (float) - Starting radial velocity [0.0]
        plot_loop (bool) - Plot the surface during each loop evaluation? [False]
    
    Returns:
        orbit (galpy Orbit object) - Orbit object representing a closed 
            orbit in the given potential for the given angular momentum
    '''
    
    # Turn off physical
    potential.turn_physical_off(pot)
    
    # Initialize starting orbit
    o = orbit.Orbit([R0,vR0,Lz/R0,0.0,0.0,0.0])
    
    # Evaluate the while loop
    loop_counter = 0
    delta_R = rtol*2.
    while delta_R > rtol:
        
        # Evaluate the crossing time so the integration can be performed 
        # for ~ long enough. Integrate for 100 crossing times.
        tdyn = 2*np.pi*(R0/np.abs(potential.evaluateRforces(pot,R0,0.0)))**0.5
        times = np.linspace(0,100*tdyn,num=10001)
        o.integrate(times,pot)    
        
        # Evaluate all points where the orbit crosses from negative to positive 
        # phi
        phis = o.phi(times) - np.pi
        shift_phis = np.roll(phis,-1)
        where_cross = (phis[:-1] < 0.)*(shift_phis[:-1] > 0.)
        R_cross = o.R(times)[:-1][where_cross]
        vR_cross = o.vR(times)[:-1][where_cross]
        
        # Plot the surface of section as a test, if asked
        if plot_loop:
            fig = plt.figure()
            ax = fig.add_subplot(111)
            ax.scatter(R_cross,vR_cross,s=20,color='Black')
            ax.set_xlabel(r'$R$')
            ax.set_ylabel(r'$v_{R}$')
            ax.set_title(   r'R='+str(round(o.R(0),6))+\
                            ', vR='+str(round(o.vR(0),6))+\
                            ', vT='+str(round(o.vT(0),6))
                        )
            fig.savefig('./loop_fig'+str(loop_counter)+'.pdf')
        ##fi
        
        # Calculate the difference in radius
        delta_R = np.abs( o.R(0) - np.average( R_cross ) )
        
        # Update the orbit
        o = orbit.Orbit( [  np.average( R_cross ),
                            np.average( vR_cross ),
                            Lz/np.average( R_cross ),
                            0.0,0.0,0.0] )
    
        # Count
        loop_counter += 1
        
    return o
#def

# ----------------------------------------------------------------------------

## Test the function
Lz = 1.0
R0 = 0.5
mwpot = potential.MWPotential2014
closed_orbit = find_closed_orbit(mwpot,Lz,R0=R0,plot_loop=True)

# Finishes in 4 loop passes

closed_orbit.R() # 1.000...
closed_orbit.vR() # 0.000...
closed_orbit.vT() # 1.000...
closed_orbit.Lz() # 1.0

# ----------------------------------------------------------------------------