import numpy as nu
import galpy.util.bovy_plot as plot
def plotRotcurve(Pot,*args,**kwargs):
    """
    NAME:
       plotRotcurve
    PURPOSE:
       plot the rotation curve for this potential (in the z=0 plane for
       non-spherical potentials)
    INPUT:
       Pot - Potential or list of Potential instances

       Rrange - 

       grid - grid in R

       savefilename - save to or restore from this savefile (pickle)
       +bovy_plot.bovy_plot args and kwargs
    OUTPUT:
       plot to output device
    BUGS:
       Should use calcRotcurve
    HISTORY:
       2010-07-10 - Written - Bovy (NYU)
    """
    if kwargs.has_key('Rrange'):
        Rrange= kwargs['Rrange']
        kwargs.pop('Rrange')
    else:
        Rrange= [0.01,5.]
    if kwargs.has_key('grid'):
        grid= kwargs['grid']
        kwargs.pop('grid')
    else:
        grid= 1001
    if kwargs.has_key('savefilename'):
        savefilename= kwargs['savefilename']
        kwargs.pop('savefilename')
    else:
        savefilename= None
    isList= isinstance(Pot,list)
    isNonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if isNonAxi:
        raise AttributeError("Rotation curve plotting for non-axisymmetric potentials is not currently supported")
    if not savefilename == None and os.path.exists(savefilename):
        print "Restoring savefile "+savefilename+" ..."
        savefile= open(savefilename,'rb')
        rotcurve= pickle.load(savefile)
        Rs= pickle.load(savefile)
        savefile.close()
    else:
        Rs= nu.linspace(Rrange[0],Rrange[1],grid)
        rotcurve= nu.zeros(grid)
        from planarPotential import evaluateplanarRforces
        for ii in range(grid):
            try:
                rotcurve[ii]= nu.sqrt(Rs[ii]*-evaluateplanarRforces(Rs[ii],Pot))

            except TypeError:
                from planarPotential import RZToplanarPotential
                Pot= RZToplanarPotential(Pot)
                rotcurve[ii]= nu.sqrt(Rs[ii]*-evaluateplanarRforces(Rs[ii],Pot))
        if not savefilename == None:
            print "Writing savefile "+savefilename+" ..."
            savefile= open(savefilename,'wb')
            pickle.dump(rotcurve,savefile)
            pickle.dump(Rs,savefile)
            savefile.close()
    if not kwargs.has_key('xlabel'):
        kwargs['xlabel']= r"$R/R_0$"
    if not kwargs.has_key('ylabel'):
        kwargs['ylabel']= r"$v_c(R)/v_c(R_0)$"
    kwargs['xrange']= Rrange
    return plot.bovy_plot(Rs,rotcurve,*args,
                          **kwargs)

def calcRotcurve(Pot,Rs):
    """
    NAME:
       calctRotcurve
    PURPOSE:
       calculate the rotation curve for this potential (in the z=0 plane for
       non-spherical potentials)
    INPUT:
       Pot - Potential or list of Potential instances

       Rs - (array of) radius(i)
    OUTPUT:
       array of vc
    HISTORY:
       2011-04-13 - Written - Bovy (NYU)
    """
    isList= isinstance(Pot,list)
    isNonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if isNonAxi:
        raise AttributeError("Rotation curve plotting for non-axisymmetric potentials is not currently supported")
    try:
        grid= len(Rs)
    except TypeError:
        grid=1
        Rs= nu.array([Rs])
    rotcurve= nu.zeros(grid)
    from planarPotential import evaluateplanarRforces
    for ii in range(grid):
        try:
            rotcurve[ii]= nu.sqrt(Rs[ii]*-evaluateplanarRforces(Rs[ii],Pot))
        except TypeError:
            from planarPotential import RZToplanarPotential
            Pot= RZToplanarPotential(Pot)
            rotcurve[ii]= nu.sqrt(Rs[ii]*-evaluateplanarRforces(Rs[ii],Pot))
    return rotcurve


@@ -1015,251 +1015,8 @@
 one%0A
-    isList= isinstance(Pot,list)%0A    isNonAxi= ((isList and Pot%5B0%5D.isNonAxi) or (not isList and Pot.isNonAxi))%0A    if isNonAxi:%0A        raise AttributeError(%22Rotation curve plotting for non-axisymmetric potentials is not currently supported%22)%0A
     
@@ -1356,430 +1356,27 @@
 ve= 
-nu.zeros(grid)%0A        from planarPotential import evaluateplanarRforces%0A        for ii in range(grid):%0A            try:%0A                rotcurve%5Bii%5D= nu.sqrt(Rs%5Bii%5D*-evaluateplanarRforces(Rs%5Bii%5D,Pot))%0A%0A            except TypeError:%0A                from planarPotential import RZToplanarPotential%0A                Pot= RZToplanarPotential(Pot)%0A                rotcurve%5Bii%5D= nu.sqrt(Rs%5Bii%5D*-evaluateplanarRforces(Rs%5Bii%5D,Pot)
+calcRotcurve(Pot,Rs
 )%0A  
@@ -1947,17 +1947,16 @@
     calc
-t
 Rotcurve

make plotRotcurve use calcRotcurve