def test_meanjr():
#This is a *very* rough test against a rough estimate of the mean
qdf= quasiisothermaldf(1./4.,0.2,0.1,1.,1.,
pot=MWPotential,aA=aAS,cutcounter=True)
assert numpy.fabs(numpy.log(qdf.meanjr(0.9,0.,mc=True))\
-2.*numpy.log(0.2)-0.2
+numpy.log(epifreq(MWPotential,0.9))) < 0.4, 'meanjr is not what is expected'
assert numpy.fabs(numpy.log(qdf.meanjr(0.5,0.,mc=True))\
-2.*numpy.log(0.2)-1.
+numpy.log(epifreq(MWPotential,0.5))) < 0.4, 'meanjr is not what is expected'
return None
def epifreq(self,R):
from galpy.potential import epifreq
if self._interpepifreq:
indx= (R >= self._rgrid[0])*(R <= self._rgrid[-1])
out= numpy.empty_like(R)
if numpy.sum(indx) > 0:
if self._logR:
out[indx]= self._epifreqInterp(numpy.log(R[indx]))
else:
out[indx]= self._epifreqInterp(R[indx])
if numpy.sum(True^indx) > 0:
out[True^indx]= epifreq(self._origPot,R[True^indx])
return out
else:
return epifreq(self._origPot,R)
def LinShuReductionFactor(axiPot,
R,
sigmar,
nonaxiPot=None,
k=None,
m=None,
OmegaP=None):
"""
NAME:
LinShuReductionFactor
PURPOSE:
Calculate the Lin & Shu (1966) reduction factor: the reduced linear response of a kinematically-warm stellar disk to a perturbation
INPUT:
axiPot - The background, axisymmetric potential
R - Cylindrical radius (can be Quantity)
sigmar - radial velocity dispersion of the population (can be Quantity)
Then either provide:
1) m= m in the perturbation's m x phi (number of arms for a spiral)
k= wavenumber (see Binney & Tremaine 2008)
OmegaP= pattern speed (can be Quantity)
2) nonaxiPot= a non-axisymmetric Potential instance (such as SteadyLogSpiralPotential) that has functions that return OmegaP, m, and wavenumber
OUTPUT:
reduction factor
HISTORY:
2014-08-23 - Written - Bovy (IAS)
"""
axiPot = flatten(axiPot)
from galpy.potential import omegac, epifreq
if nonaxiPot is None and (OmegaP is None or k is None or m is None):
raise IOError(
"Need to specify either nonaxiPot= or m=, k=, OmegaP= for LinShuReductionFactor"
)
elif not nonaxiPot is None:
OmegaP = nonaxiPot.OmegaP()
k = nonaxiPot.wavenumber(R)
m = nonaxiPot.m()
tepif = epifreq(axiPot, R)
s = m * (OmegaP - omegac(axiPot, R)) / tepif
chi = sigmar**2. * k**2. / tepif**2.
return (1.-s**2.)/nu.sin(nu.pi*s)\
*integrate.quad(lambda t: nu.exp(-chi*(1.+nu.cos(t)))\
*nu.sin(s*t)*nu.sin(t),
0.,nu.pi)[0]
def epifreq(self,R):
if self._interpepifreq:
if self._logR:
return self._epifreqInterp(numpy.log(R))
else:
return self._epifreqInterp(R)
else:
from galpy.potential import epifreq
return epifreq(self._origPot,R)
def test_sigmat_staeckel_mc():
numpy.random.seed(2)
qdf= quasiisothermaldf(1./4.,0.2,0.1,1.,1.,
pot=MWPotential,aA=aAS,cutcounter=True)
#In the mid-plane
gamma= 2.*omegac(MWPotential,0.9)/epifreq(MWPotential,0.9)
assert numpy.fabs(numpy.log(qdf.sigmaT2(0.9,0.,mc=True)/qdf.sigmaR2(0.9,0.,mc=True))+2.*numpy.log(gamma)) < 0.3, "qdf's sigmaT2/sigmaR2 deviates more than expected from input for staeckel approx."
#higher up
assert numpy.fabs(numpy.log(qdf.sigmaT2(0.9,0.2,mc=True)/qdf.sigmaR2(0.9,0.2,mc=True))+2.*numpy.log(gamma)) < 0.3, "qdf's sigmaT2/sigmaR2 deviates more than expected from input for staeckel approx."
return None
def test_actionAngleTorus_basic_freqs():
from galpy.actionAngle import actionAngleTorus
from galpy.potential import epifreq, omegac, verticalfreq, rl, \
JaffePotential, PowerSphericalPotential, HernquistPotential
tol= -3.
jr= 10.**-6.
jz= 10.**-6.
jp= JaffePotential(normalize=1.)
aAT= actionAngleTorus(pot=jp)
# at Lz=1
jphi= 1.
om= aAT.Freqs(jr,jphi,jz)
assert numpy.fabs((om[0]-epifreq(jp,rl(jp,jphi)))/om[0]) < 10.**tol, \
'Close-to-circular orbit does not have Or=kappa for actionAngleTorus'
assert numpy.fabs((om[1]-omegac(jp,rl(jp,jphi)))/om[1]) < 10.**tol, \
'Close-to-circular orbit does not have Ophi=omega for actionAngleTorus'
assert numpy.fabs((om[2]-verticalfreq(jp,rl(jp,jphi)))/om[2]) < 10.**tol, \
'Close-to-circular orbit does not have Oz=nu for actionAngleTorus'
# at Lz=1.5, w/ different potential
pp= PowerSphericalPotential(normalize=1.)
aAT= actionAngleTorus(pot=pp)
jphi= 1.5
om= aAT.Freqs(jr,jphi,jz)
assert numpy.fabs((om[0]-epifreq(pp,rl(pp,jphi)))/om[0]) < 10.**tol, \
'Close-to-circular orbit does not have Or=kappa for actionAngleTorus'
assert numpy.fabs((om[1]-omegac(pp,rl(pp,jphi)))/om[1]) < 10.**tol, \
'Close-to-circular orbit does not have Ophi=omega for actionAngleTorus'
assert numpy.fabs((om[2]-verticalfreq(pp,rl(pp,jphi)))/om[2]) < 10.**tol, \
'Close-to-circular orbit does not have Oz=nu for actionAngleTorus'
# at Lz=0.5, w/ different potential
tol= -2.5 # appears more difficult
hp= HernquistPotential(normalize=1.)
aAT= actionAngleTorus(pot=hp)
jphi= 0.5
om= aAT.Freqs(jr,jphi,jz)
assert numpy.fabs((om[0]-epifreq(hp,rl(hp,jphi)))/om[0]) < 10.**tol, \
'Close-to-circular orbit does not have Or=kappa for actionAngleTorus'
assert numpy.fabs((om[1]-omegac(hp,rl(hp,jphi)))/om[1]) < 10.**tol, \
'Close-to-circular orbit does not have Ophi=omega for actionAngleTorus'
assert numpy.fabs((om[2]-verticalfreq(hp,rl(hp,jphi)))/om[2]) < 10.**tol, \
'Close-to-circular orbit does not have Oz=nu for actionAngleTorus'
return None
def test_actionAngleTorus_basic_freqs():
from galpy.actionAngle import actionAngleTorus
from galpy.potential import epifreq, omegac, verticalfreq, rl, \
JaffePotential, PowerSphericalPotential, HernquistPotential
tol= -3.
jr= 10.**-6.
jz= 10.**-6.
jp= JaffePotential(normalize=1.)
aAT= actionAngleTorus(pot=jp)
# at Lz=1
jphi= 1.
om= aAT.Freqs(jr,jphi,jz)
assert numpy.fabs((om[0]-epifreq(jp,rl(jp,jphi)))/om[0]) < 10.**tol, \
'Close-to-circular orbit does not have Or=kappa for actionAngleTorus'
assert numpy.fabs((om[1]-omegac(jp,rl(jp,jphi)))/om[1]) < 10.**tol, \
'Close-to-circular orbit does not have Ophi=omega for actionAngleTorus'
assert numpy.fabs((om[2]-verticalfreq(jp,rl(jp,jphi)))/om[2]) < 10.**tol, \
'Close-to-circular orbit does not have Oz=nu for actionAngleTorus'
# at Lz=1.5, w/ different potential
pp= PowerSphericalPotential(normalize=1.)
aAT= actionAngleTorus(pot=pp)
jphi= 1.5
om= aAT.Freqs(jr,jphi,jz)
assert numpy.fabs((om[0]-epifreq(pp,rl(pp,jphi)))/om[0]) < 10.**tol, \
'Close-to-circular orbit does not have Or=kappa for actionAngleTorus'
assert numpy.fabs((om[1]-omegac(pp,rl(pp,jphi)))/om[1]) < 10.**tol, \
'Close-to-circular orbit does not have Ophi=omega for actionAngleTorus'
assert numpy.fabs((om[2]-verticalfreq(pp,rl(pp,jphi)))/om[2]) < 10.**tol, \
'Close-to-circular orbit does not have Oz=nu for actionAngleTorus'
# at Lz=0.5, w/ different potential
tol= -2.5 # appears more difficult
hp= HernquistPotential(normalize=1.)
aAT= actionAngleTorus(pot=hp)
jphi= 0.5
om= aAT.Freqs(jr,jphi,jz)
assert numpy.fabs((om[0]-epifreq(hp,rl(hp,jphi)))/om[0]) < 10.**tol, \
'Close-to-circular orbit does not have Or=kappa for actionAngleTorus'
assert numpy.fabs((om[1]-omegac(hp,rl(hp,jphi)))/om[1]) < 10.**tol, \
'Close-to-circular orbit does not have Ophi=omega for actionAngleTorus'
assert numpy.fabs((om[2]-verticalfreq(hp,rl(hp,jphi)))/om[2]) < 10.**tol, \
'Close-to-circular orbit does not have Oz=nu for actionAngleTorus'
return None
def LinShuReductionFactor(axiPot,R,sigmar,nonaxiPot=None,
k=None,m=None,OmegaP=None):
"""
NAME:
LinShuReductionFactor
PURPOSE:
Calculate the Lin & Shu (1966) reduction factor: the reduced linear response of a kinematically-warm stellar disk to a perturbation
INPUT:
axiPot - The background, axisymmetric potential
R - Cylindrical radius (can be Quantity)
sigmar - radial velocity dispersion of the population (can be Quantity)
Then either provide:
1) m= m in the perturbation's m x phi (number of arms for a spiral)
k= wavenumber (see Binney & Tremaine 2008)
OmegaP= pattern speed (can be Quantity)
2) nonaxiPot= a non-axisymmetric Potential instance (such as SteadyLogSpiralPotential) that has functions that return OmegaP, m, and wavenumber
OUTPUT:
reduction factor
HISTORY:
2014-08-23 - Written - Bovy (IAS)
"""
axiPot= flatten(axiPot)
from galpy.potential import omegac, epifreq
if nonaxiPot is None and (OmegaP is None or k is None or m is None):
raise IOError("Need to specify either nonaxiPot= or m=, k=, OmegaP= for LinShuReductionFactor")
elif not nonaxiPot is None:
OmegaP= nonaxiPot.OmegaP()
k= nonaxiPot.wavenumber(R)
m= nonaxiPot.m()
tepif= epifreq(axiPot,R)
s= m*(OmegaP-omegac(axiPot,R))/tepif
chi= sigmar**2.*k**2./tepif**2.
return (1.-s**2.)/nu.sin(nu.pi*s)\
*integrate.quad(lambda t: nu.exp(-chi*(1.+nu.cos(t)))\
*nu.sin(s*t)*nu.sin(t),
0.,nu.pi)[0]
def test_call_diffinoutputs():
qdf= quasiisothermaldf(1./4.,0.2,0.1,1.,1.,
pot=MWPotential,aA=aAS,cutcounter=True)
#when specifying rg etc., first get these from a previous output
val, trg, tkappa, tnu, tOmega= qdf((0.03,0.9,0.02),_return_freqs=True)
#First check that just supplying these again works
assert numpy.fabs(val-qdf((0.03,0.9,0.02),rg=trg,kappa=tkappa,nu=tnu,
Omega=tOmega)) < 10.**-8., 'qdf calls w/ rg, and frequencies specified and w/ not specified do not agrees'
#Also calculate the frequencies
assert numpy.fabs(val-qdf((0.03,0.9,0.02),rg=trg,
kappa=epifreq(MWPotential,trg),
nu=verticalfreq(MWPotential,trg),
Omega=omegac(MWPotential,trg))) < 10.**-8., 'qdf calls w/ rg, and frequencies specified and w/ not specified do not agrees'
#Also test _return_actions
val, jr,lz,jz= qdf(0.9,0.1,0.95,0.1,0.08,_return_actions=True)
assert numpy.fabs(val-qdf((jr,lz,jz))) < 10.**-8., 'qdf call w/ R,vR,... and actions specified do not agree'
acs= aAS(0.9,0.1,0.95,0.1,0.08)
assert numpy.fabs(acs[0]-jr) < 10.**-8., 'direct calculation of jr and that returned from qdf.__call__ does not agree'
assert numpy.fabs(acs[1]-lz) < 10.**-8., 'direct calculation of lz and that returned from qdf.__call__ does not agree'
assert numpy.fabs(acs[2]-jz) < 10.**-8., 'direct calculation of jz and that returned from qdf.__call__ does not agree'
#Test unbound orbits
#Find unbound orbit, new qdf s.t. we can get UnboundError (only with
taAS= actionAngleStaeckel(pot=MWPotential,c=False,delta=0.5)
qdfnc= quasiisothermaldf(1./4.,0.2,0.1,1.,1.,
pot=MWPotential,
aA=taAS,
cutcounter=True)
from galpy.actionAngle import UnboundError
try: acs= taAS(0.9,10.,-20.,0.1,10.)
except UnboundError: pass
else:
print(acs)
raise AssertionError('Test orbit in qdf that is supposed to be unbound is not')
assert qdfnc(0.9,10.,-20.,0.1,10.) < 10.**-10., 'unbound orbit does not return qdf equal to zero'
#Test negative lz
assert qdf((0.03,-0.1,0.02)) < 10.**-8., 'qdf w/ cutcounter=True and negative lz does not return 0'
assert qdf((0.03,-0.1,0.02),log=True) <= numpy.finfo(numpy.dtype(numpy.float64)).min+1., 'qdf w/ cutcounter=True and negative lz does not return 0'
#Test func
val= qdf((0.03,0.9,0.02))
fval= qdf((0.03,0.9,0.02),func=lambda x,y,z: numpy.sin(x)*numpy.cos(y)\
*numpy.exp(z))
assert numpy.fabs(val*numpy.sin(0.03)*numpy.cos(0.9)*numpy.exp(0.02)-
fval) < 10.**-8, 'qdf __call__ w/ func does not work as expected'
lfval= qdf((0.03,0.9,0.02),func=lambda x,y,z: numpy.sin(x)*numpy.cos(y)\
*numpy.exp(z),log=True)
assert numpy.fabs(numpy.log(val)+numpy.log(numpy.sin(0.03)\
*numpy.cos(0.9)\
*numpy.exp(0.02))-
lfval) < 10.**-8, 'qdf __call__ w/ func does not work as expected'
return None
def _calc_epifreq(self, r):
"""
NAME:
_calc_epifreq
PURPOSE:
calculate the epicycle frequency at r
INPUT:
r - radius
OUTPUT:
kappa
HISTORY:
2012-07-25 - Written - Bovy (IAS@MPIA)
NOTE:
takes about 0.1 ms for a Miyamoto-Nagai potential
"""
return potential.epifreq(self._pot, r)
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 _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 _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 _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,
RZPot=None,rgrid=(0.01,2.,101),zgrid=(0.,0.2,101),logR=False,
interpPot=False,interpRforce=False,interpzforce=False,
interpDens=False,
interpvcirc=False,
interpdvcircdr=False,
interpepifreq=False,interpverticalfreq=False,
use_c=False,enable_c=False,zsym=True,
numcores=None):
"""
NAME:
__init__
PURPOSE:
Initialize an interpRZPotential instance
INPUT:
RZPot - RZPotential to be interpolated
rgrid - R grid to be given to linspace
zgrid - z grid to be given to linspace
logR - if True, rgrid is in the log of R
interpPot, interpRfoce, interpzforce, interpDens,interpvcirc, interpeopifreq, interpverticalfreq, interpdvcircdr= if True, interpolate these functions
use_c= use C to speed up the calculation
enable_c= enable use of C for interpolations
zsym= if True (default), the potential is assumed to be symmetric around z=0 (so you can use, e.g., zgrid=(0.,1.,101)).
numcores= if set to an integer, use this many cores (only used for vcirc, dvcircdR, epifreq, and verticalfreq; NOT NECESSARILY FASTER, TIME TO MAKE SURE)
OUTPUT:
instance
HISTORY:
2010-07-21 - Written - Bovy (NYU)
2013-01-24 - Started with new implementation - Bovy (IAS)
"""
Potential.__init__(self,amp=1.)
self.hasC= True
self._origPot= RZPot
self._rgrid= numpy.linspace(*rgrid)
self._logR= logR
if self._logR:
self._rgrid= numpy.exp(self._rgrid)
self._logrgrid= numpy.log(self._rgrid)
self._zgrid= numpy.linspace(*zgrid)
self._interpPot= interpPot
self._interpRforce= interpRforce
self._interpzforce= interpzforce
self._interpDens= interpDens
self._interpvcirc= interpvcirc
self._interpdvcircdr= interpdvcircdr
self._interpepifreq= interpepifreq
self._interpverticalfreq= interpverticalfreq
self._enable_c= enable_c*ext_loaded
self._zsym= zsym
if interpPot:
if use_c*ext_loaded:
self._potGrid, err= calc_potential_c(self._origPot,self._rgrid,self._zgrid)
else:
from galpy.potential import evaluatePotentials
potGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
potGrid[ii,jj]= evaluatePotentials(self._rgrid[ii],self._zgrid[jj],self._origPot)
self._potGrid= potGrid
if self._logR:
self._potInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
self._potGrid,
kx=3,ky=3,s=0.)
else:
self._potInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
self._potGrid,
kx=3,ky=3,s=0.)
if enable_c*ext_loaded:
self._potGrid_splinecoeffs= calc_2dsplinecoeffs_c(self._potGrid)
if interpRforce:
if use_c*ext_loaded:
self._rforceGrid, err= calc_potential_c(self._origPot,self._rgrid,self._zgrid,rforce=True)
else:
from galpy.potential import evaluateRforces
rforceGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
rforceGrid[ii,jj]= evaluateRforces(self._rgrid[ii],self._zgrid[jj],self._origPot)
self._rforceGrid= rforceGrid
if self._logR:
self._rforceInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
self._rforceGrid,
kx=3,ky=3,s=0.)
else:
self._rforceInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
self._rforceGrid,
kx=3,ky=3,s=0.)
if enable_c*ext_loaded:
self._rforceGrid_splinecoeffs= calc_2dsplinecoeffs_c(self._rforceGrid)
if interpzforce:
if use_c*ext_loaded:
self._zforceGrid, err= calc_potential_c(self._origPot,self._rgrid,self._zgrid,zforce=True)
else:
from galpy.potential import evaluatezforces
zforceGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
zforceGrid[ii,jj]= evaluatezforces(self._rgrid[ii],self._zgrid[jj],self._origPot)
self._zforceGrid= zforceGrid
if self._logR:
self._zforceInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
self._zforceGrid,
kx=3,ky=3,s=0.)
else:
self._zforceInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
self._zforceGrid,
kx=3,ky=3,s=0.)
if enable_c*ext_loaded:
self._zforceGrid_splinecoeffs= calc_2dsplinecoeffs_c(self._zforceGrid)
if interpDens:
if False:
raise NotImplementedError("Using C to calculate an interpolation grid for the density is not supported currently")
self._densGrid, err= calc_dens_c(self._origPot,self._rgrid,self._zgrid)
else:
from galpy.potential import evaluateDensities
densGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
densGrid[ii,jj]= evaluateDensities(self._rgrid[ii],self._zgrid[jj],self._origPot)
self._densGrid= densGrid
if self._logR:
self._densInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
numpy.log(self._densGrid+10.**-10.),
kx=3,ky=3,s=0.)
else:
self._densInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
numpy.log(self._densGrid+10.**-10.),
kx=3,ky=3,s=0.)
if False:
self._densGrid_splinecoeffs= calc_2dsplinecoeffs_c(self._densGrid)
if interpvcirc:
from galpy.potential import vcirc
if not numcores is None:
self._vcircGrid= multi.parallel_map((lambda x: vcirc(self._origPot,self._rgrid[x])),
range(len(self._rgrid)),numcores=numcores)
else:
self._vcircGrid= numpy.array([vcirc(self._origPot,r) for r in self._rgrid])
if self._logR:
self._vcircInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid,self._vcircGrid,k=3)
else:
self._vcircInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid,self._vcircGrid,k=3)
if interpdvcircdr:
from galpy.potential import dvcircdR
if not numcores is None:
self._dvcircdrGrid= multi.parallel_map((lambda x: dvcircdR(self._origPot,self._rgrid[x])),
range(len(self._rgrid)),numcores=numcores)
else:
self._dvcircdrGrid= numpy.array([dvcircdR(self._origPot,r) for r in self._rgrid])
if self._logR:
self._dvcircdrInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid,self._dvcircdrGrid,k=3)
else:
self._dvcircdrInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid,self._dvcircdrGrid,k=3)
if interpepifreq:
from galpy.potential import epifreq
if not numcores is None:
self._epifreqGrid= multi.parallel_map((lambda x: epifreq(self._origPot,self._rgrid[x])),
range(len(self._rgrid)),numcores=numcores)
else:
self._epifreqGrid= numpy.array([epifreq(self._origPot,r) for r in self._rgrid])
indx= True-numpy.isnan(self._epifreqGrid)
if numpy.sum(indx) < 4:
if self._logR:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid[indx],self._epifreqGrid[indx],k=1)
else:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid[indx],self._epifreqGrid[indx],k=1)
else:
if self._logR:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid[indx],self._epifreqGrid[indx],k=3)
else:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid[indx],self._epifreqGrid[indx],k=3)
if interpverticalfreq:
from galpy.potential import verticalfreq
if not numcores is None:
self._verticalfreqGrid= multi.parallel_map((lambda x: verticalfreq(self._origPot,self._rgrid[x])),
range(len(self._rgrid)),numcores=numcores)
else:
self._verticalfreqGrid= numpy.array([verticalfreq(self._origPot,r) for r in self._rgrid])
if self._logR:
self._verticalfreqInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid,self._verticalfreqGrid,k=3)
else:
self._verticalfreqInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid,self._verticalfreqGrid,k=3)
return None
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:
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])
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(
R, z, self._pot) + 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]))
Op.append(omegac(self._pot, axiR[ii]))
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,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:
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])
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(
R, z, self._pot) + 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]))
Op.append(omegac(self._pot, axiR[ii]))
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,
RZPot=None,rgrid=(numpy.log(0.01),numpy.log(20.),101),
zgrid=(0.,1.,101),logR=True,
interpPot=False,interpRforce=False,interpzforce=False,
interpDens=False,
interpvcirc=False,
interpdvcircdr=False,
interpepifreq=False,interpverticalfreq=False,
ro=None,vo=None,
use_c=False,enable_c=False,zsym=True,
numcores=None):
"""
NAME:
__init__
PURPOSE:
Initialize an interpRZPotential instance
INPUT:
RZPot - RZPotential to be interpolated
rgrid - R grid to be given to linspace as in rs= linspace(*rgrid)
zgrid - z grid to be given to linspace as in zs= linspace(*zgrid)
logR - if True, rgrid is in the log of R so logrs= linspace(*rgrid)
interpPot, interpRforce, interpzforce, interpDens,interpvcirc, interpepifreq, interpverticalfreq, interpdvcircdr= if True, interpolate these functions
use_c= use C to speed up the calculation of the grid
enable_c= enable use of C for interpolations
zsym= if True (default), the potential is assumed to be symmetric around z=0 (so you can use, e.g., zgrid=(0.,1.,101)).
numcores= if set to an integer, use this many cores (only used for vcirc, dvcircdR, epifreq, and verticalfreq; NOT NECESSARILY FASTER, TIME TO MAKE SURE)
ro=, vo= distance and velocity scales for translation into internal units (default from configuration file)
OUTPUT:
instance
HISTORY:
2010-07-21 - Written - Bovy (NYU)
2013-01-24 - Started with new implementation - Bovy (IAS)
"""
if isinstance(RZPot,interpRZPotential):
from galpy.potential import PotentialError
raise PotentialError('Cannot setup interpRZPotential with another interpRZPotential')
# Propagate ro and vo
roSet= True
voSet= True
if ro is None:
if isinstance(RZPot,list):
ro= RZPot[0]._ro
roSet= RZPot[0]._roSet
else:
ro= RZPot._ro
roSet= RZPot._roSet
if vo is None:
if isinstance(RZPot,list):
vo= RZPot[0]._vo
voSet= RZPot[0]._voSet
else:
vo= RZPot._vo
voSet= RZPot._voSet
Potential.__init__(self,amp=1.,ro=ro,vo=vo)
# Turn off physical if it hadn't been on
if not roSet: self._roSet= False
if not voSet: self._voSet= False
self._origPot= RZPot
self._rgrid= numpy.linspace(*rgrid)
self._logR= logR
if self._logR:
self._rgrid= numpy.exp(self._rgrid)
self._logrgrid= numpy.log(self._rgrid)
self._zgrid= numpy.linspace(*zgrid)
self._interpPot= interpPot
self._interpRforce= interpRforce
self._interpzforce= interpzforce
self._interpDens= interpDens
self._interpvcirc= interpvcirc
self._interpdvcircdr= interpdvcircdr
self._interpepifreq= interpepifreq
self._interpverticalfreq= interpverticalfreq
self._enable_c= enable_c*ext_loaded
self.hasC= self._enable_c
self._zsym= zsym
if interpPot:
if use_c*ext_loaded:
self._potGrid, err= calc_potential_c(self._origPot,self._rgrid,self._zgrid)
else:
from galpy.potential import evaluatePotentials
potGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
potGrid[ii,jj]= evaluatePotentials(self._origPot,self._rgrid[ii],self._zgrid[jj])
self._potGrid= potGrid
if self._logR:
self._potInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
self._potGrid,
kx=3,ky=3,s=0.)
else:
self._potInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
self._potGrid,
kx=3,ky=3,s=0.)
if enable_c*ext_loaded:
self._potGrid_splinecoeffs= calc_2dsplinecoeffs_c(self._potGrid)
if interpRforce:
if use_c*ext_loaded:
self._rforceGrid, err= calc_potential_c(self._origPot,self._rgrid,self._zgrid,rforce=True)
else:
from galpy.potential import evaluateRforces
rforceGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
rforceGrid[ii,jj]= evaluateRforces(self._origPot,self._rgrid[ii],self._zgrid[jj])
self._rforceGrid= rforceGrid
if self._logR:
self._rforceInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
self._rforceGrid,
kx=3,ky=3,s=0.)
else:
self._rforceInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
self._rforceGrid,
kx=3,ky=3,s=0.)
if enable_c*ext_loaded:
self._rforceGrid_splinecoeffs= calc_2dsplinecoeffs_c(self._rforceGrid)
if interpzforce:
if use_c*ext_loaded:
self._zforceGrid, err= calc_potential_c(self._origPot,self._rgrid,self._zgrid,zforce=True)
else:
from galpy.potential import evaluatezforces
zforceGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
zforceGrid[ii,jj]= evaluatezforces(self._origPot,self._rgrid[ii],self._zgrid[jj])
self._zforceGrid= zforceGrid
if self._logR:
self._zforceInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
self._zforceGrid,
kx=3,ky=3,s=0.)
else:
self._zforceInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
self._zforceGrid,
kx=3,ky=3,s=0.)
if enable_c*ext_loaded:
self._zforceGrid_splinecoeffs= calc_2dsplinecoeffs_c(self._zforceGrid)
if interpDens:
from galpy.potential import evaluateDensities
densGrid= numpy.zeros((len(self._rgrid),len(self._zgrid)))
for ii in range(len(self._rgrid)):
for jj in range(len(self._zgrid)):
densGrid[ii,jj]= evaluateDensities(self._origPot,self._rgrid[ii],self._zgrid[jj])
self._densGrid= densGrid
if self._logR:
self._densInterp= interpolate.RectBivariateSpline(self._logrgrid,
self._zgrid,
numpy.log(self._densGrid+10.**-10.),
kx=3,ky=3,s=0.)
else:
self._densInterp= interpolate.RectBivariateSpline(self._rgrid,
self._zgrid,
numpy.log(self._densGrid+10.**-10.),
kx=3,ky=3,s=0.)
if interpvcirc:
from galpy.potential import vcirc
if not numcores is None:
self._vcircGrid= multi.parallel_map((lambda x: vcirc(self._origPot,self._rgrid[x])),
list(range(len(self._rgrid))),numcores=numcores)
else:
self._vcircGrid= numpy.array([vcirc(self._origPot,r) for r in self._rgrid])
if self._logR:
self._vcircInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid,self._vcircGrid,k=3)
else:
self._vcircInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid,self._vcircGrid,k=3)
if interpdvcircdr:
from galpy.potential import dvcircdR
if not numcores is None:
self._dvcircdrGrid= multi.parallel_map((lambda x: dvcircdR(self._origPot,self._rgrid[x])),
list(range(len(self._rgrid))),numcores=numcores)
else:
self._dvcircdrGrid= numpy.array([dvcircdR(self._origPot,r) for r in self._rgrid])
if self._logR:
self._dvcircdrInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid,self._dvcircdrGrid,k=3)
else:
self._dvcircdrInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid,self._dvcircdrGrid,k=3)
if interpepifreq:
from galpy.potential import epifreq
if not numcores is None:
self._epifreqGrid= numpy.array(multi.parallel_map((lambda x: epifreq(self._origPot,self._rgrid[x])),
list(range(len(self._rgrid))),numcores=numcores))
else:
self._epifreqGrid= numpy.array([epifreq(self._origPot,r) for r in self._rgrid])
indx= True^numpy.isnan(self._epifreqGrid)
if numpy.sum(indx) < 4:
if self._logR:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid[indx],self._epifreqGrid[indx],k=1)
else:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid[indx],self._epifreqGrid[indx],k=1)
else:
if self._logR:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid[indx],self._epifreqGrid[indx],k=3)
else:
self._epifreqInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid[indx],self._epifreqGrid[indx],k=3)
if interpverticalfreq:
from galpy.potential import verticalfreq
if not numcores is None:
self._verticalfreqGrid= multi.parallel_map((lambda x: verticalfreq(self._origPot,self._rgrid[x])),
list(range(len(self._rgrid))),numcores=numcores)
else:
self._verticalfreqGrid= numpy.array([verticalfreq(self._origPot,r) for r in self._rgrid])
if self._logR:
self._verticalfreqInterp= interpolate.InterpolatedUnivariateSpline(self._logrgrid,self._verticalfreqGrid,k=3)
else:
self._verticalfreqInterp= interpolate.InterpolatedUnivariateSpline(self._rgrid,self._verticalfreqGrid,k=3)
return None
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:
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")
