Beispiel #1
0
 def actionsFreqs(self,*args,**kwargs):
     """
     NAME:
        actionsFreqs
     PURPOSE:
        evaluate the actions and frequencies (jr,lz,jz,Omegar,Omegaphi,Omegaz)
     INPUT:
        Either:
           a) R,vR,vT,z,vz
           b) Orbit instance: initial condition used if that's it, orbit(t)
              if there is a time given as well
        scipy.integrate.quadrature keywords
     OUTPUT:
         (jr,lz,jz,Omegar,Omegaphi,Omegaz)
     HISTORY:
        2013-08-28 - Written - Bovy (IAS)
     """
     if ((self._c and not (kwargs.has_key('c') and not kwargs['c']))\
             or (ext_loaded and ((kwargs.has_key('c') and kwargs['c'])))) \
             and _check_c(self._pot):
         if len(args) == 5: #R,vR.vT, z, vz
             R,vR,vT, z, vz= args
         elif len(args) == 6: #R,vR.vT, z, vz, phi
             R,vR,vT, z, vz, phi= args
         else:
             meta= actionAngle(*args)
             R= meta._R
             vR= meta._vR
             vT= meta._vT
             z= meta._z
             vz= meta._vz
         if isinstance(R,float):
             R= nu.array([R])
             vR= nu.array([vR])
             vT= nu.array([vT])
             z= nu.array([z])
             vz= nu.array([vz])
         Lz= R*vT
         if self._useu0:
             #First calculate u0
             if kwargs.has_key('u0'):
                 u0= kwargs['u0']
             else:
                 E= nu.array([evaluatePotentials(R[ii],z[ii],self._pot) +vR[ii]**2./2.+vz[ii]**2./2.+vT[ii]**2./2. for ii in range(len(R))])
                 u0= actionAngleStaeckel_c.actionAngleStaeckel_calcu0(E,Lz,
                                                                      self._pot,
                                                                      self._delta)[0]
             if kwargs.has_key('u0'): kwargs.pop('u0')
         else:
             u0= None
         jr, jz, Omegar, Omegaphi, Omegaz, err= actionAngleStaeckel_c.actionAngleFreqStaeckel_c(\
             self._pot,self._delta,R,vR,vT,z,vz,u0=u0)
         if err == 0:
             return (jr,Lz,jz,Omegar,Omegaphi,Omegaz)
         else:
             raise RuntimeError("C-code for calculation actions failed; try with c=False")
     else:
         if kwargs.has_key('c') and kwargs['c'] and not self._c:
             warnings.warn("C module not used because potential does not have a C implementation",galpyWarning)
         raise NotImplementedError("actionsFreqs with c=False not implemented")
Beispiel #2
0
 def __call__(self,*args,**kwargs):
     """
     NAME:
        __call__
     PURPOSE:
        evaluate the actions (jr,lz,jz)
     INPUT:
        Either:
           a) R,vR,vT,z,vz
           b) Orbit instance: initial condition used if that's it, orbit(t)
              if there is a time given as well
         c= True/False; overrides the object's c= keyword to use C or not
        scipy.integrate.quadrature keywords
     OUTPUT:
        (jr,lz,jz)
     HISTORY:
        2012-11-27 - Written - Bovy (IAS)
     """
     if ((self._c and not (kwargs.has_key('c') and not kwargs['c']))\
             or (ext_loaded and ((kwargs.has_key('c') and kwargs['c'])))) \
             and _check_c(self._pot):
         if len(args) == 5: #R,vR.vT, z, vz
             R,vR,vT, z, vz= args
         elif len(args) == 6: #R,vR.vT, z, vz, phi
             R,vR,vT, z, vz, phi= args
         else:
             meta= actionAngle(*args)
             R= meta._R
             vR= meta._vR
             vT= meta._vT
             z= meta._z
             vz= meta._vz
         if isinstance(R,float):
             R= nu.array([R])
             vR= nu.array([vR])
             vT= nu.array([vT])
             z= nu.array([z])
             vz= nu.array([vz])
         Lz= R*vT
         if self._useu0:
             #First calculate u0
             if kwargs.has_key('u0'):
                 u0= kwargs['u0']
             else:
                 E= nu.array([evaluatePotentials(R[ii],z[ii],self._pot) +vR[ii]**2./2.+vz[ii]**2./2.+vT[ii]**2./2. for ii in range(len(R))])
                 u0= actionAngleStaeckel_c.actionAngleStaeckel_calcu0(E,Lz,
                                                                      self._pot,
                                                                      self._delta)[0]
             if kwargs.has_key('u0'): kwargs.pop('u0')
         else:
             u0= None
         jr, jz, err= actionAngleStaeckel_c.actionAngleStaeckel_c(\
             self._pot,self._delta,R,vR,vT,z,vz,u0=u0)
         if err == 0:
             return (jr,Lz,jz)
         else:
             raise RuntimeError("C-code for calculation actions failed; try with c=False")
     else:
         if kwargs.has_key('c') and kwargs['c'] and not self._c:
             warnings.warn("C module not used because potential does not have a C implementation",galpyWarning)
         if kwargs.has_key('c'): kwargs.pop('c')
         if (len(args) == 5 or len(args) == 6) \
                 and isinstance(args[0],nu.ndarray):
             ojr= nu.zeros((len(args[0])))
             olz= nu.zeros((len(args[0])))
             ojz= nu.zeros((len(args[0])))
             for ii in range(len(args[0])):
                 if len(args) == 5:
                     targs= (args[0][ii],args[1][ii],args[2][ii],
                             args[3][ii],args[4][ii])
                 elif len(args) == 6:
                     targs= (args[0][ii],args[1][ii],args[2][ii],
                             args[3][ii],args[4][ii],args[5][ii])
                 tjr,tlz,tjz= self(*targs,**copy.copy(kwargs))
                 ojr[ii]= tjr
                 ojz[ii]= tjz
                 olz[ii]= tlz
             return (ojr,olz,ojz)
         else:
             #Set up the actionAngleStaeckelSingle object
             aASingle= actionAngleStaeckelSingle(*args,pot=self._pot,
                                                  delta=self._delta)
             return (aASingle.JR(**copy.copy(kwargs)),
                     aASingle._R*aASingle._vT,
                     aASingle.Jz(**copy.copy(kwargs)))
    def __init__(self,pot=None,delta=None,Rmax=5.,
                 nE=25,npsi=25,nLz=25,numcores=1,
                 **kwargs):
        """
        NAME:
           __init__
        PURPOSE:
           initialize an actionAngleStaeckelGrid object
        INPUT:
           pot= potential or list of potentials

           delta= focus of prolate confocal coordinate system

           Rmax = Rmax for building grids

           nE=, npsi=, nLz= grid size

           numcores= number of cpus to use to parallellize

           +scipy.integrate.quad keywords
        OUTPUT:
        HISTORY:
            2012-11-29 - Written - Bovy (IAS)
        """
        if pot is None:
            raise IOError("Must specify pot= for actionAngleStaeckelGrid")
        self._pot= pot
        if delta is None:
            raise IOError("Must specify delta= for actionAngleStaeckelGrid")
        if ext_loaded and kwargs.has_key('c') and kwargs['c']:
            self._c= True
        else:
            self._c= False
        self._delta= delta
        self._Rmax= Rmax
        self._Rmin= 0.01
        #Set up the actionAngleStaeckel object that we will use to interpolate
        self._aA= actionAngleStaeckel.actionAngleStaeckel(pot=self._pot,delta=self._delta,c=self._c)
        #Build grid
        self._Lzmin= 0.01
        self._Lzs= numpy.linspace(self._Lzmin,
                                  self._Rmax\
                                      *galpy.potential.vcirc(self._pot,
                                                             self._Rmax),
                                  nLz)
        self._Lzmax= self._Lzs[-1]
        self._nLz= nLz
        #Calculate E_c(R=RL), energy of circular orbit
        self._RL= numpy.array([galpy.potential.rl(self._pot,l) for l in self._Lzs])
        self._RLInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs,
                                                                 self._RL,k=3)
        try:
            self._ERL= galpy.potential.evaluatePotentials(self._RL,numpy.zeros(self._nLz),self._pot) +self._Lzs**2./2./self._RL**2.
        except TypeError:
            self._ERL= numpy.array([galpy.potential.evaluatePotentials(self._RL[ii],0.,self._pot) +self._Lzs[ii]**2./2./self._RL[ii]**2. for ii in range(nLz)])
        self._ERLmax= numpy.amax(self._ERL)+1.
        self._ERLInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs,
                                                                  numpy.log(-(self._ERL-self._ERLmax)),k=3)
        self._Ramax= 200./8.
        try:
            self._ERa= galpy.potential.evaluatePotentials(self._Ramax,0.,self._pot) +self._Lzs**2./2./self._Ramax**2.
        except TypeError:
            self._ERa= numpy.array([galpy.potential.evaluatePotentials(self._Ramax,0.,self._pot) +self._Lzs[ii]**2./2./self._Ramax**2. for ii in range(nLz)])
        #self._EEsc= numpy.array([self._ERL[ii]+galpy.potential.vesc(self._pot,self._RL[ii])**2./4. for ii in range(nLz)])
        self._ERamax= numpy.amax(self._ERa)+1.
        self._ERaInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs,
                                                                  numpy.log(-(self._ERa-self._ERamax)),k=3)
        y= numpy.linspace(0.,1.,nE)
        self._nE= nE
        psis= numpy.linspace(0.,1.,npsi)*numpy.pi/2.
        self._npsi= npsi
        jr= numpy.zeros((nLz,nE,npsi))
        jz= numpy.zeros((nLz,nE,npsi))
        u0= numpy.zeros((nLz,nE))
        jrLzE= numpy.zeros((nLz))
        jzLzE= numpy.zeros((nLz))
        #First calculate u0
        thisLzs= (numpy.tile(self._Lzs,(nE,1)).T).flatten()
        thisERL= (numpy.tile(self._ERL,(nE,1)).T).flatten()
        thisERa= (numpy.tile(self._ERa,(nE,1)).T).flatten()
        thisy= (numpy.tile(y,(nLz,1))).flatten()
        thisE= _invEfunc(_Efunc(thisERa,thisERL)+thisy*(_Efunc(thisERL,thisERL)-_Efunc(thisERa,thisERL)),thisERL)
        if isinstance(self._pot,galpy.potential.interpRZPotential) and hasattr(self._pot,'_origPot'):
            u0pot= self._pot._origPot
        else:
            u0pot= self._pot
        if self._c:
            mu0= actionAngleStaeckel_c.actionAngleStaeckel_calcu0(thisE,thisLzs,
                                                                  u0pot,
                                                                  self._delta)[0]
        else:
            if numcores > 1:
                mu0= multi.parallel_map((lambda x: self.calcu0(thisE[x],
                                                               thisLzs[x])),
                                        range(nE*nLz),
                                        numcores=numcores)
            else:
                mu0= map((lambda x: self.calcu0(thisE[x],
                                                thisLzs[x])),
                         range(nE*nLz))
        u0= numpy.reshape(mu0,(nLz,nE))
        thisR= self._delta*numpy.sinh(u0)
        thisv= numpy.reshape(self.vatu0(thisE.flatten(),thisLzs.flatten(),
                                        u0.flatten(),
                                        thisR.flatten()),(nLz,nE))
        self.thisv= thisv
        #reshape
        thisLzs= numpy.reshape(thisLzs,(nLz,nE))
        thispsi= numpy.tile(psis,(nLz,nE,1)).flatten()
        thisLzs= numpy.tile(thisLzs.T,(npsi,1,1)).T.flatten()
        thisR= numpy.tile(thisR.T,(npsi,1,1)).T.flatten()
        thisv= numpy.tile(thisv.T,(npsi,1,1)).T.flatten()
        mjr, mlz, mjz= self._aA(thisR, #R
                                thisv*numpy.cos(thispsi), #vR
                                thisLzs/thisR, #vT
                                numpy.zeros(len(thisR)), #z
                                thisv*numpy.sin(thispsi), #vz
                                fixed_quad=True) 
        if isinstance(self._pot,galpy.potential.interpRZPotential) and hasattr(self._pot,'_origPot'):
            #Interpolated potentials have problems with extreme orbits
            indx= (mjr == 9999.99)
            indx+= (mjz == 9999.99)
            #Re-calculate these using the original potential, hopefully not too slow
            tmpaA= actionAngleStaeckel.actionAngleStaeckel(pot=self._pot._origPot,delta=self._delta,c=self._c)
            mjr[indx], dum, mjz[indx]= tmpaA(thisR[indx], #R
                                             thisv[indx]*numpy.cos(thispsi[indx]), #vR
                                             thisLzs[indx]/thisR[indx], #vT
                                             numpy.zeros(numpy.sum(indx)), #z
                                             thisv[indx]*numpy.sin(thispsi[indx]), #vz
                                             fixed_quad=True)
        jr= numpy.reshape(mjr,(nLz,nE,npsi))
        jz= numpy.reshape(mjz,(nLz,nE,npsi))
        for ii in range(nLz):
                jrLzE[ii]= numpy.amax(jr[ii,(jr[ii,:,:] != 9999.99)])#:,:])
                jzLzE[ii]= numpy.amax(jz[ii,(jz[ii,:,:] != 9999.99)])#:,:])
        jrLzE[(jrLzE == 0.)]= numpy.amin(jrLzE[(jrLzE > 0.)])
        jzLzE[(jzLzE == 0.)]= numpy.amin(jzLzE[(jzLzE > 0.)])
        for ii in range(nLz):
            jr[ii,:,:]/= jrLzE[ii]
            jz[ii,:,:]/= jzLzE[ii]
        #Deal w/ 9999.99
        jr[(jr > 1.)]= 1.
        jz[(jz > 1.)]= 1.
        #First interpolate the maxima
        self._jr= jr
        self._jz= jz
        self._u0= u0
        self._jrLzE= jrLzE
        self._jzLzE= jzLzE
        self._jrLzInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs,
                                                                   numpy.log(jrLzE+10.**-5.),k=3)
        self._jzLzInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs,
                                                                   numpy.log(jzLzE+10.**-5.),k=3)
        #Interpolate u0
        self._logu0Interp= interpolate.RectBivariateSpline(self._Lzs,
                                                           y,
                                                           numpy.log(u0),
                                                           kx=3,ky=3,s=0.)
        #spline filter jr and jz, such that they can be used with ndimage.map_coordinates
        self._jrFiltered= ndimage.spline_filter(numpy.log(self._jr+10.**-10.),order=3)
        self._jzFiltered= ndimage.spline_filter(numpy.log(self._jz+10.**-10.),order=3)
        return None