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:
__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 _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 _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 _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")
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 __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 ('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:
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 'u0' in kwargs:
u0 = nu.asarray(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]
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 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:
meta = actionAngle(*args)
R = meta._R
vR = meta._vR
vT = meta._vT
z = meta._z
vz = meta._vz
phi = meta._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(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]
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) for r in R[indx]]
Omegaphi[indx] = [omegac(self._pot, r) for r in R[indx]]
Omegaz[indx] = [verticalfreq(self._pot, r) 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")
