<commit_before>import numpy as nu
import galpy.util.bovy_coords as coords
from FullOrbit import FullOrbit
from RZOrbit import RZOrbit
from planarOrbit import planarOrbit, planarROrbit
from linearOrbit import linearOrbit
class Orbit:
    """General orbit class representing an orbit"""
    def __init__(self,vxvv=None,uvw=False,
                 radec=False,vo=235.,ro=8.5,
                 solarmotion='hogg'):
        """
        NAME:
           __init__
        PURPOSE:
           Initialize an Orbit instance
        INPUT:
           vxvv - initial conditions 
                  3D can be either
              1) in Galactocentric cylindrical coordinates
              2) [ra,dec,d,mu_ra, mu_dec,vlos] in [deg,deg,kpc,mas/yr,mas/yr,km/s] (all 2000.0)
              3) [ra,dec,d,U,V,W] in [deg,deg,kpc,km/s,km/s,kms]
        OPTIONAL INPUTS:
           radec - if True, input is 2) (or 3) above
           uvw - if True, velocities are UVW
           vo - circular velocity at ro
           ro - distance from vantage point to GC
           solarmotion - 'hogg' or 'dehnen', or 'schoenrich'
        OUTPUT:
           instance
        HISTORY:
           2010-07-20 - Written - Bovy (NYU)
        """
        if solarmotion.lower() == 'hogg':
            vsolar= nu.array([-10.1,4.0,6.7])/vo
        elif solarmotion.lower() == 'dehnen':
            vsolar= nu.array([-10.,5.25,7.17])/vo
        elif solarmotion.lower() == 'schoenrich':
            vsolar= nu.array([-11.1,12.24,7.25])/vo
        if radec:
            l,b= coords.radec_to_lb(vxvv[0],vxvv[1],degree=True)
            if uvw:
                X,Y,Z= coords.lbd_to_XYZ(l,b,vxvv[2])
                vx= vxvv[3]
                vy= vxvv[4]
                vz= vxvv[5]
            else:
                pmll, pmbb= coords.pmrapmdec_to_pmllpmbb(vxvv[3],vxvv[4],
                                                         vxvv[0],vxvv[1],
                                                         degree=True)
                X,Y,Z,vx,vy,vz= coords.sphergal_to_rectgal(l,b,vxvv[2],
                                                           vxvv[5],pmll, pmbb,
                                                           degree=True)
            X/= ro
            Y/= ro
            Z/= ro
            vx/= vo
            vy/= vo
            vz/= vo
            vsun= nu.array([0.,1.,0.,])+vsolar
            R, phi, z= coords.XYZ_to_galcencyl(X,Y,Z)
            vR, vT,vz= coords.vxvyvz_to_galcencyl(vx,vy,vz,
                                                  R,phi,z,
                                                  vsun=vsun,galcen=True)
            vxvv= [R,vR,vT,z,vz,phi]
        elif radec and uvw:
            pass
        self.vxvv= vxvv
        if len(vxvv) == 2:
            self._orb= linearOrbit(vxvv=vxvv)
        elif len(vxvv) == 3:
            self._orb= planarROrbit(vxvv=vxvv)
        elif len(vxvv) == 4:
            self._orb= planarOrbit(vxvv=vxvv)
        elif len(vxvv) == 5:
            self._orb= RZOrbit(vxvv=vxvv)
        elif len(vxvv) == 6:
            self._orb= FullOrbit(vxvv=vxvv)

    def setphi(self,phi):
        """
        NAME:
           setphi
        PURPOSE:
           set initial azimuth
        INPUT:
           phi - desired azimuth
        OUTPUT:
           (none)
        HISTORY:
           2010-08-01 - Written - Bovy (NYU)
        BUGS:
           Should perform check that this orbit has phi
        """
        if len(self.vxvv) == 2:
            raise AttributeError("One-dimensional orbit has no azimuth")
        elif len(self.vxvv) == 3:
            #Upgrade
            vxvv= [self.vxvv[0],self.vxvv[1],self.vxvv[2],phi]
            self.vxvv= vxvv
            self._orb= planarROrbit(vxvv=vxvv)
        elif len(self.vxvv) == 4:
            self.vxvv[-1]= phi
            self._orb.vxvv[-1]= phi
        elif len(self.vxvv) == 5:
            #Upgrade
            vxvv= [self.vxvv[0],self.vxvv[1],self.vxvv[2],self.vxvv[3],
                   self.vxvv[4],phi]
            self.vxvv= vxvv
            self._orb= FullOrbit(vxvv=vxvv)
        elif len(self.vxvv) == 6:
            self.vxvv[-1]= phi
            self._orb.vxvv[-1]= phi

    def integrate(self,t,pot,method='odeint'):
        """
        NAME:
           integrate
        PURPOSE:
           integrate the orbit
        INPUT:
           t - list of times at which to output (0 has to be in this!)

           pot - potential instance or list of instances
           method= 'odeint' for scipy's odeint or 'leapfrog' for a simple
                   leapfrog implementation
        OUTPUT:
           (none) (get the actual orbit using getOrbit()<commit_msg>add dim() function to Orbit<commit_after>import numpy as nu
import galpy.util.bovy_coords as coords
from FullOrbit import FullOrbit
from RZOrbit import RZOrbit
from planarOrbit import planarOrbit, planarROrbit
from linearOrbit import linearOrbit
class Orbit:
    """General orbit class representing an orbit"""
    def __init__(self,vxvv=None,uvw=False,
                 radec=False,vo=235.,ro=8.5,
                 solarmotion='hogg'):
        """
        NAME:
           __init__
        PURPOSE:
           Initialize an Orbit instance
        INPUT:
           vxvv - initial conditions 
                  3D can be either
              1) in Galactocentric cylindrical coordinates
              2) [ra,dec,d,mu_ra, mu_dec,vlos] in [deg,deg,kpc,mas/yr,mas/yr,km/s] (all 2000.0)
              3) [ra,dec,d,U,V,W] in [deg,deg,kpc,km/s,km/s,kms]
        OPTIONAL INPUTS:
           radec - if True, input is 2) (or 3) above
           uvw - if True, velocities are UVW
           vo - circular velocity at ro
           ro - distance from vantage point to GC
           solarmotion - 'hogg' or 'dehnen', or 'schoenrich'
        OUTPUT:
           instance
        HISTORY:
           2010-07-20 - Written - Bovy (NYU)
        """
        if solarmotion.lower() == 'hogg':
            vsolar= nu.array([-10.1,4.0,6.7])/vo
        elif solarmotion.lower() == 'dehnen':
            vsolar= nu.array([-10.,5.25,7.17])/vo
        elif solarmotion.lower() == 'schoenrich':
            vsolar= nu.array([-11.1,12.24,7.25])/vo
        if radec:
            l,b= coords.radec_to_lb(vxvv[0],vxvv[1],degree=True)
            if uvw:
                X,Y,Z= coords.lbd_to_XYZ(l,b,vxvv[2])
                vx= vxvv[3]
                vy= vxvv[4]
                vz= vxvv[5]
            else:
                pmll, pmbb= coords.pmrapmdec_to_pmllpmbb(vxvv[3],vxvv[4],
                                                         vxvv[0],vxvv[1],
                                                         degree=True)
                X,Y,Z,vx,vy,vz= coords.sphergal_to_rectgal(l,b,vxvv[2],
                                                           vxvv[5],pmll, pmbb,
                                                           degree=True)
            X/= ro
            Y/= ro
            Z/= ro
            vx/= vo
            vy/= vo
            vz/= vo
            vsun= nu.array([0.,1.,0.,])+vsolar
            R, phi, z= coords.XYZ_to_galcencyl(X,Y,Z)
            vR, vT,vz= coords.vxvyvz_to_galcencyl(vx,vy,vz,
                                                  R,phi,z,
                                                  vsun=vsun,galcen=True)
            vxvv= [R,vR,vT,z,vz,phi]
        elif radec and uvw:
            pass
        self.vxvv= vxvv
        if len(vxvv) == 2:
            self._orb= linearOrbit(vxvv=vxvv)
        elif len(vxvv) == 3:
            self._orb= planarROrbit(vxvv=vxvv)
        elif len(vxvv) == 4:
            self._orb= planarOrbit(vxvv=vxvv)
        elif len(vxvv) == 5:
            self._orb= RZOrbit(vxvv=vxvv)
        elif len(vxvv) == 6:
            self._orb= FullOrbit(vxvv=vxvv)

    def setphi(self,phi):
        """
        NAME:
           setphi
        PURPOSE:
           set initial azimuth
        INPUT:
           phi - desired azimuth
        OUTPUT:
           (none)
        HISTORY:
           2010-08-01 - Written - Bovy (NYU)
        BUGS:
           Should perform check that this orbit has phi
        """
        if len(self.vxvv) == 2:
            raise AttributeError("One-dimensional orbit has no azimuth")
        elif len(self.vxvv) == 3:
            #Upgrade
            vxvv= [self.vxvv[0],self.vxvv[1],self.vxvv[2],phi]
            self.vxvv= vxvv
            self._orb= planarROrbit(vxvv=vxvv)
        elif len(self.vxvv) == 4:
            self.vxvv[-1]= phi
            self._orb.vxvv[-1]= phi
        elif len(self.vxvv) == 5:
            #Upgrade
            vxvv= [self.vxvv[0],self.vxvv[1],self.vxvv[2],self.vxvv[3],
                   self.vxvv[4],phi]
            self.vxvv= vxvv
            self._orb= FullOrbit(vxvv=vxvv)
        elif len(self.vxvv) == 6:
            self.vxvv[-1]= phi
            self._orb.vxvv[-1]= phi

    def dim(self):
        """
        NAME:
           dim
        PURPOSE:
           return the dimension of the problem
        INPUT:
        OUTPUT:
           dimension
        HISTORY:
           2011-02-03 - Written - Bovy (NYU)
        """
        if len(self.vxvv) == 2:
            return 1
        elif len(self.vxvv) == 3 or len(self.vxvv) == 4:
            return 2
        elif len(self.vxvv) == 5 or len(self.vxvv) == 6:
            return 3

    def integrate(self,t,pot,method='odeint'):
        """
        NAME:
           integrate
        PURPOSE:
           integrate the orbit
        INPUT:
           t - list of times at which to output (0 has to be in this!)

           pot - potential instance or list of instances
           method= 'odeint' for scipy's odeint or 'leapfrog' for a simple
                   leapfrog implementation
        OUTPUT:
           (none) (get the actual orbit using getOrbit()
{'fraction_non_alphanumeric': 0.11136131013306039, 'fraction_numerical': 0.021084953940634595, 'mean_word_length': 1.3165007112375533, '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': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1517565', 'n_tokens_mistral': 3431, 'n_tokens_neox': 3156, 'n_words': 701}