Python orbit.Orbit类代码示例

def sample_los(options,args,dfc):
    print "Sampling the LOS ..."
    n= 0
    out= []
    while n < options.nobjects:
        thisos= dfc.sampleLOS(los=options.los,n=options.nobjects-n)
        #add distance uncertainty
        thisout= []
        for o in thisos:
            basico= Orbit([o.R(),o.vR(),o.vT(),0.,0.,o.phi()])#HACK bc ll does not work in 2D currently
            thisd= basico.dist(obs=[1.,0.,0.],ro=1.)
            #Add uncertainty logarithmically, like photometric distance uncertainty TO DO
            thisd+= nu.random.normal()*thisd*options.distuncertainty
            thisout.append(Orbit([basico.ll(obs=[1.,0.,0.],ro=1.)[0],
                                  thisd[0],
                                  basico.pmll(obs=[1.,0.,0.,0.,1.,0.],
                                              ro=1.,vo=1.)[0],
                                  basico.vlos(obs=[1.,0.,0.,0.,1.,0.],
                                              ro=1.,vo=1.)[0]],
                                 lb=True,
                                 solarmotion=[0.,0.,0.],ro=1.,vo=1.))
        ds= nu.array([o.dist(obs=[1.,0.,0.],ro=1.) for o in thisout]).flatten()
        thisos= [thisout[ii] for ii in range(len(thisout)) \
                     if (ds[ii] < options.dmax)]
        out.extend(thisos)
        n+= len(thisos)
    if len(out) > options.nobjects:
        out= out[0:options.nobjects-1]
    #Save
    picklefile= open(args[0],'wb')
    pickle.dump(out,picklefile)
    picklefile.close()
    return None                   

def test_correctInitialSolarMotion():
    age = np.pi # Half a galactic revolution
    
    # For reference: the expected solar motion as provided by Schoenrich
    REFERENCE_SCHOENRICH_SOLAR_MOTION = np.array([-11.1, -12.24, -7.25])
    ntimesteps = 10

    # vxvv is in cylindrical coordinates here. [R,vR,vT(,z,vz,phi)]
    lsr_orbit = Orbit(vxvv=[1,0,1,0,0,0],vo=220,ro=8,solarmotion='schoenrich')
    lsr_orbit.integrate(np.linspace(0.,age,ntimesteps), MWPotential2014)

    U = lsr_orbit.U(0)
    V = lsr_orbit.V(0)
    W = lsr_orbit.W(0)

    results = np.array([U, V, W]).reshape(3)
    for i, vel in enumerate('UVW'):
        print("For velocity {}:".format(vel))
        print("    expected: {}".format(REFERENCE_SCHOENRICH_SOLAR_MOTION[i]))
        print("    received: {}".format(results[i]))

    assert (np.allclose(REFERENCE_SCHOENRICH_SOLAR_MOTION, results)),\
        '!!! Using galpy version {} Need galpy version 1.1 !!!'.format(
            galpy.__version__
        )

def observed_to_xyzuvw_orbit(obs, ts, lsr_orbit=None):
    """
    Convert six-parameter astrometric solution to XYZUVW orbit.

    Parameters
    ----------
    obs :   [RA (deg), DEC (deg), pi (mas),
             mu_ra (mas/yr), mu_dec (mas/yr), vlos (km/s)]
        Current kinematics
    ts : [ntimes] array
        times (in Myr) to traceback to

    lsr_orbit : Orbit
        the orbit of the local standard of rest for comparison, if None can
        calculate on the fly

    XYZUVW : [ntimes, 6] array
        The space position and velocities of the star in a co-rotating frame
        centred on the LSR
    """
    #convert times from Myr into bovy_time
    bovy_ts = ts / (bovy_conversion.time_in_Gyr(220.,8.)*1000) # bovy-time/Myr
    logging.info("max time in Myr: {}".format(np.max(ts)))
    logging.info("max time in Bovy time: {}".format(np.max(bovy_ts)))

    o = Orbit(vxvv=obs, radec=True, solarmotion='schoenrich')
    o.integrate(bovy_ts,mp,method='odeint')
    data = o.getOrbit()
    XYZUVW = galpy_coords_to_xyzuvw(data, bovy_ts)
    return XYZUVW

def traceback2(params,times):
    """Trace forward a cluster. First column of returned array is the
    position of the cluster at a given age.

    Parameters
    ----------
    times: float array
        Times to trace forward, in Myr. Note that positive numbers are going
        forward in time.
    params: float array
        [RA,DE,Plx,PM(RA),PM(DE),RV]
        RA = Right Ascension (Deg)
        DE = Declination (Deg)
        Plx = Paralax (Mas)
        PM(RA) = Proper motion (Right Ascension) (mas/yr)
        PM(DE) = Proper motion (Declination) (mas/yr)
        RV = Radial Velocity (km/s)
    age: Age of cluster, in Myr
    
    """
    #FIXME: This is very out of date and should be deleted!!!
    #Times in Myr
    ts = -(times/1e3)/bovy_conversion.time_in_Gyr(220.,8.)
    nts = len(times)

    #Positions and velocities in the co-rotating solar reference frame.
    xyzuvw = np.zeros( (1,nts,6) )

    #Trace forward the local standard of rest.
    lsr_orbit= Orbit(vxvv=[1.,0,1,0,0.,0],vo=220,ro=8)
    lsr_orbit.integrate(ts,MWPotential2014)#,method='odeint')

    xyzuvw = integrate_xyzuvw(params,ts,lsr_orbit,MWPotential2014)
    return xyzuvw

def test_estimateDelta():

    #_____initialize some KKSPot_____
    Delta = 1.0
    pot = KuzminKutuzovStaeckelPotential(ac=20.,Delta=Delta,normalize=True)

    #_____initialize an orbit (twice)_____
    vxvv = [1.,0.1,1.1,0.01,0.1]
    o= Orbit(vxvv=vxvv)

    #_____integrate the orbit with C_____
    ts= numpy.linspace(0,101,100)
    o.integrate(ts,pot,method='leapfrog_c')


    #____estimate Focal length Delta_____
    #for each time step individually:
    deltas_estimate = numpy.zeros(len(ts))
    for ii in range(len(ts)):
        deltas_estimate[ii] = estimateDeltaStaeckel(pot,o.R(ts[ii]),o.z(ts[ii]))

    assert numpy.all(numpy.fabs(deltas_estimate - Delta) < 10.**-8), \
            'Focal length Delta estimated along the orbit is not constant.'

    #for all time steps together:
    delta_estimate = estimateDeltaStaeckel(pot,o.R(ts),o.z(ts))
    
    assert numpy.fabs(delta_estimate - Delta) < 10.**-8, \
            'Focal length Delta estimated from the orbit is not the same as the input focal length.'

    return None

def trace_cartesian_orbit(xyzuvw_start, times=None, single_age=True,
                          potential=MWPotential2014, ro=8., vo=220.):
    """
    Given a star's XYZUVW relative to the LSR (at any time), project its
    orbit forward (or backward) to each of the times listed in *times*

    Positive times --> traceforward
    Negative times --> traceback

    TODO: Primary source of inefficiencies, 1366.2 (s)

    Parameters
    ----------
    xyzuvw : [pc,pc,pc,km/s,km/s,km/s]
    times : (float) or ([ntimes] float array)
        Myr - time of 0.0 must be present in the array. Times need not be
        spread linearly.
    single_age: (bool) {True}
        Set this flag if only providing a single age to trace to

    Returns
    -------
    xyzuvw_tf : [ntimes, 6] array
        [pc, pc, pc, km/s, km/s, km/s] - the traced orbit with positions
        and velocities

    Notes
    -----
    Profiling comments have been left in for future reference, but note
    that the profiling was done with previous versions of coordinate
    functions - ones that utilised astropy.units (and thus symbolic algebra)
    """
    if single_age:
        # replace 0 with some tiny number
        if times == 0.:
            times = 1e-15
        times = np.array([0., times])
    else:
        times = np.array(times)

    xyzuvw_start = np.copy(xyzuvw_start).astype(np.float)

    bovy_times = convert_myr2bovytime(times)

    # since the LSR is constant in chron coordinates, the starting point
    # is always treated as time 0
    galpy_coords = convert_cart2galpycoords(xyzuvw_start, ts=0.,
                                            ro=ro, vo=vo)
    o = Orbit(vxvv=galpy_coords, ro=ro, vo=vo)
    o.integrate(bovy_times, potential, method='odeint')

    xyzuvw = convert_galpycoords2cart(o.getOrbit(), bovy_times,
                                      ro=ro, vo=vo)
    if single_age:
        return xyzuvw[-1]
    return xyzuvw

def evolveorbit(icon, ti, tau, pot):
    global x
    global y
    o = Orbit(vxvv=icon) # [R,vR,vT,z,vz,phi]
    tf = ti+tau
    ts = np.linspace(ti,tf,100)
    o.integrate(ts, pot, method = 'leapfrog')
    x.append(o.x(ts[0]))
    y.append(o.y(ts[0]))
    return [o.R(tf),o.vR(tf),o.vT(tf),o.z(tf),o.vz(tf),o.phi(tf)]

def test_orbitint():
    import numpy
    from galpy.potential import MWPotential2014
    from galpy.potential import evaluatePotentials as evalPot
    from galpy.orbit import Orbit
    E, Lz= -1.25, 0.6
    o1= Orbit([0.8,0.,Lz/0.8,0.,numpy.sqrt(2.*(E-evalPot(0.8,0.,MWPotential2014)-(Lz/0.8)**2./2.)),0.])
    ts= numpy.linspace(0.,100.,2001)
    o1.integrate(ts,MWPotential2014)
    o1.plot(xrange=[0.3,1.],yrange=[-0.2,0.2],color='k')
    o2= Orbit([0.8,0.3,Lz/0.8,0.,numpy.sqrt(2.*(E-evalPot(0.8,0.,MWPotential2014)-(Lz/0.8)**2./2.-0.3**2./2.)),0.])
    o2.integrate(ts,MWPotential2014)
    o2.plot(xrange=[0.3,1.],yrange=[-0.2,0.2],color='k')
    return None

def get_lsr_orbit(times, Potential=MWPotential2014):
    """Integrate the local standard of rest backwards in time, in order to provide
    a reference point to our XYZUVW coordinates.
    
    Parameters
    ----------
    times: numpy float array
        Times at which we want the orbit of the local standard of rest.
        
    Potential: galpy potential object.
    """
    ts = -(times/1e3)/bovy_conversion.time_in_Gyr(220.,8.)
    lsr_orbit= Orbit(vxvv=[1.,0,1,0,0.,0],vo=220,ro=8, solarmotion='schoenrich')
    lsr_orbit.integrate(ts,Potential)#,method='odeint')
    return lsr_orbit

def calc_delta_MWPotential2014(savefilename,plotfilename):
    Lmin, Lmax= 0.01, 10.
    if not os.path.exists(savefilename):
        #Setup grid
        nL, nE= 101,101
        Ls= 10.**numpy.linspace(numpy.log10(Lmin),numpy.log10(Lmax),nL)
        #Integration times
        ts= numpy.linspace(0.,20.,1001)
        deltas= numpy.empty((nL,nE))
        Einf= evalPot(10.**12.,0.,MWPotential2014)
        print Einf
        for ii in range(nL):
            #Calculate Ec
            Rc= rl(MWPotential2014,Ls[ii])
            print ii, "Rc = ", Rc*8.
            Ec= evalPot(Rc,0.,MWPotential2014)+0.5*Ls[ii]**2./Rc**2.
            Es= Ec+(Einf-Ec)*10.**numpy.linspace(-2.,0.,nE)
            for jj in range(nE):
                #Setup an orbit with this energy and angular momentum
                Er= 2.*(Es[jj]-Ec) #Random energy times 2 = vR^2 + vz^2
                vR= numpy.sqrt(4./5.*Er)
                vz= numpy.sqrt(1./5.*Er)
                o= Orbit([Rc,vR,Ls[ii]/Rc,0.,vz,0.])
                o.integrate(ts,MWPotential2014,method='symplec4_c')
                deltas[ii,jj]= estimateDeltaStaeckel(o.R(ts),o.z(ts),
                                                     pot=MWPotential2014)
        #Save
        save_pickles(savefilename,deltas)
    else:
        savefile= open(savefilename,'rb')
        deltas= pickle.load(savefile)
        savefile.close()
    #Plot
    print numpy.nanmax(deltas)
    bovy_plot.bovy_print()
    bovy_plot.bovy_dens2d(deltas.T,origin='lower',cmap='jet',
                          xrange=[numpy.log10(Lmin),numpy.log10(Lmax)],
                          yrange=[-2,0.],
                          xlabel=r'$\log_{10} L$',
                          ylabel=r'$\log_{10}\left(\frac{E-E_c(L)}{E(\infty)-E_c(L)}\right)$',
                          colorbar=True,shrink=0.78,
                          zlabel=r'$\mathrm{Approximate\ focal\ length}$',
#                          interpolation='nearest',
                          vmin=0.,vmax=0.6)
    bovy_plot.bovy_end_print(plotfilename)
    return None

def test_actionConservation():

    #_____initialize some KKSPot_____
    Delta = 1.0
    pot = KuzminKutuzovStaeckelPotential(ac=20.,Delta=Delta,normalize=True)

    #_____initialize an orbit (twice)_____
    vxvv = [1.,0.1,1.1,0.01,0.1]
    o= Orbit(vxvv=vxvv)

    #_____integrate the orbit with C_____
    ts= numpy.linspace(0,101,100)
    o.integrate(ts,pot,method='leapfrog_c')

    #_____Setup ActionAngle object and calculate actions (Staeckel approximation)_____
    aAS = actionAngleStaeckel(pot=pot,delta=Delta,c=True)
    jrs,lzs,jzs = aAS(o.R(ts),o.vR(ts),o.vT(ts),o.z(ts),o.vz(ts))

    assert numpy.all(numpy.fabs(jrs - jrs[0]) < 10.**-8.), \
        'Radial action is not conserved along orbit.'

    assert numpy.all(numpy.fabs(lzs - lzs[0]) < 10.**-8.), \
        'Angular momentum is not conserved along orbit.'

    assert numpy.all(numpy.fabs(jzs - jzs[0]) < 10.**-8.), \
        'Vertical action is not conserved along orbit.'

    return None

def trace_orbit_xyzuvw(xyzuvw_then, times):
    """
    Given a star's XYZUVW relative to the LSR (at any time), project its
    orbit forward to each of the times listed in *times*

    Parameters
    ----------
    xyzuvw : [pc,pc,pc,km/s,km/s,km/s]
    times : [ntimes] array
        Myr

    Returns
    -------
    xyzuvw_tf : [ntimes, 6] array
        the traced-forward positions and velocities
    """
    XYZUVW_sun_now = np.array([0.,0.,0.025,11.1,12.24,7.25])

    # convert positions to kpc
    xyzuvw_then[:3] *= 1e-3

    # convert times from Myr to bovy units
    bovy_times = times*1e-3 / bovy_conversion.time_in_Gyr(220., 8.)
    logging.debug("Tracing up to {} Myr".format(times[-1]))
    logging.debug("Tracing up to {} Bovy yrs".format(bovy_times[-1]))

    # Galpy coordinates are from the vantage point of the sun.
    # So even though the sun wasn't where it is,  we still "observe" the star
    # from a solar height and motion
    XYZUVW_gp_tf = xyzuvw_then - XYZUVW_sun_now
    logging.debug("Galpy vector: {}".format(XYZUVW_gp_tf))

    l,b,dist = lbdist_from_XYZ(XYZUVW_gp_tf)
    vxvv = [l,b,dist,XYZUVW_gp_tf[3],XYZUVW_gp_tf[4],XYZUVW_gp_tf[5]]
    logging.debug("vxvv: {}".format(vxvv))
    otf = Orbit(vxvv=vxvv, lb=True, uvw=True, solarmotion='schoenrich')

    otf.integrate(bovy_times,mp,method='odeint')
    data_tf = otf.getOrbit()
    XYZUVW_tf = galpy_coords_to_xyzuvw(data_tf, bovy_times)
    logging.debug("Started orbit at {}".format(XYZUVW_tf[0]))
    logging.debug("Finished orbit at {}".format(XYZUVW_tf[-1]))
    return XYZUVW_tf

def get_initial_condition(to):
    """Find an initial condition near apocenter that ends up today near the end of the stream and is close to a gyrfalcon stepsize
    to= Find an apocenter that is just more recent than this time"""
    vo, ro= 220., 8.
    # Center of the stream
    o= Orbit([0.10035165,-0.81302488,0.80986668,0.58024425,0.92753945,
               0.88763126],ro=ro,vo=vo)
    #Flip for backwards integration
    of= o.flip()
    ts= numpy.linspace(0.,to/bovy_conversion.time_in_Gyr(vo,ro),10001)
    of.integrate(ts,MWPotential2014)
    # Find the closest apocenter to the time to
    rs= numpy.sqrt(of.R(ts)**2.+of.z(ts)**2.)
    drs= rs-numpy.roll(rs,1)
    nearApo= (drs < 0.)*(numpy.roll(drs,1) > 0.)
    tsNearApo= ts[nearApo]
    tsNearApo*= bovy_conversion.time_in_Gyr(vo,ro)
    tfgf= numpy.amax(tsNearApo)
    #Round to nearest 2.**-8.
    tfgf= round(tfgf/2.**-8.)*2.**-8
    tf= tfgf/bovy_conversion.time_in_Gyr(vo,ro)
    print 'Current: %s,%s,%s,%s,%s,%s' % (of.x()[0], of.y()[0], of.z(), -of.vx()[0], -of.vy()[0], -of.vz())
    print 'At %g Gyr: %s,%s,%s,%s,%s,%s' % (tf*bovy_conversion.time_in_Gyr(vo,ro),of.x(tf)[0], of.y(tf)[0], of.z(tf), -of.vx(tf,use_physical=False)[0]*bovy_conversion.velocity_in_kpcGyr(vo,ro), -of.vy(tf,use_physical=False)[0]*bovy_conversion.velocity_in_kpcGyr(vo,ro), -of.vz(tf,use_physical=False)*bovy_conversion.velocity_in_kpcGyr(vo,ro))
    return None

def integrate_xyzuvw(params,ts,lsr_orbit,MWPotential2014):
    """Convenience function"""
    vxvv = params.copy()
    vxvv[2]=1.0/params[2]
    o = Orbit(vxvv=vxvv, radec=True, solarmotion='schoenrich')
    o.integrate(ts,MWPotential2014,method='odeint')
    xyzuvw = np.zeros( (len(ts),6) )
    xyzuvw[:,0] = 1e3*(o.x(ts)-lsr_orbit.x(ts))
    xyzuvw[:,1] = 1e3*(o.y(ts)-lsr_orbit.y(ts))
    xyzuvw[:,2] = 1e3*(o.z(ts))
    xyzuvw[:,3] = o.U(ts) - lsr_orbit.U(ts)
    xyzuvw[:,4] = o.V(ts) - lsr_orbit.V(ts)
    xyzuvw[:,5] = o.W(ts)
    return xyzuvw

def test_actionAngleTorus_orbit():
    from galpy.actionAngle import actionAngleTorus
    from galpy.potential import MWPotential2014
    from galpy.orbit import Orbit
    # Set up instance
    aAT= actionAngleTorus(pot=MWPotential2014,tol=10.**-5.)
    jr,jphi,jz= 0.05,1.1,0.025
    # First calculate frequencies and the initial RvR
    RvRom= aAT.xvFreqs(jr,jphi,jz,
                       numpy.array([0.]),
                       numpy.array([1.]),
                       numpy.array([2.]))
    om= RvRom[1:]
    # Angles along an orbit
    ts= numpy.linspace(0.,100.,1001)
    angler= ts*om[0]
    anglephi= 1.+ts*om[1]
    anglez= 2.+ts*om[2]
    # Calculate the orbit using actionAngleTorus
    RvR= aAT(jr,jphi,jz,angler,anglephi,anglez).T
    # Calculate the orbit using orbit integration
    orb= Orbit([RvRom[0][:,0],RvRom[0][:,1],RvRom[0][:,2],
                RvRom[0][:,3],RvRom[0][:,4],RvRom[0][:,5]])
    orb.integrate(ts,MWPotential2014)
    # Compare
    tol= -3.
    assert numpy.all(numpy.fabs(orb.R(ts)-RvR[0]) < 10.**tol), \
        'Integrated orbit does not agree with torus orbit in R'
    assert numpy.all(numpy.fabs(orb.vR(ts)-RvR[1]) < 10.**tol), \
        'Integrated orbit does not agree with torus orbit in vR'
    assert numpy.all(numpy.fabs(orb.vT(ts)-RvR[2]) < 10.**tol), \
        'Integrated orbit does not agree with torus orbit in vT'
    assert numpy.all(numpy.fabs(orb.z(ts)-RvR[3]) < 10.**tol), \
        'Integrated orbit does not agree with torus orbit in z'
    assert numpy.all(numpy.fabs(orb.vz(ts)-RvR[4]) < 10.**tol), \
        'Integrated orbit does not agree with torus orbit in vz'
    assert numpy.all(numpy.fabs(orb.phi(ts)-RvR[5]) < 10.**tol), \
        'Integrated orbit does not agree with torus orbit in phi'
    return None