<commit_before>import numpy as nu
import galpy.util.bovy_plot as plot
from Potential import PotentialError, Potential
from plotRotcurve import plotRotcurve
class planarPotential:
    """Class representing 2D (R,\phi) potentials"""
    def __init__(self,amp=1.):
        self._amp= 1.
        self.dim= 2
        self.isNonAxi= True #Gets reset by planarAxiPotential
        self.isRZ= False
        return None

    def __call__(self,R,phi=0.):
        """
        NAME:
           __call__
        PURPOSE:
           evaluate the potential
        INPUT: 
           Either: R or R,phi [rad]
        OUTPUT:
           Phi(R(,phi)))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        try:
            return self._amp*self._evaluate(R,phi=phi)
        except AttributeError:
            raise PotentialError("'_evaluate' function not implemented for this potential")

    def Rforce(self,R,phi=0.):
        """
        NAME:
           Rforce
        PURPOSE:
           evaluate the radial force
        INPUT:
           Either: R or R,phi [rad]         
        OUTPUT:
           F_R(R,(\phi)))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        try:
            return self._amp*self._Rforce(R,phi=phi)
        except AttributeError:
            raise PotentialError("'_Rforce' function not implemented for this potential")

    def phiforce(self,R,phi=0.):
        """
        NAME:
           phiforce
        PURPOSE:
           evaluate the phi force
        INPUT:
           Either: R or R,phi [rad]         
        OUTPUT:
           F_\phi(R,(\phi)))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        try:
            return self._amp*self._phiforce(R,phi=phi)
        except AttributeError:
            raise PotentialError("'_phiforce' function not implemented for this potential")


class planarAxiPotential(planarPotential):
    """Class representing axisymmetric planar potentials"""
    def __init__(self,amp=1.):
        planarPotential.__init__(self,amp=amp)
        self.isNonAxi= False
        return None
    
    def _phiforce(self,R,phi=0.):
        return 0.

    def plot(self,*args,**kwargs):
        """
        NAME:
           plot
        PURPOSE:
           plot the potential
        INPUT:
           Rrange - range
           grid - number of points to plot
           savefilename - save to or restore from this savefile (pickle)
           +bovy_plot(*args,**kwargs)
        OUTPUT:
           plot to output device
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        plotplanarPotentials(self,*args,**kwargs)

    def plotRotcurve(self,*args,**kwargs):
        """
        NAME:
           plotRotcurve
        PURPOSE:
           plot the rotation curve for this potential
        INPUT:
           Rrange - range
           grid - number of points to plot
           savefilename - save to or restore from this savefile (pickle)
           +bovy_plot(*args,**kwargs)
        OUTPUT:
           plot to output device
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        plotRotcurve(self,*args,**kwargs)

class planarPotentialFromRZPotential(planarAxiPotential):
    """Class that represents an axisymmetic planar potential derived from a 
    RZPotential"""
    def __init__(self,RZPot):
        """
        NAME:
           __init__
        PURPOSE:
           Initialize
        INPUT:
           RZPot - RZPotential instance
        OUTPUT:
           planarAxiPotential instance
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        planarAxiPotential.__init__(self,amp=1.)
        self._RZPot= RZPot
        return None

    def _evaluate(self,R,phi=0.):
        """
        NAME:
           _evaluate
        PURPOSE:
           evaluate the potential
        INPUT:
           Either: R or R,phi [rad]      
        OUTPUT:
          Pot(R(,\phi))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        return self._RZPot(R,0.)
            
    def _Rforce(self,R,phi=0.):
        """
        NAME:
           _Rforce
        PURPOSE:
           evaluate the radial force
        INPUT:
           Either: R or R,phi [rad]      
        OUTPUT:
          F_R(R(,\phi))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        return self._RZPot.Rforce(R,0.)
            
def RZToplanarPotential(RZPot):
    """
    NAME:
       RZToplanarPotential
    PURPOSE:
       convert an RZPotential to a planarPotential in the mid-plane (z=0)
    INPUT:
       RZPot - RZPotential instance or list of such instances (existing 
               planarPotential instances are just copied to the output)
    OUTPUT:
       planarPotential instance(s)
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    if isinstance(RZPot,list):
        out= []
        for pot in RZPot:
            if isinstance(pot,planarPotential):
                out.append(pot)
            else:
                out.append(planarPotentialFromRZPotential(pot))
        return out
    elif isinstance(RZPot,Potential):
        return planarPotentialFromRZPotential(RZPot)
    elif isinstance(RZPot,planarPotential):
        return RZPot
    else:
        raise PotentialError("Input to 'RZToplanarPotential' is neither an RZPotential-instance or a list of such instances")

def evaluateplanarPotentials(R,Pot,phi=0.):
    """
    NAME:
       evaluateplanarPotentials
    PURPOSE:
       evaluate a (list of) planarPotential instance(s)
    INPUT:
       R (+phi optional)
       Pot - (list of) planarPotential instance(s)
    OUTPUT:
       Phi(R(,phi))
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    isList= isinstance(Pot,list)
    nonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if nonAxi and not hasphi:
        raise PotentialError("The (list of) planarPotential instances is non-axisymmetric, but you did not provide phi")
    if isinstance(Pot,list):
        sum= 0.
        for pot in Pot:
            if nonAxi:
                sum+= pot(R,phi=phi)
            else:
                sum+= pot(R)
        return sum
    elif isinstance(Pot,planarPotential):
        if nonAxi:
            return Pot(R,phi=phi)
        else:
            return Pot(R)
    else:
        raise PotentialError("Input to 'evaluateplanarPotentials' is neither a Potential-instance or a list of such instances")

def evaluateplanarRforces(R,Pot,phi=0.):
    """
    NAME:
       evaluateplanarRforces
    PURPOSE:
       evaluate the Rforce of a (list of) planarPotential instance(s)
    INPUT:
       R (+phi optional)
       Pot - (list of) planarPotential instance(s)
    OUTPUT:
       F_R(R(,phi))
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    isList= isinstance(Pot,list)
    nonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if nonAxi and not hasphi:
        raise PotentialError("The (list of) planarPotential instances is non-axisymmetric, but you did not provide phi")
    if isinstance(Pot,list):
        sum= 0.
        for pot in Pot:
            if nonAxi:
                sum+= pot.Rforce(R,phi=phi)
            else:
                sum+= pot.Rforce(R)
        return sum
    elif isinstance(Pot,planarPotential):
        if nonAxi:
            return Pot.Rforce(R,phi=phi)
        else:
            return Pot.Rforce(R)
    else:
        raise PotentialError("Input to 'evaluateplanarRforces' is neither a Potential-instance or a list of such instances")

def evaluateplanarphiforces(R,Pot,phi=0.):
    """
    NAME:
       evaluateplanarphiforces
    PURPOSE:
       evaluate the phiforce of a (list of) planarPotential instance(s)
    INPUT:
       R (+phi optional)
       Pot - (list of) planarPotential instance(s)
    OUTPUT:
       F_phi(R(,phi))
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    isList= isinstance(Pot,list)
    nonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if nonAxi and not hasphi:
        raise PotentialError("The (list of) planarPotential instances is non-axisymmetric, but you did not provide phi")
    if isinstance(Pot,list):
        sum= 0.
        for pot in Pot:
            if nonAxi:
                sum+= pot.phiforce(R,phi=phi)
            else:
                sum+= pot.phiforce(R)
        return sum
    elif isinstance(Pot,planarPotential):
        if nonAxi:
            return Pot.phiforce(R,phi=phi)
        else:
            return Pot.phiforce(R)
    else:
        raise PotentialError("Input to 'evaluateplanarphiforces' is neither a Potential-instance or a list of such instances")

def plotplanarPotentials(Pot,*args,**kwargs):
    """
    NAME:
       plotplanarPotentials
    PURPOSE:
       plot a planar potential
    INPUT:<commit_msg>add time-dependence to planar potentials<commit_after>import numpy as nu
import galpy.util.bovy_plot as plot
from Potential import PotentialError, Potential
from plotRotcurve import plotRotcurve
class planarPotential:
    """Class representing 2D (R,\phi) potentials"""
    def __init__(self,amp=1.):
        self._amp= 1.
        self.dim= 2
        self.isNonAxi= True #Gets reset by planarAxiPotential
        self.isRZ= False
        return None

    def __call__(self,R,phi=0.,t=0.):
        """
        NAME:
           __call__
        PURPOSE:
           evaluate the potential
        INPUT: 
           R
           phi
           t
        OUTPUT:
           Phi(R(,phi,t)))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        try:
            return self._amp*self._evaluate(R,phi=phi,t=t)
        except AttributeError:
            raise PotentialError("'_evaluate' function not implemented for this potential")

    def Rforce(self,R,phi=0.,t=0.):
        """
        NAME:
           Rforce
        PURPOSE:
           evaluate the radial force
        INPUT:
           R
           phi
           t
        OUTPUT:
           F_R(R,(\phi,t)))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        try:
            return self._amp*self._Rforce(R,phi=phi,t=t)
        except AttributeError:
            raise PotentialError("'_Rforce' function not implemented for this potential")

    def phiforce(self,R,phi=0.,t=0.):
        """
        NAME:
           phiforce
        PURPOSE:
           evaluate the phi force
        INPUT:
           R
           phi
           t
        OUTPUT:
           F_\phi(R,(\phi,t)))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        try:
            return self._amp*self._phiforce(R,phi=phi,t=t)
        except AttributeError:
            raise PotentialError("'_phiforce' function not implemented for this potential")


class planarAxiPotential(planarPotential):
    """Class representing axisymmetric planar potentials"""
    def __init__(self,amp=1.):
        planarPotential.__init__(self,amp=amp)
        self.isNonAxi= False
        return None
    
    def _phiforce(self,R,phi=0.,t=0.):
        return 0.

    def plot(self,*args,**kwargs):
        """
        NAME:
           plot
        PURPOSE:
           plot the potential
        INPUT:
           Rrange - range
           grid - number of points to plot
           savefilename - save to or restore from this savefile (pickle)
           +bovy_plot(*args,**kwargs)
        OUTPUT:
           plot to output device
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        plotplanarPotentials(self,*args,**kwargs)

    def plotRotcurve(self,*args,**kwargs):
        """
        NAME:
           plotRotcurve
        PURPOSE:
           plot the rotation curve for this potential
        INPUT:
           Rrange - range
           grid - number of points to plot
           savefilename - save to or restore from this savefile (pickle)
           +bovy_plot(*args,**kwargs)
        OUTPUT:
           plot to output device
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        plotRotcurve(self,*args,**kwargs)

class planarPotentialFromRZPotential(planarAxiPotential):
    """Class that represents an axisymmetic planar potential derived from a 
    RZPotential"""
    def __init__(self,RZPot):
        """
        NAME:
           __init__
        PURPOSE:
           Initialize
        INPUT:
           RZPot - RZPotential instance
        OUTPUT:
           planarAxiPotential instance
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        planarAxiPotential.__init__(self,amp=1.)
        self._RZPot= RZPot
        return None

    def _evaluate(self,R,phi=0.,t=0.):
        """
        NAME:
           _evaluate
        PURPOSE:
           evaluate the potential
        INPUT:
           R
           phi
           t
        OUTPUT:
          Pot(R(,\phi,t))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        return self._RZPot(R,0.,t=t)
            
    def _Rforce(self,R,phi=0.,t=0.):
        """
        NAME:
           _Rforce
        PURPOSE:
           evaluate the radial force
        INPUT:
           R
           phi
           t
        OUTPUT:
          F_R(R(,\phi,t))
        HISTORY:
           2010-07-13 - Written - Bovy (NYU)
        """
        return self._RZPot.Rforce(R,0.,t=t)
            
def RZToplanarPotential(RZPot):
    """
    NAME:
       RZToplanarPotential
    PURPOSE:
       convert an RZPotential to a planarPotential in the mid-plane (z=0)
    INPUT:
       RZPot - RZPotential instance or list of such instances (existing 
               planarPotential instances are just copied to the output)
    OUTPUT:
       planarPotential instance(s)
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    if isinstance(RZPot,list):
        out= []
        for pot in RZPot:
            if isinstance(pot,planarPotential):
                out.append(pot)
            else:
                out.append(planarPotentialFromRZPotential(pot))
        return out
    elif isinstance(RZPot,Potential):
        return planarPotentialFromRZPotential(RZPot)
    elif isinstance(RZPot,planarPotential):
        return RZPot
    else:
        raise PotentialError("Input to 'RZToplanarPotential' is neither an RZPotential-instance or a list of such instances")

def evaluateplanarPotentials(R,Pot,phi=0.,t=0.):
    """
    NAME:
       evaluateplanarPotentials
    PURPOSE:
       evaluate a (list of) planarPotential instance(s)
    INPUT:
       R (+phi optional)
       Pot - (list of) planarPotential instance(s)
       t - time (optional)
    OUTPUT:
       Phi(R(,phi,t))
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    isList= isinstance(Pot,list)
    nonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if nonAxi and not hasphi:
        raise PotentialError("The (list of) planarPotential instances is non-axisymmetric, but you did not provide phi")
    if isinstance(Pot,list):
        sum= 0.
        for pot in Pot:
            if nonAxi:
                sum+= pot(R,phi=phi,t=t)
            else:
                sum+= pot(R,t=t)
        return sum
    elif isinstance(Pot,planarPotential):
        if nonAxi:
            return Pot(R,phi=phi,t=t)
        else:
            return Pot(R,t=t)
    else:
        raise PotentialError("Input to 'evaluateplanarPotentials' is neither a Potential-instance or a list of such instances")

def evaluateplanarRforces(R,Pot,phi=0.,t=0.):
    """
    NAME:
       evaluateplanarRforces
    PURPOSE:
       evaluate the Rforce of a (list of) planarPotential instance(s)
    INPUT:
       R (+phi optional)
       Pot - (list of) planarPotential instance(s)
       t - time (optional)
    OUTPUT:
       F_R(R(,phi,t))
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    isList= isinstance(Pot,list)
    nonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if nonAxi and not hasphi:
        raise PotentialError("The (list of) planarPotential instances is non-axisymmetric, but you did not provide phi")
    if isinstance(Pot,list):
        sum= 0.
        for pot in Pot:
            if nonAxi:
                sum+= pot.Rforce(R,phi=phi,t=t)
            else:
                sum+= pot.Rforce(R,t=t)
        return sum
    elif isinstance(Pot,planarPotential):
        if nonAxi:
            return Pot.Rforce(R,phi=phi,t=t)
        else:
            return Pot.Rforce(R,t=t)
    else:
        raise PotentialError("Input to 'evaluateplanarRforces' is neither a Potential-instance or a list of such instances")

def evaluateplanarphiforces(R,Pot,phi=0.,t=0.):
    """
    NAME:
       evaluateplanarphiforces
    PURPOSE:
       evaluate the phiforce of a (list of) planarPotential instance(s)
    INPUT:
       R (+phi optional)
       Pot - (list of) planarPotential instance(s)
       t - time (optional)
    OUTPUT:
       F_phi(R(,phi,t))
    HISTORY:
       2010-07-13 - Written - Bovy (NYU)
    """
    isList= isinstance(Pot,list)
    nonAxi= ((isList and Pot[0].isNonAxi) or (not isList and Pot.isNonAxi))
    if nonAxi and not hasphi:
        raise PotentialError("The (list of) planarPotential instances is non-axisymmetric, but you did not provide phi")
    if isinstance(Pot,list):
        sum= 0.
        for pot in Pot:
            if nonAxi:
                sum+= pot.phiforce(R,phi=phi,t=t)
            else:
                sum+= pot.phiforce(R,t=t)
        return sum
    elif isinstance(Pot,planarPotential):
        if nonAxi:
            return Pot.phiforce(R,phi=phi,t=t)
        else:
            return Pot.phiforce(R,t=t)
    else:
        raise PotentialError("Input to 'evaluateplanarphiforces' is neither a Potential-instance or a list of such instances")

def plotplanarPotentials(Pot,*args,**kwargs):
    """
    NAME:
       plotplanarPotentials
    PURPOSE:
       plot a planar potential
    INPUT:
{'fraction_non_alphanumeric': 0.0944055944055944, 'fraction_numerical': 0.01393051393051393, 'mean_word_length': 1.8018970611102472, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2962897', 'n_tokens_mistral': 5982, 'n_tokens_neox': 5435, 'n_words': 1428}