def Jz(self,*args,**kwargs): """ NAME: Jz PURPOSE: evaluate the action 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 scipy.integrate.quadrature keywords OUTPUT: jz HISTORY: 2012-07-30 - Written - Bovy (IAS@MPIA) """ self._parse_eval_args(*args) Phi= _evaluatePotentials(self._pot,self._eval_R,self._eval_z) Phio= _evaluatePotentials(self._pot,self._eval_R,0.) Ez= Phi-Phio+self._eval_vz**2./2. #Bigger than Ezzmax? thisEzZmax= numpy.exp(self._EzZmaxsInterp(self._eval_R)) if self._eval_R > self._Rmax or self._eval_R < self._Rmin or (Ez != 0. and numpy.log(Ez) > thisEzZmax): #Outside of the grid if _PRINTOUTSIDEGRID: #pragma: no cover print("Outside of grid in Ez") jz= self._aA(self._eval_R,0.,1.,#these two r dummies 0.,math.sqrt(2.*Ez), _justjz=True, **kwargs)[2] else: jz= (self._jzInterp(self._eval_R,Ez/thisEzZmax)\ *(numpy.exp(self._jzEzmaxInterp(self._eval_R))-10.**-5.))[0][0] return jz
def Jz(self,*args,**kwargs): """ NAME: Jz PURPOSE: evaluate the action 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 scipy.integrate.quadrature keywords OUTPUT: jz HISTORY: 2012-07-30 - Written - Bovy (IAS@MPIA) """ self._parse_eval_args(*args) Phi= _evaluatePotentials(self._pot,self._eval_R,self._eval_z) Phio= _evaluatePotentials(self._pot,self._eval_R,0.) Ez= Phi-Phio+self._eval_vz**2./2. #Bigger than Ezzmax? thisEzZmax= numpy.exp(self._EzZmaxsInterp(self._eval_R)) if self._eval_R > self._Rmax or self._eval_R < self._Rmin or (Ez != 0. and numpy.log(Ez) > thisEzZmax): #Outside of the grid if _PRINTOUTSIDEGRID: #pragma: no cover print("Outside of grid in Ez") jz= self._aA(self._eval_R,0.,1.,#these two r dummies 0.,math.sqrt(2.*Ez), _justjz=True, **kwargs)[2] else: jz= (self._jzInterp(self._eval_R,Ez/thisEzZmax)\ *(numpy.exp(self._jzEzmaxInterp(self._eval_R))-10.**-5.))[0][0] return jz
def calcELStaeckel(R,vR,vT,z,vz,pot,vc=1.,ro=1.): """ NAME: calcELStaeckel PURPOSE: calculate the energy and angular momentum INPUT: R - Galactocentric radius (/ro) vR - radial part of the velocity (/vc) vT - azimuthal part of the velocity (/vc) vc - circular velocity ro - reference radius OUTPUT: (E,L) HISTORY: 2012-11-30 - Written - Bovy (IAS) """ return (_evaluatePotentials(pot,R,z)+vR**2./2.+vT**2./2.+vz**2./2.,R*vT)
def potentialStaeckel(u,v,pot,delta): """ NAME: potentialStaeckel PURPOSE: return the potential INPUT: u - confocal u v - confocal v pot - potential delta - focus OUTPUT: Phi(u,v) HISTORY: 2012-11-29 - Written - Bovy (IAS) """ R,z= bovy_coords.uv_to_Rz(u,v,delta=delta) return _evaluatePotentials(pot,R,z)
def potentialStaeckel(u, v, pot, delta): """ NAME: potentialStaeckel PURPOSE: return the potential INPUT: u - confocal u v - confocal v pot - potential delta - focus OUTPUT: Phi(u,v) HISTORY: 2012-11-29 - Written - Bovy (IAS) """ R, z = bovy_coords.uv_to_Rz(u, v, delta=delta) return _evaluatePotentials(pot, R, z)
def calcELStaeckel(R,vR,vT,z,vz,pot,vc=1.,ro=1.): """ NAME: calcELStaeckel PURPOSE: calculate the energy and angular momentum INPUT: R - Galactocentric radius (/ro) vR - radial part of the velocity (/vc) vT - azimuthal part of the velocity (/vc) vc - circular velocity ro - reference radius OUTPUT: (E,L) HISTORY: 2012-11-30 - Written - Bovy (IAS) """ return (_evaluatePotentials(pot,R,z)+vR**2./2.+vT**2./2.+vz**2./2.,R*vT)
def _evaluate(self,*args,**kwargs): """ NAME: _evaluate 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 scipy.integrate.quadrature keywords OUTPUT: (jr,lz,jz) HISTORY: 2012-07-27 - Written - Bovy (IAS@MPIA) NOTE: For a Miyamoto-Nagai potential, this seems accurate to 0.1% and takes ~0.13 ms For a MWPotential, this takes ~ 0.17 ms """ 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: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_vz #First work on the vertical action Phi= _evaluatePotentials(self._pot,R,z) try: Phio= _evaluatePotentials(self._pot,R,numpy.zeros(len(R))) except TypeError: Phio= _evaluatePotentials(self._pot,R,0.) Ez= Phi-Phio+vz**2./2. #Bigger than Ezzmax? thisEzZmax= numpy.exp(self._EzZmaxsInterp(R)) if isinstance(R,numpy.ndarray): indx= (R > self._Rmax) indx+= (R < self._Rmin) indx+= (Ez != 0.)*(numpy.log(Ez) > thisEzZmax) indxc= True^indx jz= numpy.empty(R.shape) if numpy.sum(indxc) > 0: jz[indxc]= (self._jzInterp.ev(R[indxc],Ez[indxc]/thisEzZmax[indxc])\ *(numpy.exp(self._jzEzmaxInterp(R[indxc]))-10.**-5.)) if numpy.sum(indx) > 0: jz[indx]= self._aA(R[indx], numpy.zeros(numpy.sum(indx)), numpy.ones(numpy.sum(indx)),#these two r dummies numpy.zeros(numpy.sum(indx)), numpy.sqrt(2.*Ez[indx]), _justjz=True, **kwargs)[2] else: if R > self._Rmax or R < self._Rmin or (Ez != 0 and numpy.log(Ez) > thisEzZmax): #Outside of the grid if _PRINTOUTSIDEGRID: #pragma: no cover print("Outside of grid in Ez", R > self._Rmax , R < self._Rmin , (Ez != 0 and numpy.log(Ez) > thisEzZmax)) jz= self._aA(R,0.,1.,#these two r dummies 0.,math.sqrt(2.*Ez), _justjz=True, **kwargs)[2] else: jz= (self._jzInterp(R,Ez/thisEzZmax)\ *(numpy.exp(self._jzEzmaxInterp(R))-10.**-5.))[0][0] #Radial action ERLz= numpy.fabs(R*vT)+self._gamma*jz ER= Phio+vR**2./2.+ERLz**2./2./R**2. thisRL= self._RLInterp(ERLz) thisERRL= -numpy.exp(self._ERRLInterp(ERLz))+self._ERRLmax thisERRa= -numpy.exp(self._ERRaInterp(ERLz))+self._ERRamax if isinstance(R,numpy.ndarray): indx= ((ER-thisERRa)/(thisERRL-thisERRa) > 1.)\ *(((ER-thisERRa)/(thisERRL-thisERRa)-1.) < 10.**-2.) ER[indx]= thisERRL[indx] indx= ((ER-thisERRa)/(thisERRL-thisERRa) < 0.)\ *((ER-thisERRa)/(thisERRL-thisERRa) > -10.**-2.) ER[indx]= thisERRa[indx] indx= (ERLz < self._Lzmin) indx+= (ERLz > self._Lzmax) indx+= ((ER-thisERRa)/(thisERRL-thisERRa) > 1.) indx+= ((ER-thisERRa)/(thisERRL-thisERRa) < 0.) indxc= True^indx jr= numpy.empty(R.shape) if numpy.sum(indxc) > 0: jr[indxc]= (self._jrInterp.ev(ERLz[indxc], (ER[indxc]-thisERRa[indxc])/(thisERRL[indxc]-thisERRa[indxc]))\ *(numpy.exp(self._jrERRaInterp(ERLz[indxc]))-10.**-5.)) if numpy.sum(indx) > 0: jr[indx]= self._aA(thisRL[indx], numpy.sqrt(2.*(ER[indx]-_evaluatePotentials(self._pot,thisRL[indx],0.))-ERLz[indx]**2./thisRL[indx]**2.), ERLz[indx]/thisRL[indx], numpy.zeros(len(thisRL)), numpy.zeros(len(thisRL)), _justjr=True, **kwargs)[0] else: if (ER-thisERRa)/(thisERRL-thisERRa) > 1. \ and ((ER-thisERRa)/(thisERRL-thisERRa)-1.) < 10.**-2.: ER= thisERRL elif (ER-thisERRa)/(thisERRL-thisERRa) < 0. \ and (ER-thisERRa)/(thisERRL-thisERRa) > -10.**-2.: ER= thisERRa #Outside of grid? if ERLz < self._Lzmin or ERLz > self._Lzmax \ or (ER-thisERRa)/(thisERRL-thisERRa) > 1. \ or (ER-thisERRa)/(thisERRL-thisERRa) < 0.: if _PRINTOUTSIDEGRID: #pragma: no cover print("Outside of grid in ER/Lz", ERLz < self._Lzmin , ERLz > self._Lzmax \ , (ER-thisERRa)/(thisERRL-thisERRa) > 1. \ , (ER-thisERRa)/(thisERRL-thisERRa) < 0., ER, thisERRL, thisERRa, (ER-thisERRa)/(thisERRL-thisERRa)) jr= self._aA(thisRL[0], numpy.sqrt(2.*(ER-_evaluatePotentials(self._pot,thisRL,0.))-ERLz**2./thisRL**2.)[0], (ERLz/thisRL)[0], 0.,0., _justjr=True, **kwargs)[0] else: jr= (self._jrInterp(ERLz, (ER-thisERRa)/(thisERRL-thisERRa))\ *(numpy.exp(self._jrERRaInterp(ERLz))-10.**-5.))[0][0] return (jr,R*vT,jz)
def _evaluate(self, *args, **kwargs): """ NAME: __call__ (_evaluate) PURPOSE: evaluate the actions (jr,lz,jz) INPUT: Either: a) R,vR,vT,z,vz[,phi]: 1) floats: phase-space value for single object (phi is optional) (each can be a Quantity) 2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity) b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument delta= (object-wide default) can be used to override the object-wide focal length; can also be an array with length N to allow different delta for different phase-space points u0= (None) if object-wide option useu0 is set, u0 to use (if useu0 and useu0 is None, a good value will be computed) c= (object-wide default, bool) True/False to override the object-wide setting for whether or not to use the C implementation order= (object-wide default, int) number of points to use in the Gauss-Legendre numerical integration of the relevant action integrals When not using C: fixed_quad= (False) if True, use Gaussian quadrature (scipy.integrate.fixed_quad instead of scipy.integrate.quad) scipy.integrate.fixed_quad or .quad keywords OUTPUT: (jr,lz,jz) HISTORY: 2012-11-27 - Written - Bovy (IAS) 2017-12-27 - Allowed individual delta for each point - Bovy (UofT) """ delta = kwargs.pop('delta', self._delta) order = kwargs.get('order', self._order) if ((self._c and not ('c' in kwargs and not kwargs['c']))\ or (ext_loaded and (('c' in kwargs 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: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_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 'u0' in kwargs: u0 = nu.asarray(kwargs['u0']) else: E = nu.array([ _evaluatePotentials(self._pot, R[ii], z[ii]) + 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,delta)[0] kwargs.pop('u0', None) else: u0 = None jr, jz, err= actionAngleStaeckel_c.actionAngleStaeckel_c(\ self._pot,delta,R,vR,vT,z,vz,u0=u0,order=order) if err == 0: return (jr, Lz, jz) else: #pragma: no cover raise RuntimeError( "C-code for calculation actions failed; try with c=False") else: if 'c' in kwargs and kwargs['c'] and not self._c: #pragma: no cover warnings.warn( "C module not used because potential does not have a C implementation", galpyWarning) kwargs.pop('c', None) 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]) tkwargs = copy.copy(kwargs) try: tkwargs['delta'] = delta[ii] except TypeError: tkwargs['delta'] = delta tjr, tlz, tjz = self(*targs, **tkwargs) 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=delta) return (aASingle.JR(**copy.copy(kwargs)), aASingle._R * aASingle._vT, aASingle.Jz(**copy.copy(kwargs)))
def _uminumaxvmin(self, *args, **kwargs): """ NAME: _uminumaxvmin PURPOSE: evaluate u_min, u_max, and v_min 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 OUTPUT: (umin,umax,vmin) HISTORY: 2017-12-12 - Written - Bovy (UofT) """ delta = kwargs.pop('delta', self._delta) if ((self._c and not ('c' in kwargs and not kwargs['c']))\ or (ext_loaded and (('c' in kwargs 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: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_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 'u0' in kwargs: u0 = nu.asarray(kwargs['u0']) else: E = nu.array([ _evaluatePotentials(self._pot, R[ii], z[ii]) + 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,delta)[0] kwargs.pop('u0', None) else: u0 = None umin, umax, vmin, err= \ actionAngleStaeckel_c.actionAngleUminUmaxVminStaeckel_c(\ self._pot,delta,R,vR,vT,z,vz,u0=u0) if err == 0: return (umin, umax, vmin) else: #pragma: no cover raise RuntimeError( "C-code for calculation actions failed; try with c=False") else: if 'c' in kwargs and kwargs['c'] and not self._c: #pragma: no cover warnings.warn( "C module not used because potential does not have a C implementation", galpyWarning) kwargs.pop('c', None) if (len(args) == 5 or len(args) == 6) \ and isinstance(args[0],nu.ndarray): oumin = nu.zeros((len(args[0]))) oumax = nu.zeros((len(args[0]))) ovmin = 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]) tkwargs = copy.copy(kwargs) try: tkwargs['delta'] = delta[ii] except TypeError: tkwargs['delta'] = delta tumin,tumax,tvmin= self._uminumaxvmin(\ *targs,**tkwargs) oumin[ii] = tumin oumax[ii] = tumax ovmin[ii] = tvmin return (oumin, oumax, ovmin) else: #Set up the actionAngleStaeckelSingle object aASingle = actionAngleStaeckelSingle(*args, pot=self._pot, delta=delta) umin, umax = aASingle.calcUminUmax() vmin = aASingle.calcVmin() return (umin, umax, vmin)
def _actionsFreqs(self,*args,**kwargs): """ NAME: actionsFreqs (_actionsFreqs) PURPOSE: evaluate the actions and frequencies (jr,lz,jz,Omegar,Omegaphi,Omegaz) INPUT: Either: a) R,vR,vT,z,vz[,phi]: 1) floats: phase-space value for single object (phi is optional) (each can be a Quantity) 2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity) b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument fixed_quad= (False) if True, use n=10 fixed_quad integration scipy.integrate.quadrature or .fixed_quad keywords OUTPUT: (jr,lz,jz,Omegar,Omegaphi,Omegaz) HISTORY: 2013-12-28 - Written - Bovy (IAS) """ fixed_quad= kwargs.pop('fixed_quad',False) 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: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_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]) if self._c: #pragma: no cover pass else: Lz= R*vT Lx= -z*vT Ly= z*vR-R*vz L2= Lx*Lx+Ly*Ly+Lz*Lz E= _evaluatePotentials(self._pot,R,z)+vR**2./2.+vT**2./2.+vz**2./2. L= nu.sqrt(L2) #Actions Jphi= Lz Jz= L-nu.fabs(Lz) #Jr requires some more work #Set up an actionAngleAxi object for EL and rap/rperi calculations axiR= nu.sqrt(R**2.+z**2.) axivT= L/axiR axivR= (R*vR+z*vz)/axiR Jr= [] Or= [] Op= [] for ii in range(len(axiR)): axiaA= actionAngleAxi(axiR[ii],axivR[ii],axivT[ii], pot=self._2dpot) (rperi,rap)= axiaA.calcRapRperi() EL= axiaA.calcEL() E, L= EL Jr.append(self._calc_jr(rperi,rap,E,L,fixed_quad,**kwargs)) #Radial period if Jr[-1] < 10.**-9.: #Circular orbit Or.append(epifreq(self._pot,axiR[ii],use_physical=False)) Op.append(omegac(self._pot,axiR[ii],use_physical=False)) continue Rmean= m.exp((m.log(rperi)+m.log(rap))/2.) Or.append(self._calc_or(Rmean,rperi,rap,E,L,fixed_quad,**kwargs)) Op.append(self._calc_op(Or[-1],Rmean,rperi,rap,E,L,fixed_quad,**kwargs)) Op= nu.array(Op) Oz= copy.copy(Op) Op[vT < 0.]*= -1. return (nu.array(Jr),Jphi,Jz,nu.array(Or),Op,Oz)
def _actionsFreqsAngles(self,*args,**kwargs): """ NAME: _actionsFreqsAngles PURPOSE: evaluate the actions, frequencies, and angles (jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez) INPUT: Either: a) R,vR,vT,z,vz,phi (MUST HAVE PHI) 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,angler,anglephi,anglez) HISTORY: 2013-08-28 - Written - Bovy (IAS) """ if ((self._c and not ('c' in kwargs and not kwargs['c']))\ or (ext_loaded and (('c' in kwargs and kwargs['c'])))) \ and _check_c(self._pot): if len(args) == 5: #R,vR.vT, z, vz pragma: no cover raise IOError("Must specify phi") elif len(args) == 6: #R,vR.vT, z, vz, phi R,vR,vT, z, vz, phi= args else: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_vz phi= self._eval_phi if isinstance(R,float): R= nu.array([R]) vR= nu.array([vR]) vT= nu.array([vT]) z= nu.array([z]) vz= nu.array([vz]) phi= nu.array([phi]) Lz= R*vT if self._useu0: #First calculate u0 if 'u0' in kwargs: u0= nu.asarray(kwargs['u0']) else: E= nu.array([_evaluatePotentials(self._pot,R[ii],z[ii]) +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] kwargs.pop('u0',None) else: u0= None jr, jz, Omegar, Omegaphi, Omegaz, angler, anglephi,anglez, err= actionAngleStaeckel_c.actionAngleFreqAngleStaeckel_c(\ self._pot,self._delta,R,vR,vT,z,vz,phi,u0=u0) # Adjustements for close-to-circular orbits indx= nu.isnan(Omegar)*(jr < 10.**-3.)+nu.isnan(Omegaz)*(jz < 10.**-3.) #Close-to-circular and close-to-the-plane orbits if nu.sum(indx) > 0: Omegar[indx]= [epifreq(self._pot,r,use_physical=False) for r in R[indx]] Omegaphi[indx]= [omegac(self._pot,r,use_physical=False) for r in R[indx]] Omegaz[indx]= [verticalfreq(self._pot,r,use_physical=False) for r in R[indx]] if err == 0: return (jr,Lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez) else: raise RuntimeError("C-code for calculation actions failed; try with c=False") #pragma: no cover else: #pragma: no cover if 'c' in kwargs and kwargs['c'] and not self._c: #pragma: no cover warnings.warn("C module not used because potential does not have a C implementation",galpyWarning) raise NotImplementedError("actionsFreqs with c=False not implemented")
def _EccZmaxRperiRap(self,*args,**kwargs): """ NAME: EccZmaxRperiRap (_EccZmaxRperiRap) PURPOSE: evaluate the eccentricity, maximum height above the plane, peri- and apocenter in the Staeckel approximation INPUT: Either: a) R,vR,vT,z,vz[,phi]: 1) floats: phase-space value for single object (phi is optional) (each can be a Quantity) 2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity) b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument OUTPUT: (e,zmax,rperi,rap) HISTORY: 2017-12-15 - Written - Bovy (UofT) """ 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: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_vz Lz= R*vT Phi= _evaluatePotentials(self._pot,R,z) E= Phi+vR**2./2.+vT**2./2.+vz**2./2. thisERL= -numpy.exp(self._ERLInterp(Lz))+self._ERLmax thisERa= -numpy.exp(self._ERaInterp(Lz))+self._ERamax if isinstance(R,numpy.ndarray): indx= ((E-thisERa)/(thisERL-thisERa) > 1.)\ *(((E-thisERa)/(thisERL-thisERa)-1.) < 10.**-2.) E[indx]= thisERL[indx] indx= ((E-thisERa)/(thisERL-thisERa) < 0.)\ *((E-thisERa)/(thisERL-thisERa) > -10.**-2.) E[indx]= thisERa[indx] indx= (Lz < self._Lzmin) indx+= (Lz > self._Lzmax) indx+= ((E-thisERa)/(thisERL-thisERa) > 1.) indx+= ((E-thisERa)/(thisERL-thisERa) < 0.) indxc= True^indx ecc= numpy.empty(R.shape) zmax= numpy.empty(R.shape) rperi= numpy.empty(R.shape) rap= numpy.empty(R.shape) if numpy.sum(indxc) > 0: u0= numpy.exp(self._logu0Interp.ev(Lz[indxc], (_Efunc(E[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc]))/(_Efunc(thisERL[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc])))) sinh2u0= numpy.sinh(u0)**2. thisEr= self.Er(R[indxc],z[indxc],vR[indxc],vz[indxc], E[indxc],Lz[indxc],sinh2u0,u0) thisEz= self.Ez(R[indxc],z[indxc],vR[indxc],vz[indxc], E[indxc],Lz[indxc],sinh2u0,u0) thisv2= self.vatu0(E[indxc],Lz[indxc],u0,self._delta*numpy.sinh(u0),retv2=True) cos2psi= 2.*thisEr/thisv2/(1.+sinh2u0) #latter is cosh2u0 cos2psi[(cos2psi > 1.)*(cos2psi < 1.+10.**-5.)]= 1. indxCos2psi= (cos2psi > 1.) indxCos2psi+= (cos2psi < 0.) indxc[indxc]= True^indxCos2psi#Handle these two cases as off-grid indx= True^indxc psi= numpy.arccos(numpy.sqrt(cos2psi[True^indxCos2psi])) coords= numpy.empty((3,numpy.sum(indxc))) coords[0,:]= (Lz[indxc]-self._Lzmin)/(self._Lzmax-self._Lzmin)*(self._nLz-1.) y= (_Efunc(E[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc]))/(_Efunc(thisERL[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc])) coords[1,:]= y*(self._nE-1.) coords[2,:]= psi/numpy.pi*2.*(self._npsi-1.) ecc[indxc]= (numpy.exp(ndimage.interpolation.map_coordinates(self._eccFiltered, coords, order=3, prefilter=False))-10.**-10.) rperi[indxc]= (numpy.exp(ndimage.interpolation.map_coordinates(self._rperiFiltered, coords, order=3, prefilter=False))-10.**-10.)*(numpy.exp(self._rperiLzInterp(Lz[indxc]))-10.**-5.) # We do rap below with zmax #Switch to Ez-calculated psi sin2psi= 2.*thisEz[True^indxCos2psi]/thisv2[True^indxCos2psi]/(1.+sinh2u0[True^indxCos2psi]) #latter is cosh2u0 sin2psi[(sin2psi > 1.)*(sin2psi < 1.+10.**-5.)]= 1. indxSin2psi= (sin2psi > 1.) indxSin2psi+= (sin2psi < 0.) indxc[indxc]= True^indxSin2psi#Handle these two cases as off-grid indx= True^indxc psiz= numpy.arcsin(numpy.sqrt(sin2psi[True^indxSin2psi])) newcoords= numpy.empty((3,numpy.sum(indxc))) newcoords[0:2,:]= coords[0:2,True^indxSin2psi] newcoords[2,:]= psiz/numpy.pi*2.*(self._npsi-1.) zmax[indxc]= (numpy.exp(ndimage.interpolation.map_coordinates(self._zmaxFiltered, newcoords, order=3, prefilter=False))-10.**-10.)*(numpy.exp(self._zmaxLzInterp(Lz[indxc]))-10.**-5.) rap[indxc]= (numpy.exp(ndimage.interpolation.map_coordinates(self._rapFiltered, newcoords, order=3, prefilter=False))-10.**-10.)*(numpy.exp(self._rapLzInterp(Lz[indxc]))-10.**-5.) if numpy.sum(indx) > 0: eccindiv, zmaxindiv, rperiindiv, rapindiv=\ self._aA.EccZmaxRperiRap(R[indx], vR[indx], vT[indx], z[indx], vz[indx], **kwargs) ecc[indx]= eccindiv zmax[indx]= zmaxindiv rperi[indx]= rperiindiv rap[indx]= rapindiv else: ecc,zmax,rperi,rap= self.EccZmaxRperiRap(numpy.array([R]), numpy.array([vR]), numpy.array([vT]), numpy.array([z]), numpy.array([vz]), **kwargs) return (ecc[0],zmax[0],rperi[0],rap[0]) ecc[ecc < 0.]= 0. zmax[zmax < 0.]= 0. rperi[rperi < 0.]= 0. rap[rap < 0.]= 0. return (ecc,zmax,rperi,rap)
def _evaluate(self, *args, **kwargs): """ NAME: _evaluate 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 fixed_quad= (False) if True, use n=10 fixed_quad integration scipy.integrate.quadrature keywords OUTPUT: (jr,lz,jz) HISTORY: 2013-12-28 - Written - Bovy (IAS) """ fixed_quad = kwargs.pop('fixed_quad', False) 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: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_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]) if self._c: #pragma: no cover pass else: Lz = R * vT Lx = -z * vT Ly = z * vR - R * vz L2 = Lx * Lx + Ly * Ly + Lz * Lz E= _evaluatePotentials(self._pot,R,z)\ +vR**2./2.+vT**2./2.+vz**2./2. L = nu.sqrt(L2) #Actions Jphi = Lz Jz = L - nu.fabs(Lz) #Jr requires some more work #Set up an actionAngleAxi object for EL and rap/rperi calculations axiR = nu.sqrt(R**2. + z**2.) axivT = L / axiR axivR = (R * vR + z * vz) / axiR Jr = [] for ii in range(len(axiR)): axiaA = actionAngleAxi(axiR[ii], axivR[ii], axivT[ii], pot=self._2dpot) (rperi, rap) = axiaA.calcRapRperi() EL = axiaA.calcEL() E, L = EL Jr.append(self._calc_jr(rperi, rap, E, L, fixed_quad, **kwargs)) return (nu.array(Jr), Jphi, Jz)
def _actionsFreqsAngles(self, *args, **kwargs): """ NAME: _actionsFreqsAngles PURPOSE: evaluate the actions, frequencies, and angles (jr,lz,jz,Omegar,Omegaphi,Omegaz,ar,ap,az) 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 fixed_quad= (False) if True, use n=10 fixed_quad integration scipy.integrate.quadrature keywords OUTPUT: (jr,lz,jz,Omegar,Omegaphi,Omegaz,ar,aphi,az) HISTORY: 2013-12-29 - Written - Bovy (IAS) """ fixed_quad = kwargs.pop('fixed_quad', False) if len(args) == 5: #R,vR.vT, z, vz pragma: no cover raise IOError("You need to provide phi when calculating angles") elif len(args) == 6: #R,vR.vT, z, vz, phi R, vR, vT, z, vz, phi = args else: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_vz phi = self._eval_phi if isinstance(R, float): R = nu.array([R]) vR = nu.array([vR]) vT = nu.array([vT]) z = nu.array([z]) vz = nu.array([vz]) phi = nu.array([phi]) if self._c: #pragma: no cover pass else: Lz = R * vT Lx = -z * vT Ly = z * vR - R * vz L2 = Lx * Lx + Ly * Ly + Lz * Lz E = _evaluatePotentials( self._pot, R, z) + vR**2. / 2. + vT**2. / 2. + vz**2. / 2. L = nu.sqrt(L2) #Actions Jphi = Lz Jz = L - nu.fabs(Lz) #Jr requires some more work #Set up an actionAngleAxi object for EL and rap/rperi calculations axiR = nu.sqrt(R**2. + z**2.) axivT = L / axiR #these are really spherical coords, called axi bc they go in actionAngleAxi axivR = (R * vR + z * vz) / axiR axivz = (z * vR - R * vz) / axiR Jr = [] Or = [] Op = [] ar = [] az = [] #Calculate the longitude of the ascending node asc = self._calc_long_asc(z, R, axivz, phi, Lz, L) for ii in range(len(axiR)): axiaA = actionAngleAxi(axiR[ii], axivR[ii], axivT[ii], pot=self._2dpot) (rperi, rap) = axiaA.calcRapRperi() EL = axiaA.calcEL() E, L = EL Jr.append(self._calc_jr(rperi, rap, E, L, fixed_quad, **kwargs)) #Radial period Rmean = m.exp((m.log(rperi) + m.log(rap)) / 2.) if Jr[-1] < 10.**-9.: #Circular orbit Or.append(epifreq(self._pot, axiR[ii], use_physical=False)) Op.append(omegac(self._pot, axiR[ii], use_physical=False)) else: Or.append( self._calc_or(Rmean, rperi, rap, E, L, fixed_quad, **kwargs)) Op.append( self._calc_op(Or[-1], Rmean, rperi, rap, E, L, fixed_quad, **kwargs)) #Angles ar.append( self._calc_angler(Or[-1], axiR[ii], Rmean, rperi, rap, E, L, axivR[ii], fixed_quad, **kwargs)) az.append( self._calc_anglez(Or[-1], Op[-1], ar[-1], z[ii], axiR[ii], Rmean, rperi, rap, E, L, Lz[ii], axivR[ii], axivz[ii], fixed_quad, **kwargs)) Op = nu.array(Op) Oz = copy.copy(Op) Op[vT < 0.] *= -1. ap = copy.copy(asc) ar = nu.array(ar) az = nu.array(az) ap[vT < 0.] -= az[vT < 0.] ap[vT >= 0.] += az[vT >= 0.] ar = ar % (2. * nu.pi) ap = ap % (2. * nu.pi) az = az % (2. * nu.pi) return (nu.array(Jr), Jphi, Jz, nu.array(Or), Op, Oz, ar, ap, az)
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 fixed_quad= (False) if True, use n=10 fixed_quad integration scipy.integrate.quadrature keywords OUTPUT: (jr,lz,jz,Omegar,Omegaphi,Omegaz) HISTORY: 2013-12-28 - Written - Bovy (IAS) """ fixed_quad = kwargs.pop('fixed_quad', False) 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: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_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]) if self._c: #pragma: no cover pass else: Lz = R * vT Lx = -z * vT Ly = z * vR - R * vz L2 = Lx * Lx + Ly * Ly + Lz * Lz E = _evaluatePotentials( self._pot, R, z) + vR**2. / 2. + vT**2. / 2. + vz**2. / 2. L = nu.sqrt(L2) #Actions Jphi = Lz Jz = L - nu.fabs(Lz) #Jr requires some more work #Set up an actionAngleAxi object for EL and rap/rperi calculations axiR = nu.sqrt(R**2. + z**2.) axivT = L / axiR axivR = (R * vR + z * vz) / axiR Jr = [] Or = [] Op = [] for ii in range(len(axiR)): axiaA = actionAngleAxi(axiR[ii], axivR[ii], axivT[ii], pot=self._2dpot) (rperi, rap) = axiaA.calcRapRperi() EL = axiaA.calcEL() E, L = EL Jr.append(self._calc_jr(rperi, rap, E, L, fixed_quad, **kwargs)) #Radial period if Jr[-1] < 10.**-9.: #Circular orbit Or.append(epifreq(self._pot, axiR[ii], use_physical=False)) Op.append(omegac(self._pot, axiR[ii], use_physical=False)) continue Rmean = m.exp((m.log(rperi) + m.log(rap)) / 2.) Or.append( self._calc_or(Rmean, rperi, rap, E, L, fixed_quad, **kwargs)) Op.append( self._calc_op(Or[-1], Rmean, rperi, rap, E, L, fixed_quad, **kwargs)) Op = nu.array(Op) Oz = copy.copy(Op) Op[vT < 0.] *= -1. return (nu.array(Jr), Jphi, Jz, nu.array(Or), Op, Oz)
def _evaluate(self,*args,**kwargs): """ NAME: __call__ (_evaluate) PURPOSE: evaluate the actions (jr,lz,jz) INPUT: Either: a) R,vR,vT,z,vz[,phi]: 1) floats: phase-space value for single object (phi is optional) (each can be a Quantity) 2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity) b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument fixed_quad= (False) if True, use n=10 fixed_quad integration scipy.integrate.quadrature or .fixed_quad keywords OUTPUT: (jr,lz,jz) HISTORY: 2013-12-28 - Written - Bovy (IAS) """ fixed_quad= kwargs.pop('fixed_quad',False) 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: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_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]) if self._c: #pragma: no cover pass else: Lz= R*vT Lx= -z*vT Ly= z*vR-R*vz L2= Lx*Lx+Ly*Ly+Lz*Lz E= _evaluatePotentials(self._pot,R,z)\ +vR**2./2.+vT**2./2.+vz**2./2. L= nu.sqrt(L2) #Actions Jphi= Lz Jz= L-nu.fabs(Lz) #Jr requires some more work #Set up an actionAngleAxi object for EL and rap/rperi calculations axiR= nu.sqrt(R**2.+z**2.) axivT= L/axiR axivR= (R*vR+z*vz)/axiR Jr= [] for ii in range(len(axiR)): axiaA= actionAngleAxi(axiR[ii],axivR[ii],axivT[ii], pot=self._2dpot) (rperi,rap)= axiaA.calcRapRperi() EL= axiaA.calcEL() E, L= EL Jr.append(self._calc_jr(rperi,rap,E,L,fixed_quad,**kwargs)) return (nu.array(Jr),Jphi,Jz)
def _actionsFreqsAngles(self,*args,**kwargs): """ NAME: actionsFreqsAngles (_actionsFreqsAngles) PURPOSE: evaluate the actions, frequencies, and angles (jr,lz,jz,Omegar,Omegaphi,Omegaz,ar,ap,az) INPUT: Either: a) R,vR,vT,z,vz[,phi]: 1) floats: phase-space value for single object (phi is optional) (each can be a Quantity) 2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity) b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument fixed_quad= (False) if True, use n=10 fixed_quad integration scipy.integrate.quadrature or .fixed_quad keywords OUTPUT: (jr,lz,jz,Omegar,Omegaphi,Omegaz,ar,aphi,az) HISTORY: 2013-12-29 - Written - Bovy (IAS) """ fixed_quad= kwargs.pop('fixed_quad',False) if len(args) == 5: #R,vR.vT, z, vz pragma: no cover raise IOError("You need to provide phi when calculating angles") elif len(args) == 6: #R,vR.vT, z, vz, phi R,vR,vT, z, vz, phi= args else: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_vz phi= self._eval_phi if isinstance(R,float): R= nu.array([R]) vR= nu.array([vR]) vT= nu.array([vT]) z= nu.array([z]) vz= nu.array([vz]) phi= nu.array([phi]) if self._c: #pragma: no cover pass else: Lz= R*vT Lx= -z*vT Ly= z*vR-R*vz L2= Lx*Lx+Ly*Ly+Lz*Lz E= _evaluatePotentials(self._pot,R,z)+vR**2./2.+vT**2./2.+vz**2./2. L= nu.sqrt(L2) #Actions Jphi= Lz Jz= L-nu.fabs(Lz) #Jr requires some more work #Set up an actionAngleAxi object for EL and rap/rperi calculations axiR= nu.sqrt(R**2.+z**2.) axivT= L/axiR #these are really spherical coords, called axi bc they go in actionAngleAxi axivR= (R*vR+z*vz)/axiR axivz= (z*vR-R*vz)/axiR Jr= [] Or= [] Op= [] ar= [] az= [] #Calculate the longitude of the ascending node asc= self._calc_long_asc(z,R,axivz,phi,Lz,L) for ii in range(len(axiR)): axiaA= actionAngleAxi(axiR[ii],axivR[ii],axivT[ii], pot=self._2dpot) (rperi,rap)= axiaA.calcRapRperi() EL= axiaA.calcEL() E, L= EL Jr.append(self._calc_jr(rperi,rap,E,L,fixed_quad,**kwargs)) #Radial period Rmean= m.exp((m.log(rperi)+m.log(rap))/2.) if Jr[-1] < 10.**-9.: #Circular orbit Or.append(epifreq(self._pot,axiR[ii],use_physical=False)) Op.append(omegac(self._pot,axiR[ii],use_physical=False)) else: Or.append(self._calc_or(Rmean,rperi,rap,E,L,fixed_quad,**kwargs)) Op.append(self._calc_op(Or[-1],Rmean,rperi,rap,E,L,fixed_quad,**kwargs)) #Angles ar.append(self._calc_angler(Or[-1],axiR[ii],Rmean,rperi,rap, E,L,axivR[ii],fixed_quad,**kwargs)) az.append(self._calc_anglez(Or[-1],Op[-1],ar[-1], z[ii],axiR[ii], Rmean,rperi,rap,E,L,Lz[ii], axivR[ii],axivz[ii], fixed_quad,**kwargs)) Op= nu.array(Op) Oz= copy.copy(Op) Op[vT < 0.]*= -1. ap= copy.copy(asc) ar= nu.array(ar) az= nu.array(az) ap[vT < 0.]-= az[vT < 0.] ap[vT >= 0.]+= az[vT >= 0.] ar= ar % (2.*nu.pi) ap= ap % (2.*nu.pi) az= az % (2.*nu.pi) return (nu.array(Jr),Jphi,Jz,nu.array(Or),Op,Oz, ar,ap,az)
def __init__(self,pot=None,zmax=1.,gamma=1.,Rmax=5., nR=16,nEz=16,nEr=31,nLz=31,numcores=1, **kwargs): """ NAME: __init__ PURPOSE: initialize an actionAngleAdiabaticGrid object INPUT: pot= potential or list of potentials zmax= zmax for building Ez grid Rmax = Rmax for building grids gamma= (default=1.) replace Lz by Lz+gamma Jz in effective potential nEz=, nEr=, nLz, nR= grid size numcores= number of cpus to use to parallellize c= if True, use C to calculate actions ro= distance from vantage point to GC (kpc; can be Quantity) vo= circular velocity at ro (km/s; can be Quantity) +scipy.integrate.quad keywords OUTPUT: instance HISTORY: 2012-07-27 - Written - Bovy (IAS@MPIA) """ actionAngle.__init__(self, ro=kwargs.get('ro',None),vo=kwargs.get('vo',None)) if pot is None: #pragma: no cover raise IOError("Must specify pot= for actionAngleAxi") self._c= kwargs.pop('c',False) self._gamma= gamma self._pot= pot self._zmax= zmax self._Rmax= Rmax self._Rmin= 0.01 #Set up the actionAngleAdiabatic object that we will use to interpolate self._aA= actionAngleAdiabatic(pot=self._pot,gamma=self._gamma, c=self._c) #Build grid for Ez, first calculate Ez(zmax;R) function self._Rs= numpy.linspace(self._Rmin,self._Rmax,nR) self._EzZmaxs= _evaluatePotentials(self._pot,self._Rs, self._zmax*numpy.ones(nR))\ -_evaluatePotentials(self._pot,self._Rs,numpy.zeros(nR)) self._EzZmaxsInterp= interpolate.InterpolatedUnivariateSpline(self._Rs,numpy.log(self._EzZmaxs),k=3) y= numpy.linspace(0.,1.,nEz) jz= numpy.zeros((nR,nEz)) jzEzzmax= numpy.zeros(nR) thisRs= (numpy.tile(self._Rs,(nEz,1)).T).flatten() thisEzZmaxs= (numpy.tile(self._EzZmaxs,(nEz,1)).T).flatten() thisy= (numpy.tile(y,(nR,1))).flatten() if self._c: jz= self._aA(thisRs, numpy.zeros(len(thisRs)), numpy.ones(len(thisRs)),#these two r dummies numpy.zeros(len(thisRs)), numpy.sqrt(2.*thisy*thisEzZmaxs), **kwargs)[2] jz= numpy.reshape(jz,(nR,nEz)) jzEzzmax[0:nR]= jz[:,nEz-1] else: if numcores > 1: jz= multi.parallel_map((lambda x: self._aA(thisRs[x],0.,1.,#these two r dummies 0.,math.sqrt(2.*thisy[x]*thisEzZmaxs[x]), _justjz=True, **kwargs)[2]), range(nR*nEz),numcores=numcores) jz= numpy.reshape(jz,(nR,nEz)) jzEzzmax[0:nR]= jz[:,nEz-1] else: for ii in range(nR): for jj in range(nEz): #Calculate Jz jz[ii,jj]= self._aA(self._Rs[ii],0.,1.,#these two r dummies 0.,numpy.sqrt(2.*y[jj]*self._EzZmaxs[ii]), _justjz=True,**kwargs)[2] if jj == nEz-1: jzEzzmax[ii]= jz[ii,jj] for ii in range(nR): jz[ii,:]/= jzEzzmax[ii] #First interpolate Ez=Ezmax self._jzEzmaxInterp= interpolate.InterpolatedUnivariateSpline(self._Rs,numpy.log(jzEzzmax+10.**-5.),k=3) self._jz= jz self._jzInterp= interpolate.RectBivariateSpline(self._Rs, y, jz, kx=3,ky=3,s=0.) #JR 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] #Calculate ER(vr=0,R=RL) 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) self._ERRL= _evaluatePotentials(self._pot,self._RL,numpy.zeros(nLz)) +self._Lzs**2./2./self._RL**2. self._ERRLmax= numpy.amax(self._ERRL)+1. self._ERRLInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs, numpy.log(-(self._ERRL-self._ERRLmax)),k=3) self._Ramax= 99. self._ERRa= _evaluatePotentials(self._pot,self._Ramax,0.) +self._Lzs**2./2./self._Ramax**2. self._ERRamax= numpy.amax(self._ERRa)+1. self._ERRaInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs, numpy.log(-(self._ERRa-self._ERRamax)),k=3) y= numpy.linspace(0.,1.,nEr) jr= numpy.zeros((nLz,nEr)) jrERRa= numpy.zeros(nLz) thisRL= (numpy.tile(self._RL,(nEr-1,1)).T).flatten() thisLzs= (numpy.tile(self._Lzs,(nEr-1,1)).T).flatten() thisERRL= (numpy.tile(self._ERRL,(nEr-1,1)).T).flatten() thisERRa= (numpy.tile(self._ERRa,(nEr-1,1)).T).flatten() thisy= (numpy.tile(y[0:-1],(nLz,1))).flatten() if self._c: mjr= self._aA(thisRL, numpy.sqrt(2.*(thisERRa+thisy*(thisERRL-thisERRa)-_evaluatePotentials(self._pot,thisRL,numpy.zeros((nEr-1)*nLz)))-thisLzs**2./thisRL**2.), thisLzs/thisRL, numpy.zeros(len(thisRL)), numpy.zeros(len(thisRL)), **kwargs)[0] jr[:,0:-1]= numpy.reshape(mjr,(nLz,nEr-1)) jrERRa[0:nLz]= jr[:,0] else: if numcores > 1: mjr= multi.parallel_map((lambda x: self._aA(thisRL[x], numpy.sqrt(2.*(thisERRa[x]+thisy[x]*(thisERRL[x]-thisERRa[x])-_evaluatePotentials(self._pot,thisRL[x],0.))-thisLzs[x]**2./thisRL[x]**2.), thisLzs[x]/thisRL[x], 0.,0., _justjr=True, **kwargs)[0]), range((nEr-1)*nLz), numcores=numcores) jr[:,0:-1]= numpy.reshape(mjr,(nLz,nEr-1)) jrERRa[0:nLz]= jr[:,0] else: for ii in range(nLz): for jj in range(nEr-1): #Last one is zero by construction try: jr[ii,jj]= self._aA(self._RL[ii], numpy.sqrt(2.*(self._ERRa[ii]+y[jj]*(self._ERRL[ii]-self._ERRa[ii])-_evaluatePotentials(self._pot,self._RL[ii],0.))-self._Lzs[ii]**2./self._RL[ii]**2.), self._Lzs[ii]/self._RL[ii], 0.,0., _justjr=True, **kwargs)[0] except UnboundError: #pragma: no cover raise if jj == 0: jrERRa[ii]= jr[ii,jj] for ii in range(nLz): jr[ii,:]/= jrERRa[ii] #First interpolate Ez=Ezmax self._jr= jr self._jrERRaInterp= interpolate.InterpolatedUnivariateSpline(self._Lzs, numpy.log(jrERRa+10.**-5.),k=3) self._jrInterp= interpolate.RectBivariateSpline(self._Lzs, y, jr, kx=3,ky=3,s=0.) # Check the units self._check_consistent_units() return None
def _actionsFreqsAngles(self, *args, **kwargs): """ NAME: _actionsFreqsAngles PURPOSE: evaluate the actions, frequencies, and angles (jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez) INPUT: Either: a) R,vR,vT,z,vz,phi (MUST HAVE PHI) 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,angler,anglephi,anglez) HISTORY: 2013-08-28 - Written - Bovy (IAS) """ if ((self._c and not ('c' in kwargs and not kwargs['c']))\ or (ext_loaded and (('c' in kwargs and kwargs['c'])))) \ and _check_c(self._pot): if len(args) == 5: #R,vR.vT, z, vz pragma: no cover raise IOError("Must specify phi") elif len(args) == 6: #R,vR.vT, z, vz, phi R, vR, vT, z, vz, phi = args else: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_vz phi = self._eval_phi if isinstance(R, float): R = nu.array([R]) vR = nu.array([vR]) vT = nu.array([vT]) z = nu.array([z]) vz = nu.array([vz]) phi = nu.array([phi]) Lz = R * vT if self._useu0: #First calculate u0 if 'u0' in kwargs: u0 = nu.asarray(kwargs['u0']) else: E = nu.array([ _evaluatePotentials(self._pot, R[ii], z[ii]) + 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] kwargs.pop('u0', None) else: u0 = None jr, jz, Omegar, Omegaphi, Omegaz, angler, anglephi,anglez, err= actionAngleStaeckel_c.actionAngleFreqAngleStaeckel_c(\ self._pot,self._delta,R,vR,vT,z,vz,phi,u0=u0) # Adjustements for close-to-circular orbits indx = nu.isnan(Omegar) * (jr < 10.**-3.) + nu.isnan(Omegaz) * ( jz < 10.**-3. ) #Close-to-circular and close-to-the-plane orbits if nu.sum(indx) > 0: Omegar[indx] = [ epifreq(self._pot, r, use_physical=False) for r in R[indx] ] Omegaphi[indx] = [ omegac(self._pot, r, use_physical=False) for r in R[indx] ] Omegaz[indx] = [ verticalfreq(self._pot, r, use_physical=False) for r in R[indx] ] if err == 0: return (jr, Lz, jz, Omegar, Omegaphi, Omegaz, angler, anglephi, anglez) else: raise RuntimeError( "C-code for calculation actions failed; try with c=False" ) #pragma: no cover else: #pragma: no cover if 'c' in kwargs and kwargs['c'] and not self._c: #pragma: no cover warnings.warn( "C module not used because potential does not have a C implementation", galpyWarning) raise NotImplementedError( "actionsFreqs with c=False not implemented")
def _evaluate(self, *args, **kwargs): """ NAME: _evaluate 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 scipy.integrate.quadrature keywords OUTPUT: (jr,lz,jz) HISTORY: 2012-07-27 - Written - Bovy (IAS@MPIA) NOTE: For a Miyamoto-Nagai potential, this seems accurate to 0.1% and takes ~0.13 ms For a MWPotential, this takes ~ 0.17 ms """ 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: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_vz #First work on the vertical action Phi = _evaluatePotentials(self._pot, R, z) try: Phio = _evaluatePotentials(self._pot, R, numpy.zeros(len(R))) except TypeError: Phio = _evaluatePotentials(self._pot, R, 0.) Ez = Phi - Phio + vz**2. / 2. #Bigger than Ezzmax? thisEzZmax = numpy.exp(self._EzZmaxsInterp(R)) if isinstance(R, numpy.ndarray): indx = (R > self._Rmax) indx += (R < self._Rmin) indx += (Ez != 0.) * (numpy.log(Ez) > thisEzZmax) indxc = True ^ indx jz = numpy.empty(R.shape) if numpy.sum(indxc) > 0: jz[indxc]= (self._jzInterp.ev(R[indxc],Ez[indxc]/thisEzZmax[indxc])\ *(numpy.exp(self._jzEzmaxInterp(R[indxc]))-10.**-5.)) if numpy.sum(indx) > 0: jz[indx] = self._aA( R[indx], numpy.zeros(numpy.sum(indx)), numpy.ones(numpy.sum(indx)), #these two r dummies numpy.zeros(numpy.sum(indx)), numpy.sqrt(2. * Ez[indx]), _justjz=True, **kwargs)[2] else: if R > self._Rmax or R < self._Rmin or ( Ez != 0 and numpy.log(Ez) > thisEzZmax): #Outside of the grid if _PRINTOUTSIDEGRID: #pragma: no cover print("Outside of grid in Ez", R > self._Rmax, R < self._Rmin, (Ez != 0 and numpy.log(Ez) > thisEzZmax)) jz = self._aA( R, 0., 1., #these two r dummies 0., math.sqrt(2. * Ez), _justjz=True, **kwargs)[2] else: jz= (self._jzInterp(R,Ez/thisEzZmax)\ *(numpy.exp(self._jzEzmaxInterp(R))-10.**-5.))[0][0] #Radial action ERLz = numpy.fabs(R * vT) + self._gamma * jz ER = Phio + vR**2. / 2. + ERLz**2. / 2. / R**2. thisRL = self._RLInterp(ERLz) thisERRL = -numpy.exp(self._ERRLInterp(ERLz)) + self._ERRLmax thisERRa = -numpy.exp(self._ERRaInterp(ERLz)) + self._ERRamax if isinstance(R, numpy.ndarray): indx= ((ER-thisERRa)/(thisERRL-thisERRa) > 1.)\ *(((ER-thisERRa)/(thisERRL-thisERRa)-1.) < 10.**-2.) ER[indx] = thisERRL[indx] indx= ((ER-thisERRa)/(thisERRL-thisERRa) < 0.)\ *((ER-thisERRa)/(thisERRL-thisERRa) > -10.**-2.) ER[indx] = thisERRa[indx] indx = (ERLz < self._Lzmin) indx += (ERLz > self._Lzmax) indx += ((ER - thisERRa) / (thisERRL - thisERRa) > 1.) indx += ((ER - thisERRa) / (thisERRL - thisERRa) < 0.) indxc = True ^ indx jr = numpy.empty(R.shape) if numpy.sum(indxc) > 0: jr[indxc]= (self._jrInterp.ev(ERLz[indxc], (ER[indxc]-thisERRa[indxc])/(thisERRL[indxc]-thisERRa[indxc]))\ *(numpy.exp(self._jrERRaInterp(ERLz[indxc]))-10.**-5.)) if numpy.sum(indx) > 0: jr[indx] = self._aA( thisRL[indx], numpy.sqrt( 2. * (ER[indx] - _evaluatePotentials(self._pot, thisRL[indx], 0.)) - ERLz[indx]**2. / thisRL[indx]**2.), ERLz[indx] / thisRL[indx], numpy.zeros(len(thisRL)), numpy.zeros(len(thisRL)), _justjr=True, **kwargs)[0] else: if (ER-thisERRa)/(thisERRL-thisERRa) > 1. \ and ((ER-thisERRa)/(thisERRL-thisERRa)-1.) < 10.**-2.: ER = thisERRL elif (ER-thisERRa)/(thisERRL-thisERRa) < 0. \ and (ER-thisERRa)/(thisERRL-thisERRa) > -10.**-2.: ER = thisERRa #Outside of grid? if ERLz < self._Lzmin or ERLz > self._Lzmax \ or (ER-thisERRa)/(thisERRL-thisERRa) > 1. \ or (ER-thisERRa)/(thisERRL-thisERRa) < 0.: if _PRINTOUTSIDEGRID: #pragma: no cover print("Outside of grid in ER/Lz", ERLz < self._Lzmin , ERLz > self._Lzmax \ , (ER-thisERRa)/(thisERRL-thisERRa) > 1. \ , (ER-thisERRa)/(thisERRL-thisERRa) < 0., ER, thisERRL, thisERRa, (ER-thisERRa)/(thisERRL-thisERRa)) jr = self._aA( thisRL[0], numpy.sqrt( 2. * (ER - _evaluatePotentials(self._pot, thisRL, 0.)) - ERLz**2. / thisRL**2.)[0], (ERLz / thisRL)[0], 0., 0., _justjr=True, **kwargs)[0] else: jr= (self._jrInterp(ERLz, (ER-thisERRa)/(thisERRL-thisERRa))\ *(numpy.exp(self._jrERRaInterp(ERLz))-10.**-5.))[0][0] return (jr, R * vT, jz)
def _evaluate(self,*args,**kwargs): """ NAME: _evaluate PURPOSE: evaluate the actions (jr,lz,jz) INPUT: Either: R,vR,vT,z,vz scipy.integrate.quadrature keywords (for off-the-grid calcs) OUTPUT: (jr,lz,jz) HISTORY: 2012-11-29 - Written - Bovy (IAS) """ 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: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_vz Lz= R*vT Phi= _evaluatePotentials(self._pot,R,z) E= Phi+vR**2./2.+vT**2./2.+vz**2./2. thisERL= -numpy.exp(self._ERLInterp(Lz))+self._ERLmax thisERa= -numpy.exp(self._ERaInterp(Lz))+self._ERamax if isinstance(R,numpy.ndarray): indx= ((E-thisERa)/(thisERL-thisERa) > 1.)\ *(((E-thisERa)/(thisERL-thisERa)-1.) < 10.**-2.) E[indx]= thisERL[indx] indx= ((E-thisERa)/(thisERL-thisERa) < 0.)\ *((E-thisERa)/(thisERL-thisERa) > -10.**-2.) E[indx]= thisERa[indx] indx= (Lz < self._Lzmin) indx+= (Lz > self._Lzmax) indx+= ((E-thisERa)/(thisERL-thisERa) > 1.) indx+= ((E-thisERa)/(thisERL-thisERa) < 0.) indxc= True-indx jr= numpy.empty(R.shape) jz= numpy.empty(R.shape) if numpy.sum(indxc) > 0: u0= numpy.exp(self._logu0Interp.ev(Lz[indxc], (_Efunc(E[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc]))/(_Efunc(thisERL[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc])))) sinh2u0= numpy.sinh(u0)**2. thisEr= self.Er(R[indxc],z[indxc],vR[indxc],vz[indxc], E[indxc],Lz[indxc],sinh2u0,u0) thisEz= self.Ez(R[indxc],z[indxc],vR[indxc],vz[indxc], E[indxc],Lz[indxc],sinh2u0,u0) thisv2= self.vatu0(E[indxc],Lz[indxc],u0,self._delta*numpy.sinh(u0),retv2=True) cos2psi= 2.*thisEr/thisv2/(1.+sinh2u0) #latter is cosh2u0 cos2psi[(cos2psi > 1.)*(cos2psi < 1.+10.**-5.)]= 1. indxCos2psi= (cos2psi > 1.) indxCos2psi+= (cos2psi < 0.) indxc[indxc]= True-indxCos2psi#Handle these two cases as off-grid indx= True-indxc psi= numpy.arccos(numpy.sqrt(cos2psi[True-indxCos2psi])) coords= numpy.empty((3,numpy.sum(indxc))) coords[0,:]= (Lz[indxc]-self._Lzmin)/(self._Lzmax-self._Lzmin)*(self._nLz-1.) y= (_Efunc(E[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc]))/(_Efunc(thisERL[indxc],thisERL[indxc])-_Efunc(thisERa[indxc],thisERL[indxc])) coords[1,:]= y*(self._nE-1.) coords[2,:]= psi/numpy.pi*2.*(self._npsi-1.) jr[indxc]= (numpy.exp(ndimage.interpolation.map_coordinates(self._jrFiltered, coords, order=3, prefilter=False))-10.**-10.)*(numpy.exp(self._jrLzInterp(Lz[indxc]))-10.**-5.) #Switch to Ez-calculated psi sin2psi= 2.*thisEz[True-indxCos2psi]/thisv2[True-indxCos2psi]/(1.+sinh2u0[True-indxCos2psi]) #latter is cosh2u0 sin2psi[(sin2psi > 1.)*(sin2psi < 1.+10.**-5.)]= 1. indxSin2psi= (sin2psi > 1.) indxSin2psi+= (sin2psi < 0.) indxc[indxc]= True-indxSin2psi#Handle these two cases as off-grid indx= True-indxc psiz= numpy.arcsin(numpy.sqrt(sin2psi[True-indxSin2psi])) newcoords= numpy.empty((3,numpy.sum(indxc))) newcoords[0:2,:]= coords[0:2,True-indxSin2psi] newcoords[2,:]= psiz/numpy.pi*2.*(self._npsi-1.) jz[indxc]= (numpy.exp(ndimage.interpolation.map_coordinates(self._jzFiltered, newcoords, order=3, prefilter=False))-10.**-10.)*(numpy.exp(self._jzLzInterp(Lz[indxc]))-10.**-5.) if numpy.sum(indx) > 0: jrindiv, lzindiv, jzindiv= self._aA(R[indx], vR[indx], vT[indx], z[indx], vz[indx], **kwargs) jr[indx]= jrindiv jz[indx]= jzindiv """ jrindiv= numpy.empty(numpy.sum(indx)) jzindiv= numpy.empty(numpy.sum(indx)) for ii in range(numpy.sum(indx)): try: thisaA= actionAngleStaeckel.actionAngleStaeckelSingle(\ R[indx][ii], #R vR[indx][ii], #vR vT[indx][ii], #vT z[indx][ii], #z vz[indx][ii], #vz pot=self._pot,delta=self._delta) jrindiv[ii]= thisaA.JR(fixed_quad=True)[0] jzindiv[ii]= thisaA.Jz(fixed_quad=True)[0] except (UnboundError,OverflowError): jrindiv[ii]= numpy.nan jzindiv[ii]= numpy.nan jr[indx]= jrindiv jz[indx]= jzindiv """ else: jr,Lz, jz= self(numpy.array([R]), numpy.array([vR]), numpy.array([vT]), numpy.array([z]), numpy.array([vz]), **kwargs) return (jr[0],Lz[0],jz[0]) jr[jr < 0.]= 0. jz[jz < 0.]= 0. return (jr,R*vT,jz)
def _evaluate(self, *args, **kwargs): """ NAME: _evaluate 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 ('c' in kwargs and not kwargs['c']))\ or (ext_loaded and (('c' in kwargs 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: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_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 'u0' in kwargs: u0 = nu.asarray(kwargs['u0']) else: E = nu.array([ _evaluatePotentials(self._pot, R[ii], z[ii]) + 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] kwargs.pop('u0', None) 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: #pragma: no cover raise RuntimeError( "C-code for calculation actions failed; try with c=False") else: if 'c' in kwargs and kwargs['c'] and not self._c: #pragma: no cover warnings.warn( "C module not used because potential does not have a C implementation", galpyWarning) kwargs.pop('c', None) 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 _evaluate(self,*args,**kwargs): """ NAME: _evaluate 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 ('c' in kwargs and not kwargs['c']))\ or (ext_loaded and (('c' in kwargs 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: self._parse_eval_args(*args) R= self._eval_R vR= self._eval_vR vT= self._eval_vT z= self._eval_z vz= self._eval_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 'u0' in kwargs: u0= nu.asarray(kwargs['u0']) else: E= nu.array([_evaluatePotentials(self._pot,R[ii],z[ii]) +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] kwargs.pop('u0',None) 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: #pragma: no cover raise RuntimeError("C-code for calculation actions failed; try with c=False") else: if 'c' in kwargs and kwargs['c'] and not self._c: #pragma: no cover warnings.warn("C module not used because potential does not have a C implementation",galpyWarning) kwargs.pop('c',None) 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=30,numcores=1, **kwargs): """ NAME: __init__ PURPOSE: initialize an actionAngleStaeckelGrid object INPUT: pot= potential or list of potentials delta= focus of prolate confocal coordinate system (can be Quantity) Rmax = Rmax for building grids (natural units) nE=, npsi=, nLz= grid size numcores= number of cpus to use to parallellize ro= distance from vantage point to GC (kpc; can be Quantity) vo= circular velocity at ro (km/s; can be Quantity) OUTPUT: instance HISTORY: 2012-11-29 - Written - Bovy (IAS) """ actionAngle.__init__(self, ro=kwargs.get('ro',None),vo=kwargs.get('vo',None)) 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 'c' in kwargs and kwargs['c']: self._c= True else: self._c= False self._delta= delta if _APY_LOADED and isinstance(self._delta,units.Quantity): self._delta= self._delta.to(units.kpc).value/self._ro 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) self._ERL= _evaluatePotentials(self._pot,self._RL, numpy.zeros(self._nLz))\ +self._Lzs**2./2./self._RL**2. 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. self._ERa= _evaluatePotentials(self._pot,self._Ramax,0.) +self._Lzs**2./2./self._Ramax**2. #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= list(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.nanmax(jr[ii,(jr[ii,:,:] != 9999.99)])#:,:]) jzLzE[ii]= numpy.nanmax(jz[ii,(jz[ii,:,:] != 9999.99)])#:,:]) jrLzE[(jrLzE == 0.)]= numpy.nanmin(jrLzE[(jrLzE > 0.)]) jzLzE[(jzLzE == 0.)]= numpy.nanmin(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. #Deal w/ NaN jr[numpy.isnan(jr)]= 0. jz[numpy.isnan(jz)]= 0. #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) # Check the units self._check_consistent_units() return None
def __init__(self, pot=None, zmax=1., gamma=1., Rmax=5., nR=16, nEz=16, nEr=31, nLz=31, numcores=1, **kwargs): """ NAME: __init__ PURPOSE: initialize an actionAngleAdiabaticGrid object INPUT: pot= potential or list of potentials zmax= zmax for building Ez grid Rmax = Rmax for building grids gamma= (default=1.) replace Lz by Lz+gamma Jz in effective potential nEz=, nEr=, nLz, nR= grid size numcores= number of cpus to use to parallellize c= if True, use C to calculate actions ro= distance from vantage point to GC (kpc; can be Quantity) vo= circular velocity at ro (km/s; can be Quantity) +scipy.integrate.quad keywords OUTPUT: instance HISTORY: 2012-07-27 - Written - Bovy (IAS@MPIA) """ actionAngle.__init__(self, ro=kwargs.get('ro', None), vo=kwargs.get('vo', None)) if pot is None: #pragma: no cover raise IOError("Must specify pot= for actionAngleAxi") self._c = kwargs.pop('c', False) self._gamma = gamma self._pot = pot self._zmax = zmax self._Rmax = Rmax self._Rmin = 0.01 #Set up the actionAngleAdiabatic object that we will use to interpolate self._aA = actionAngleAdiabatic(pot=self._pot, gamma=self._gamma, c=self._c) #Build grid for Ez, first calculate Ez(zmax;R) function self._Rs = numpy.linspace(self._Rmin, self._Rmax, nR) self._EzZmaxs= _evaluatePotentials(self._pot,self._Rs, self._zmax*numpy.ones(nR))\ -_evaluatePotentials(self._pot,self._Rs,numpy.zeros(nR)) self._EzZmaxsInterp = interpolate.InterpolatedUnivariateSpline( self._Rs, numpy.log(self._EzZmaxs), k=3) y = numpy.linspace(0., 1., nEz) jz = numpy.zeros((nR, nEz)) jzEzzmax = numpy.zeros(nR) thisRs = (numpy.tile(self._Rs, (nEz, 1)).T).flatten() thisEzZmaxs = (numpy.tile(self._EzZmaxs, (nEz, 1)).T).flatten() thisy = (numpy.tile(y, (nR, 1))).flatten() if self._c: jz = self._aA( thisRs, numpy.zeros(len(thisRs)), numpy.ones(len(thisRs)), #these two r dummies numpy.zeros(len(thisRs)), numpy.sqrt(2. * thisy * thisEzZmaxs), **kwargs)[2] jz = numpy.reshape(jz, (nR, nEz)) jzEzzmax[0:nR] = jz[:, nEz - 1] else: if numcores > 1: jz = multi.parallel_map( ( lambda x: self._aA( thisRs[x], 0., 1., #these two r dummies 0., math.sqrt(2. * thisy[x] * thisEzZmaxs[x]), _justjz=True, **kwargs)[2]), range(nR * nEz), numcores=numcores) jz = numpy.reshape(jz, (nR, nEz)) jzEzzmax[0:nR] = jz[:, nEz - 1] else: for ii in range(nR): for jj in range(nEz): #Calculate Jz jz[ii, jj] = self._aA( self._Rs[ii], 0., 1., #these two r dummies 0., numpy.sqrt(2. * y[jj] * self._EzZmaxs[ii]), _justjz=True, **kwargs)[2] if jj == nEz - 1: jzEzzmax[ii] = jz[ii, jj] for ii in range(nR): jz[ii, :] /= jzEzzmax[ii] #First interpolate Ez=Ezmax self._jzEzmaxInterp = interpolate.InterpolatedUnivariateSpline( self._Rs, numpy.log(jzEzzmax + 10.**-5.), k=3) self._jz = jz self._jzInterp = interpolate.RectBivariateSpline(self._Rs, y, jz, kx=3, ky=3, s=0.) #JR 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] #Calculate ER(vr=0,R=RL) 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) self._ERRL = _evaluatePotentials( self._pot, self._RL, numpy.zeros(nLz)) + self._Lzs**2. / 2. / self._RL**2. self._ERRLmax = numpy.amax(self._ERRL) + 1. self._ERRLInterp = interpolate.InterpolatedUnivariateSpline( self._Lzs, numpy.log(-(self._ERRL - self._ERRLmax)), k=3) self._Ramax = 99. self._ERRa = _evaluatePotentials( self._pot, self._Ramax, 0.) + self._Lzs**2. / 2. / self._Ramax**2. self._ERRamax = numpy.amax(self._ERRa) + 1. self._ERRaInterp = interpolate.InterpolatedUnivariateSpline( self._Lzs, numpy.log(-(self._ERRa - self._ERRamax)), k=3) y = numpy.linspace(0., 1., nEr) jr = numpy.zeros((nLz, nEr)) jrERRa = numpy.zeros(nLz) thisRL = (numpy.tile(self._RL, (nEr - 1, 1)).T).flatten() thisLzs = (numpy.tile(self._Lzs, (nEr - 1, 1)).T).flatten() thisERRL = (numpy.tile(self._ERRL, (nEr - 1, 1)).T).flatten() thisERRa = (numpy.tile(self._ERRa, (nEr - 1, 1)).T).flatten() thisy = (numpy.tile(y[0:-1], (nLz, 1))).flatten() if self._c: mjr = self._aA( thisRL, numpy.sqrt(2. * (thisERRa + thisy * (thisERRL - thisERRa) - _evaluatePotentials( self._pot, thisRL, numpy.zeros( (nEr - 1) * nLz))) - thisLzs**2. / thisRL**2.), thisLzs / thisRL, numpy.zeros(len(thisRL)), numpy.zeros(len(thisRL)), **kwargs)[0] jr[:, 0:-1] = numpy.reshape(mjr, (nLz, nEr - 1)) jrERRa[0:nLz] = jr[:, 0] else: if numcores > 1: mjr = multi.parallel_map((lambda x: self._aA( thisRL[x], numpy.sqrt(2. * (thisERRa[x] + thisy[x] * (thisERRL[x] - thisERRa[x]) - _evaluatePotentials(self._pot, thisRL[ x], 0.)) - thisLzs[x]**2. / thisRL[x]**2.), thisLzs[x] / thisRL[x], 0., 0., _justjr=True, **kwargs)[0]), range((nEr - 1) * nLz), numcores=numcores) jr[:, 0:-1] = numpy.reshape(mjr, (nLz, nEr - 1)) jrERRa[0:nLz] = jr[:, 0] else: for ii in range(nLz): for jj in range(nEr - 1): #Last one is zero by construction try: jr[ii, jj] = self._aA( self._RL[ii], numpy.sqrt(2. * (self._ERRa[ii] + y[jj] * (self._ERRL[ii] - self._ERRa[ii]) - _evaluatePotentials( self._pot, self._RL[ii], 0.)) - self._Lzs[ii]**2. / self._RL[ii]**2.), self._Lzs[ii] / self._RL[ii], 0., 0., _justjr=True, **kwargs)[0] except UnboundError: #pragma: no cover raise if jj == 0: jrERRa[ii] = jr[ii, jj] for ii in range(nLz): jr[ii, :] /= jrERRa[ii] #First interpolate Ez=Ezmax self._jr = jr self._jrERRaInterp = interpolate.InterpolatedUnivariateSpline( self._Lzs, numpy.log(jrERRa + 10.**-5.), k=3) self._jrInterp = interpolate.RectBivariateSpline(self._Lzs, y, jr, kx=3, ky=3, s=0.) # Check the units self._check_consistent_units() return None
def _actionsFreqsAngles(self, *args, **kwargs): """ NAME: actionsFreqsAngles (_actionsFreqsAngles) PURPOSE: evaluate the actions, frequencies, and angles (jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez) INPUT: Either: a) R,vR,vT,z,vz[,phi]: 1) floats: phase-space value for single object (phi is optional) (each can be a Quantity) 2) numpy.ndarray: [N] phase-space values for N objects (each can be a Quantity) b) Orbit instance: initial condition used if that's it, orbit(t) if there is a time given as well as the second argument delta= (object-wide default) can be used to override the object-wide focal length; can also be an array with length N to allow different delta for different phase-space points u0= (None) if object-wide option useu0 is set, u0 to use (if useu0 and useu0 is None, a good value will be computed) c= (object-wide default, bool) True/False to override the object-wide setting for whether or not to use the C implementation order= (10) number of points to use in the Gauss-Legendre numerical integration of the relevant action, frequency, and angle integrals When not using C: fixed_quad= (False) if True, use Gaussian quadrature (scipy.integrate.fixed_quad instead of scipy.integrate.quad) scipy.integrate.fixed_quad or .quad keywords OUTPUT: (jr,lz,jz,Omegar,Omegaphi,Omegaz,angler,anglephi,anglez) HISTORY: 2013-08-28 - Written - Bovy (IAS) """ delta = kwargs.pop('delta', self._delta) order = kwargs.get('order', self._order) if ((self._c and not ('c' in kwargs and not kwargs['c']))\ or (ext_loaded and (('c' in kwargs and kwargs['c'])))) \ and _check_c(self._pot): if len(args) == 5: #R,vR.vT, z, vz pragma: no cover raise IOError("Must specify phi") elif len(args) == 6: #R,vR.vT, z, vz, phi R, vR, vT, z, vz, phi = args else: self._parse_eval_args(*args) R = self._eval_R vR = self._eval_vR vT = self._eval_vT z = self._eval_z vz = self._eval_vz phi = self._eval_phi if isinstance(R, float): R = nu.array([R]) vR = nu.array([vR]) vT = nu.array([vT]) z = nu.array([z]) vz = nu.array([vz]) phi = nu.array([phi]) Lz = R * vT if self._useu0: #First calculate u0 if 'u0' in kwargs: u0 = nu.asarray(kwargs['u0']) else: E = nu.array([ _evaluatePotentials(self._pot, R[ii], z[ii]) + 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,delta)[0] kwargs.pop('u0', None) else: u0 = None jr, jz, Omegar, Omegaphi, Omegaz, angler, anglephi,anglez, err= actionAngleStaeckel_c.actionAngleFreqAngleStaeckel_c(\ self._pot,delta,R,vR,vT,z,vz,phi,u0=u0,order=order) # Adjustements for close-to-circular orbits indx = nu.isnan(Omegar) * (jr < 10.**-3.) + nu.isnan(Omegaz) * ( jz < 10.**-3. ) #Close-to-circular and close-to-the-plane orbits if nu.sum(indx) > 0: Omegar[indx] = [ epifreq(self._pot, r, use_physical=False) for r in R[indx] ] Omegaphi[indx] = [ omegac(self._pot, r, use_physical=False) for r in R[indx] ] Omegaz[indx] = [ verticalfreq(self._pot, r, use_physical=False) for r in R[indx] ] if err == 0: return (jr, Lz, jz, Omegar, Omegaphi, Omegaz, angler, anglephi, anglez) else: raise RuntimeError( "C-code for calculation actions failed; try with c=False" ) #pragma: no cover else: #pragma: no cover if 'c' in kwargs and kwargs['c'] and not self._c: #pragma: no cover warnings.warn( "C module not used because potential does not have a C implementation", galpyWarning) raise NotImplementedError( "actionsFreqs with c=False not implemented")