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")
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