def _get_gravity_accel(x, y, z, galpy_pot=False):
"""
:param x:
:param y:
:param z:
:param galpy_pot:
:return:
"""
r_gal = np.sqrt(x**2 + y**2) / 1e3 # in pc -> kpc
z_gal = z / 1e3 # in pc -> kpc as required by the function
dr_gal = 2.5e-4 # in kpc
if galpy_pot:
dpot = (evaluatePotentials(MWPotential2014, (r_gal + dr_gal) * un.kpc,
z_gal * un.kpc,
phi=0,
t=0,
dR=0) -
evaluatePotentials(MWPotential2014, (r_gal - dr_gal) * un.kpc,
z_gal * un.kpc,
phi=0,
t=0,
dR=0))
kms_to_pcyr = (un.km**2 / un.s**2).to(un.pc**2 / un.yr**2)
dpot *= kms_to_pcyr
else:
dpot = (_galactic_potential(r_gal + dr_gal, z_gal) -
_galactic_potential(r_gal - dr_gal, z_gal))
dpot_dr = dpot / (
2. * dr_gal * 1e3
) # dr_gal converted to pc as dpot is returned in pc2/yr2
return -1. * dpot_dr # pc/yr2
def escape_velocity(self, gal, interpolants):
"""
Calculates the escape velocity for each particle at their respective radius.
"""
print('Calculating particle escape velocities...\n')
# Read in Rvals, assuming particles start in the plane
R_vals = self.R
z_vals = np.zeros_like(R_vals)
# specify radius and height at "infinity" in kpc
R_inf = 1000 * u.kpc
z_inf = 1000 * u.kpc
Vescs = []
for idx, (t0, R, z) in enumerate(zip(self.t0, R_vals, z_vals)):
if interpolants:
potential = interpolants[t0]
else:
potential = gal.full_potentials[t0]
pot_at_inf = evaluatePotentials(potential, R_inf, z_inf).value
Vescs.append(
np.sqrt(
2 *
(pot_at_inf - evaluatePotentials(potential, R, z).value)))
self.Vesc = np.asarray(Vescs) * u.km / u.s
return
def v_esc(pot, r):
"""Calculate the escape velocity at r
Parameters
----------
pot : galpy.potential.Potential or list of Potentials
The potential used to calculate the escape velocity.
r : array_like
Initial position vector of the desired orbit in Galactocentric
cylindrical coordinates, of the form [R, z, phi] in [pc, pc, rad].
Returns
-------
ve : float
The escape velocity at r.
"""
R, z, phi = r
E_min = evaluatePotentials(pot,
R / _pc,
z / _pc,
phi=phi,
use_physical=False)
E_max = evaluatePotentials(pot, 1e+12, 0, phi=phi, use_physical=False)
ve = (2 * (E_max - E_min))**0.5
return ve * _kms
def _evaluate(self,R,z,phi=0.,t=0.,dR=0,dphi=0):
if self._interpPot and self._enable_c:
if isinstance(R,float):
R= numpy.array([R])
if isinstance(z,float):
z= numpy.array([z])
if self._zsym:
return eval_potential_c(self,R,numpy.fabs(z))[0]
else:
return eval_potential_c(self,R,z)[0]
from galpy.potential import evaluatePotentials
if self._interpPot:
if isinstance(R,float):
return self._evaluate(numpy.array([R]),numpy.array([z]))
out= numpy.empty_like(R)
indx= (R >= self._rgrid[0])*(R <= self._rgrid[-1])
if numpy.sum(indx) > 0:
if self._zsym:
if self._logR:
out[indx]= self._potInterp.ev(numpy.log(R[indx]),numpy.fabs(z[indx]))
else:
out[indx]= self._potInterp.ev(R[indx],numpy.fabs(z[indx]))
else:
if self._logR:
out[indx]= self._potInterp.ev(numpy.log(R[indx]),z[indx])
else:
out[indx]= self._potInterp.ev(R[indx],z[indx])
if numpy.sum(True-indx) > 0:
out[True-indx]= evaluatePotentials(R[True-indx],
z[True-indx],
self._origPot)
return out
else:
return evaluatePotentials(R,z,self._origPot)
def Energy(coord, custom_potential=None):
"""
NAME:
Energy
PURPOSE:
Given an array of 6 coordinates for the position and velocity of stars
in Cartesian coordinate, return a list of energy per mass.
Assumes input and out put are in natrual unit.
INPUT:
coord= a numpy array of coordinate
OUTPUT:
energy = a list of total energy per mass
HISTORY:
2018-05-25 - Written - Samuel Wong
2018-07-24 - Changed to an array of points - Samuel Wong
"""
if custom_potential == None:
potential_obj = MWPotential2014
else:
potential_obj = custom_potential
x, y, z, vx, vy, vz = coord.T
R = np.sqrt(x**2 + y**2)
potential = evaluatePotentials(potential_obj, R, z)
kinetic = (vx**2 + vy**2 + vz**2) / 2.
energy = kinetic + potential
return energy
def test_mwpotential2014():
from galpy.potential import (MWPotential2014,
evaluateDensities,
evaluatezforces,
evaluatePotentials)
# Here these have to be default:
ro = 8 * u.kpc
vo = 220 * u.km/u.s
gala_pot = BovyMWPotential2014()
bovy_pot = MWPotential2014
for x in bovy_pot:
x.turn_physical_on()
Rs = np.random.uniform(1, 15, size=ntest) * u.kpc
zs = np.random.uniform(1, 15, size=ntest) * u.kpc
xyz = np.zeros((3, Rs.size)) * u.kpc
xyz[0] = Rs
xyz[2] = zs
assert np.allclose(gala_pot.density(xyz).to_value(u.Msun/u.pc**3),
evaluateDensities(bovy_pot, R=Rs.to_value(ro), z=zs.to_value(ro)))
assert np.allclose(gala_pot.energy(xyz).to_value((u.km / u.s)**2),
evaluatePotentials(bovy_pot, R=Rs.to_value(ro), z=zs.to_value(ro)))
assert np.allclose(gala_pot.gradient(xyz).to_value((u.km/u.s) * u.pc/u.Myr / u.pc)[2],
-evaluatezforces(bovy_pot, R=Rs.to_value(ro), z=zs.to_value(ro)))
def actionsFreqs(self,*args,**kwargs):
"""
NAME:
actionsFreqs
PURPOSE:
evaluate the actions and frequencies (jr,lz,jz,Omegar,Omegaphi,Omegaz)
INPUT:
Either:
a) R,vR,vT,z,vz
b) Orbit instance: initial condition used if that's it, orbit(t)
if there is a time given as well
scipy.integrate.quadrature keywords
OUTPUT:
(jr,lz,jz,Omegar,Omegaphi,Omegaz)
HISTORY:
2013-08-28 - Written - Bovy (IAS)
"""
if ((self._c and not (kwargs.has_key('c') and not kwargs['c']))\
or (ext_loaded and ((kwargs.has_key('c') and kwargs['c'])))) \
and _check_c(self._pot):
if len(args) == 5: #R,vR.vT, z, vz
R,vR,vT, z, vz= args
elif len(args) == 6: #R,vR.vT, z, vz, phi
R,vR,vT, z, vz, phi= args
else:
meta= actionAngle(*args)
R= meta._R
vR= meta._vR
vT= meta._vT
z= meta._z
vz= meta._vz
if isinstance(R,float):
R= nu.array([R])
vR= nu.array([vR])
vT= nu.array([vT])
z= nu.array([z])
vz= nu.array([vz])
Lz= R*vT
if self._useu0:
#First calculate u0
if kwargs.has_key('u0'):
u0= kwargs['u0']
else:
E= nu.array([evaluatePotentials(R[ii],z[ii],self._pot) +vR[ii]**2./2.+vz[ii]**2./2.+vT[ii]**2./2. for ii in range(len(R))])
u0= actionAngleStaeckel_c.actionAngleStaeckel_calcu0(E,Lz,
self._pot,
self._delta)[0]
if kwargs.has_key('u0'): kwargs.pop('u0')
else:
u0= None
jr, jz, Omegar, Omegaphi, Omegaz, err= actionAngleStaeckel_c.actionAngleFreqStaeckel_c(\
self._pot,self._delta,R,vR,vT,z,vz,u0=u0)
if err == 0:
return (jr,Lz,jz,Omegar,Omegaphi,Omegaz)
else:
raise RuntimeError("C-code for calculation actions failed; try with c=False")
else:
if kwargs.has_key('c') and kwargs['c'] and not self._c:
warnings.warn("C module not used because potential does not have a C implementation",galpyWarning)
raise NotImplementedError("actionsFreqs with c=False not implemented")
def test_v_esc(self):
R = z = phi = 1
ve = v_esc(MWPotential2014, [R, z, phi])
orb = Orbit(vxvv=[R/_pc, 0, ve/_kms, z/_pc, 0, phi])
E = orb.E(pot=MWPotential2014)
E_inf = evaluatePotentials(MWPotential2014, 1e+12, 0)
self.assertAlmostEqual(E, E_inf, 10, msg="v_esc returns incorrect "
"escape velocity")
def get_potential_at_point(self,eps,x,y,z):
R=np.sqrt(x.value_in(units.kpc)**2.+y.value_in(units.kpc)**2.)
zed=z.value_in(units.kpc)
phi=np.arctan2(y.value_in(units.kpc),x.value_in(units.kpc))
pot=potential.evaluatePotentials(self.pot,R/self.ro,zed/self.ro,phi=phi,t=self.tgalpy,ro=self.ro,vo=self.vo) | units.km**2*units.s**-2
return pot
def _evaluate(self, R, z, phi=0.0, t=0.0):
from galpy.potential import evaluatePotentials
if self._interpPot:
out = numpy.empty_like(R)
indx = (R >= self._rgrid[0]) * (R <= self._rgrid[-1]) * (z <= self._zgrid[-1]) * (z >= self._zgrid[0])
if numpy.sum(indx) > 0:
if self._enable_c:
out[indx] = eval_potential_c(self, R[indx], z[indx])[0] / self._amp
else:
if self._logR:
out[indx] = self._potInterp.ev(numpy.log(R[indx]), z[indx])
else:
out[indx] = self._potInterp.ev(R[indx], z[indx])
if numpy.sum(True - indx) > 0:
out[True - indx] = evaluatePotentials(R[True - indx], z[True - indx], self._origPot)
return out
else:
return evaluatePotentials(R, z, self._origPot)
def test_vcirc():
#test the circular velocity of the KuzminKutuzovStaeckelPotential
#using parameters from Batsleer & Dejonghe 1994, fig. 1-3
#and their formula eq. (10)
#_____model parameters______
#surface ratios of disk and halo:
ac_Ds = [50. ,50. ,50. ,50. ,50. ,50. ,40. ,40. ,40. ]
ac_Hs = [1.005,1.005,1.005,1.01,1.01,1.01 ,1.02,1.02,1.02 ]
#disk contribution to total mass:
ks = [0.05 ,0.06 ,0.07 ,0.07,0.1 ,0.125,0.1 ,0.12,0.125]
#focal distance:
Delta = 1.
for ii in range(9):
ac_D = ac_Ds[ii]
ac_H = ac_Hs[ii]
k = ks[ii]
#_____setup potential_____
#first try, not normalized:
V_D = KuzminKutuzovStaeckelPotential(amp= k ,ac=ac_D,Delta=Delta,normalize=False)
V_H = KuzminKutuzovStaeckelPotential(amp=(1.-k),ac=ac_H,Delta=Delta,normalize=False)
pot = [V_D,V_H]
#normalization according to Batsleer & Dejonghe 1994:
V00 = evaluatePotentials(pot,0.,0.)
#second try, normalized:
V_D = KuzminKutuzovStaeckelPotential(amp= k / (-V00),ac=ac_D,Delta=Delta,normalize=False)
V_H = KuzminKutuzovStaeckelPotential(amp=(1.-k) / (-V00),ac=ac_H,Delta=Delta,normalize=False)
pot = [V_D,V_H]
#_____calculate rotation curve_____
Rs = numpy.linspace(0.,20.,100)
z = 0.
vcirc_calc = numpy.sqrt(-Rs * evaluateRforces(pot,Rs,z))
#_____vcirc by Batsleer & Dejonghe eq. (10) (with proper Jacobian)_____
def vc2w(R):
g_D = Delta**2 / (1.-ac_D**2)
a_D = g_D - Delta**2
g_H = Delta**2 / (1.-ac_H**2)
a_H = g_H - Delta**2
l = R**2 - a_D
q = a_H - a_D
termD = numpy.sqrt(l )*(numpy.sqrt(l ) + numpy.sqrt(-g_D ))**2
termH = numpy.sqrt(l-q)*(numpy.sqrt(l-q) + numpy.sqrt(-g_D-q))**2
return R**2 * (k / termD + (1.-k) / termH)
vcirc_formula = numpy.sqrt(vc2w(Rs)/(-V00))
assert numpy.all(numpy.fabs(vcirc_calc - vcirc_formula) < 10**-8.), \
'Calculated circular velocity for KuzminKutuzovStaeckelPotential '+ \
'does not agree with eq. (10) (corrected by proper Jacobian) '+ \
'by Batsleer & Dejonghe (1994)'
return None
def _evaluate(self,R,z,phi=0.,t=0.):
from galpy.potential import evaluatePotentials
if self._interpPot:
out= numpy.empty_like(R)
indx= (R >= self._rgrid[0])*(R <= self._rgrid[-1])\
*(z <= self._zgrid[-1])*(z >= self._zgrid[0])
if numpy.sum(indx) > 0:
if self._enable_c:
out[indx]= eval_potential_c(self,R[indx],z[indx])[0]/self._amp
else:
if self._logR:
out[indx]= self._potInterp.ev(numpy.log(R[indx]),z[indx])
else:
out[indx]= self._potInterp.ev(R[indx],z[indx])
if numpy.sum(True^indx) > 0:
out[True^indx]= evaluatePotentials(self._origPot,
R[True^indx],
z[True^indx])
return out
else:
return evaluatePotentials(self._origPot,R,z)
def jz(z,vz,R,pot=MWPotential):
"""
NAME:
jz
PURPOSE:
calculate the vertical action
INPUT:
z - height (/R0)
vz - velocity (/V0)
R - Galactocentric radius (/Ro)
pot= potential to use
OUTPUT:
HISTORY:
2012-03-07 - Written - Bovy (IAS)
"""
#Calculate Ez, /vz for stability
potzero= evaluatePotentials(R,0.,pot)/vz**2.
Ez= 0.5+evaluatePotentials(R,z,pot)/vz**2.-potzero
#Integrate to get Jz
out= 2.*numpy.fabs(vz)/math.pi*integrate.quad(_jzIntegrand,0.,numpy.inf,
args=(R,pot,
Ez,potzero,vz))[0]
#if out == 0.: print Ez, z, vz
return out
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(R, z, pot)
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(R,z,pot)+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(R,z,pot)
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(R, z, pot) + vR**2. / 2. + vT**2. / 2. +
vz**2. / 2., R * vT)
def __call__(self,*args,**kwargs):
"""
NAME:
__call__
PURPOSE:
evaluate the actions (jr,lz,jz)
INPUT:
Either:
a) R,vR,vT,z,vz
b) Orbit instance: initial condition used if that's it, orbit(t)
if there is a time given as well
c= True/False; overrides the object's c= keyword to use C or not
scipy.integrate.quadrature keywords
OUTPUT:
(jr,lz,jz)
HISTORY:
2012-11-27 - Written - Bovy (IAS)
"""
if ((self._c and not (kwargs.has_key('c') and not kwargs['c']))\
or (ext_loaded and ((kwargs.has_key('c') and kwargs['c'])))) \
and _check_c(self._pot):
if len(args) == 5: #R,vR.vT, z, vz
R,vR,vT, z, vz= args
elif len(args) == 6: #R,vR.vT, z, vz, phi
R,vR,vT, z, vz, phi= args
else:
meta= actionAngle(*args)
R= meta._R
vR= meta._vR
vT= meta._vT
z= meta._z
vz= meta._vz
if isinstance(R,float):
R= nu.array([R])
vR= nu.array([vR])
vT= nu.array([vT])
z= nu.array([z])
vz= nu.array([vz])
Lz= R*vT
if self._useu0:
#First calculate u0
if kwargs.has_key('u0'):
u0= kwargs['u0']
else:
E= nu.array([evaluatePotentials(R[ii],z[ii],self._pot) +vR[ii]**2./2.+vz[ii]**2./2.+vT[ii]**2./2. for ii in range(len(R))])
u0= actionAngleStaeckel_c.actionAngleStaeckel_calcu0(E,Lz,
self._pot,
self._delta)[0]
if kwargs.has_key('u0'): kwargs.pop('u0')
else:
u0= None
jr, jz, err= actionAngleStaeckel_c.actionAngleStaeckel_c(\
self._pot,self._delta,R,vR,vT,z,vz,u0=u0)
if err == 0:
return (jr,Lz,jz)
else:
raise RuntimeError("C-code for calculation actions failed; try with c=False")
else:
if kwargs.has_key('c') and kwargs['c'] and not self._c:
warnings.warn("C module not used because potential does not have a C implementation",galpyWarning)
if kwargs.has_key('c'): kwargs.pop('c')
if (len(args) == 5 or len(args) == 6) \
and isinstance(args[0],nu.ndarray):
ojr= nu.zeros((len(args[0])))
olz= nu.zeros((len(args[0])))
ojz= nu.zeros((len(args[0])))
for ii in range(len(args[0])):
if len(args) == 5:
targs= (args[0][ii],args[1][ii],args[2][ii],
args[3][ii],args[4][ii])
elif len(args) == 6:
targs= (args[0][ii],args[1][ii],args[2][ii],
args[3][ii],args[4][ii],args[5][ii])
tjr,tlz,tjz= self(*targs,**copy.copy(kwargs))
ojr[ii]= tjr
ojz[ii]= tjz
olz[ii]= tlz
return (ojr,olz,ojz)
else:
#Set up the actionAngleStaeckelSingle object
aASingle= actionAngleStaeckelSingle(*args,pot=self._pot,
delta=self._delta)
return (aASingle.JR(**copy.copy(kwargs)),
aASingle._R*aASingle._vT,
aASingle.Jz(**copy.copy(kwargs)))
def __init__(self,
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 process_single(i):
vxvv = [ra[i], dec[i], distance[i], pmra[i], pmdec[i], rv[i]]
if not np.all(np.isfinite(vxvv)) or not np.all(np.isfinite(covariance[i])):
return np.ones(56) * np.nan
samp = np.random.multivariate_normal(vxvv, covariance[i], size=1000)
if not (np.all(samp[:, 1] < 90.) & np.all(samp[:, 1] > -90.)):
return np.ones(56) * np.nan
sXYZ, svxyz, sRpz, svRvTvz = obs_to_galcen(samp[:, 0],
samp[:, 1],
samp[:, 2],
samp[:, 3],
samp[:, 4],
samp[:, 5],
ro=_R0,
vo=_v0,
zo=_z0)
deltas = estimateDeltaStaeckel(MWPotential2014, np.median(sRpz[:, 0]), np.median(sRpz[:, 2]), no_median=True)
aAS = actionAngleStaeckel(pot=MWPotential2014, delta=np.mean(deltas))
e, zmax, rperi, rap = aAS.EccZmaxRperiRap(sRpz[:, 0],
svRvTvz[:, 0],
svRvTvz[:, 1],
sRpz[:, 2],
svRvTvz[:, 2],
sRpz[:, 1], delta=deltas)
tcov = np.cov(np.dstack([e, zmax, rperi, rap])[0].T)
errs = np.sqrt(np.diag(tcov))
e_zmax_corr = tcov[0, 1] / (errs[0] * errs[1])
e_rperi_corr = tcov[0, 2] / (errs[0] * errs[2])
e_rap_corr = tcov[0, 3] / (errs[0] * errs[3])
zmax_rperi_corr = tcov[1, 2] / (errs[1] * errs[2])
zmax_rap_corr = tcov[1, 3] / (errs[1] * errs[3])
rperi_rap_corr = tcov[2, 3] / (errs[2] * errs[3])
e_err, zmax_err, rperi_err, rap_err = errs
action = np.array(
aAS.actionsFreqsAngles(sRpz[:, 0], svRvTvz[:, 0], svRvTvz[:, 1], sRpz[:, 2], svRvTvz[:, 2], sRpz[:, 1]))
tcov_after_action = np.cov(action)
errs = np.sqrt(np.diag(tcov_after_action))
jr_lz_corr = tcov_after_action[0, 1] / (errs[0] * errs[1])
jr_jz_corr = tcov_after_action[0, 2] / (errs[0] * errs[2])
lz_jz_corr = tcov_after_action[1, 2] / (errs[1] * errs[2])
jr_err, lz_err, jz_err, or_err, op_err, oz_err, tr_err, tphi_err, tz_err = errs
Rc = np.array([rl(MWPotential2014, lz) for lz in action[1]]) * _R0
Ec = (evaluatePotentials(MWPotential2014, Rc, 0.) + 0.5 * (action[1]) ** 2. / Rc ** 2.) * _v0 ** 2
E = (evaluatePotentials(MWPotential2014, sRpz[:, 0], sRpz[:, 2], phi=sRpz[:, 1]) + np.sum(svRvTvz ** 2 / 2.,
axis=1)) * _v0 ** 2
# galactocentric coord and vel uncertainty
galcen_tcov = np.cov(np.dstack([sRpz[:, 0], sRpz[:, 1], sRpz[:, 2]])[0].T)
galcen_errs = np.sqrt(np.diag(galcen_tcov))
galr_err, galphi_err, galz_err = galcen_errs
galcenv_tcov = np.cov(np.dstack([svRvTvz[:, 0], svRvTvz[:, 1], svRvTvz[:, 2]])[0].T)
galcenv_errs = np.sqrt(np.diag(galcenv_tcov))
galvr_err, galvt_err, galvz_err = galcenv_errs
galvr_galvt_corr = galcenv_tcov[0, 1] / (galcenv_errs[0] * galcenv_errs[1])
galvr_galvz_corr = galcenv_tcov[0, 2] / (galcenv_errs[0] * galcenv_errs[2])
galvt_galvz_corr = galcenv_tcov[1, 2] / (galcenv_errs[1] * galcenv_errs[2])
return np.nanmean(e), e_err, np.nanmean(zmax) * _R0, zmax_err * _R0, np.nanmean(
rperi) * _R0, rperi_err * _R0, np.nanmean(rap) * _R0, rap_err * _R0, \
e_zmax_corr, e_rperi_corr, e_rap_corr, zmax_rperi_corr, zmax_rap_corr, rperi_rap_corr, \
np.nanmean(action[0]) * _R0 * _v0, jr_err * _R0 * _v0, np.nanmean(
action[1]) * _R0 * _v0, lz_err * _R0 * _v0, np.nanmean(action[2]) * _R0 * _v0, jz_err * _R0 * _v0, \
jr_lz_corr, jr_jz_corr, lz_jz_corr, \
np.nanmean(action[3]) * _freq, or_err * _freq, np.nanmean(action[4]) * _freq, op_err * _freq, np.nanmean(
action[5]) * _freq, oz_err * _freq, \
np.nanmean(action[6]), tr_err, np.nanmean(action[7]), tphi_err, np.nanmean(action[8]), tz_err, \
np.nanmean(Rc), np.nanstd(Rc), np.nanmean(E), np.nanstd(E), np.nanmean(E - Ec), np.nanstd(
E - Ec), \
np.nanmean(sRpz[:, 0]) * _R0, \
np.nanmean(sRpz[:, 1]), \
np.nanmean(sRpz[:, 2]) * _R0, \
np.nanmean(svRvTvz[:, 0]) * _v0, \
np.nanmean(svRvTvz[0, 1]) * _v0, \
np.nanmean(svRvTvz[:, 2]) * _v0, \
galr_err * _R0, \
galphi_err, \
galz_err * _R0, \
galvr_err * _v0, \
galvt_err * _v0, \
galvz_err * _v0, \
galvr_galvt_corr, \
galvr_galvz_corr, \
galvt_galvz_corr
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
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:
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 = []
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:
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 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 __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 process_single(i):
vxvv = [ra[i], dec[i], distance[i], pmra[i], pmdec[i], rv[i]]
if not np.all(np.isfinite(vxvv)) or not np.all(np.isfinite(covariance[i])):
return np.ones(56) * np.nan
samp = np.random.multivariate_normal(vxvv, covariance[i], size=1000)
if not (np.all(samp[:, 1] < 90.) & np.all(samp[:, 1] > -90.)):
return np.ones(56) * np.nan
os = Orbit(np.array([
samp[:, 0], samp[:, 1], samp[:, 2], samp[:, 3], samp[:, 4], samp[:, 5]
]).T,
radec=True,
ro=_R0,
vo=_v0,
zo=_z0,
solarmotion=[-11.1, 25.7, 7.25])
sXYZ = np.dstack([os.x(), os.y(), os.z()])[0] / _R0
sRpz = np.dstack([os.R() / _R0, os.phi(), os.z() / _R0])[0]
svRvTvz = np.dstack([os.vR(), os.vT(), os.vz()])[0] / _v0
deltas = estimateDeltaStaeckel(MWPotential2014,
np.median(sRpz[:, 0]),
np.median(sRpz[:, 2]),
no_median=True)
aAS = actionAngleStaeckel(pot=MWPotential2014, delta=np.mean(deltas))
e, zmax, rperi, rap = aAS.EccZmaxRperiRap(sRpz[:, 0],
svRvTvz[:, 0],
svRvTvz[:, 1],
sRpz[:, 2],
svRvTvz[:, 2],
sRpz[:, 1],
delta=deltas)
tcov = np.cov(np.dstack([e, zmax, rperi, rap])[0].T)
errs = np.sqrt(np.diag(tcov))
e_zmax_corr = tcov[0, 1] / (errs[0] * errs[1])
e_rperi_corr = tcov[0, 2] / (errs[0] * errs[2])
e_rap_corr = tcov[0, 3] / (errs[0] * errs[3])
zmax_rperi_corr = tcov[1, 2] / (errs[1] * errs[2])
zmax_rap_corr = tcov[1, 3] / (errs[1] * errs[3])
rperi_rap_corr = tcov[2, 3] / (errs[2] * errs[3])
e_err, zmax_err, rperi_err, rap_err = errs
action = np.array(
aAS.actionsFreqsAngles(sRpz[:, 0], svRvTvz[:, 0], svRvTvz[:, 1],
sRpz[:, 2], svRvTvz[:, 2], sRpz[:, 1]))
tcov_after_action = np.cov(action)
errs = np.sqrt(np.diag(tcov_after_action))
jr_lz_corr = tcov_after_action[0, 1] / (errs[0] * errs[1])
jr_jz_corr = tcov_after_action[0, 2] / (errs[0] * errs[2])
lz_jz_corr = tcov_after_action[1, 2] / (errs[1] * errs[2])
jr_err, lz_err, jz_err, or_err, op_err, oz_err, tr_err, tphi_err, tz_err = errs
Rc = np.array([rl(MWPotential2014, lz) for lz in action[1]]) * _R0
Ec = (evaluatePotentials(MWPotential2014, Rc, 0.) + 0.5 *
(action[1])**2. / Rc**2.) * _v0**2
E = os.E(pot=MWPotential2014)
# galactocentric coord and vel uncertainty
galcen_tcov = np.cov(np.dstack([sRpz[:, 0], sRpz[:, 1], sRpz[:, 2]])[0].T)
galcen_errs = np.sqrt(np.diag(galcen_tcov))
galr_err, galphi_err, galz_err = galcen_errs
galcenv_tcov = np.cov(
np.dstack([svRvTvz[:, 0], svRvTvz[:, 1], svRvTvz[:, 2]])[0].T)
galcenv_errs = np.sqrt(np.diag(galcenv_tcov))
galvr_err, galvt_err, galvz_err = galcenv_errs
galvr_galvt_corr = galcenv_tcov[0, 1] / (galcenv_errs[0] * galcenv_errs[1])
galvr_galvz_corr = galcenv_tcov[0, 2] / (galcenv_errs[0] * galcenv_errs[2])
galvt_galvz_corr = galcenv_tcov[1, 2] / (galcenv_errs[1] * galcenv_errs[2])
# galr mean to avoid issue near GC, error propagation
x_mean = np.nanmean(sXYZ[:, 0])
y_mean = np.nanmean(sXYZ[:, 1])
x_err = np.nanstd(sXYZ[:, 0])
y_err = np.nanstd(sXYZ[:, 1])
x_err_percentage = x_err / np.nanmean(sXYZ[:, 0])
y_err_percentage = y_err / np.nanmean(sXYZ[:, 1])
x2_err = (x_mean**2) * (2 * x_err_percentage)
y2_err = (y_mean**2) * (2 * y_err_percentage)
galr = np.sqrt(x_mean**2 + y_mean**2)
galr2_err = np.sqrt(x2_err**2 + y2_err**2)
galr2_err_percentage = galr2_err / (galr**2)
galr_err = galr * (0.5 * galr2_err_percentage)
# galphi mean to avoid issue near GC, error propagation
galphi = np.arctan2(y_mean, x_mean)
galphi_err_x = (
y_mean /
(x_mean**2 + y_mean**2)) * x_err # error propagation from x_mean
galphi_err_y = (
-x_mean /
(x_mean**2 + y_mean**2)) * y_err # error propagation from y_mean
galphi_err = np.sqrt(galphi_err_x**2 + galphi_err_y**2) # add them up
return np.nanmean(e), \
e_err, \
np.nanmean(zmax) * _R0, \
zmax_err * _R0, \
np.nanmean(rperi) * _R0, \
rperi_err * _R0, \
np.nanmean(rap) * _R0, \
rap_err * _R0, \
e_zmax_corr, \
e_rperi_corr, \
e_rap_corr, \
zmax_rperi_corr, \
zmax_rap_corr, \
rperi_rap_corr, \
np.nanmean(action[0]) * _R0 * _v0, \
jr_err * _R0 * _v0, \
np.nanmean(action[1]) * _R0 * _v0, \
lz_err * _R0 * _v0, \
np.nanmean(action[2]) * _R0 * _v0, \
jz_err * _R0 * _v0, \
jr_lz_corr, \
jr_jz_corr, \
lz_jz_corr, \
np.nanmean(action[3]) * _freq, or_err * _freq, \
np.nanmean(action[4]) * _freq, op_err * _freq, \
np.nanmean(action[5]) * _freq, oz_err * _freq, \
np.nanmean(action[6]), \
tr_err, \
np.nanmean(action[7]), \
tphi_err, \
np.nanmean(action[8]), \
tz_err, \
np.nanmean(Rc), \
np.nanstd(Rc), \
np.nanmean(E), \
np.nanstd(E), \
np.nanmean(E - Ec), \
np.nanstd(E - Ec), \
galr * _R0, \
galphi, \
np.nanmean(sRpz[:, 2]) * _R0, \
np.nanmean(svRvTvz[:, 0]) * _v0, \
np.nanmean(svRvTvz[0, 1]) * _v0, \
np.nanmean(svRvTvz[:, 2]) * _v0, \
galr_err * _R0, \
galphi_err, \
galz_err * _R0, \
galvr_err * _v0, \
galvt_err * _v0, \
galvz_err * _v0, \
galvr_galvt_corr, \
galvr_galvz_corr, \
galvt_galvz_corr
def _jzIntegrand(z,R,pot,Ez,potzero,vz):
vz2= 2.*Ez-evaluatePotentials(R,z,pot)/vz**2.+potzero
if vz2 < 0.: return 0. #Such that we don't have to specify the upper limit
else: return numpy.sqrt(vz2)
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")
def plotForceField(savefig: Optional[str] = None):
"""Plot MW Force Field.
Parameters
----------
savefic: str, optional
"""
p_b15_pal5gd1_voro = fit(fitc=True,
c=None,
addpal5=True,
addgd1=True,
fitvoro=True,
mc16=True)
# Set up potential
p_b15 = p_b15_pal5gd1_voro[0]
ro, vo = REFR0, REFV0
pot = setup_potential(p_b15, p_b15_pal5gd1_voro[0][-1], False, False, ro,
vo)
# Compute force field
Rs = np.linspace(0.01, 20.0, 51)
zs = np.linspace(-20.0, 20.0, 151)
mRs, mzs = np.meshgrid(Rs, zs, indexing="ij")
forces = np.zeros((len(Rs), len(zs), 2))
potvals = np.zeros((len(Rs), len(zs)))
for ii in tqdm(range(len(Rs))):
for jj in tqdm(range(len(zs))):
forces[ii, jj, 0] = potential.evaluateRforces(
pot,
mRs[ii, jj] / ro,
mzs[ii, jj] / ro,
use_physical=True,
ro=ro,
vo=vo,
)
forces[ii, jj, 1] = potential.evaluatezforces(
pot,
mRs[ii, jj] / ro,
mzs[ii, jj] / ro,
use_physical=True,
ro=ro,
vo=vo,
)
potvals[ii, jj] = potential.evaluatePotentials(
pot,
mRs[ii, jj] / ro,
mzs[ii, jj] / ro,
use_physical=True,
ro=ro,
vo=vo,
)
fig = plt.figure(figsize=(8, 16))
skip = 10 # Make sure to keep zs symmetric!!
scale = 35.0
# Don't plot these
# forces[(mRs < 5.)*(np.fabs(mzs) < 4.)]= np.nan
forces[(mRs < 2.0) * (np.fabs(mzs) < 5.0)] = np.nan
bovy_plot.bovy_dens2d(
potvals.T,
origin="lower",
cmap="viridis",
xrange=[Rs[0], Rs[-1]],
yrange=[zs[0], zs[-1]],
xlabel=r"$R\,(\mathrm{kpc})$",
ylabel=r"$Z\,(\mathrm{kpc})$",
contours=True,
aspect=1.0,
)
plt.quiver(
mRs[1::skip, 5:-1:skip],
mzs[1::skip, 5:-1:skip],
forces[1::skip, 5:-1:skip, 0],
forces[1::skip, 5:-1:skip, 1],
scale=scale,
)
# Add a few lines pointing to the GC
for angle in tqdm(np.linspace(0.0, np.pi / 2.0, 8)):
plt.plot((0.0, 100.0 * np.cos(angle)), (0.0, 100.0 * np.sin(angle)),
"k:")
plt.plot((0.0, 100.0 * np.cos(angle)), (0.0, -100.0 * np.sin(angle)),
"k:")
# Add measurements
# Pal 5
plt.quiver((8.0, ), (16.0, ), (-0.8, ), (-1.82, ),
color="w",
zorder=10,
scale=scale)
# GD-1
plt.quiver(
(12.5, ),
(6.675, ),
(-2.51, ),
(-1.47, ),
color="w",
zorder=10,
scale=scale,
)
# Disk + flat APOGEE rotation curve:
# Use Bovy & Tremaine (2012) method for translating F_R in the plane to F_R
# at 1.1 kpc: dFr/dz = dFz / dR
diskrs = np.linspace(5.5, 8.5, 3)
diskfzs = (-67.0 * np.exp(-(diskrs - 8.0) / 2.7) /
bovy_conversion.force_in_2piGmsolpc2(220.0, 8.0) *
bovy_conversion.force_in_kmsMyr(220.0, 8.0))
diskfrs = (-(218.0**2.0 / diskrs) * bovy_conversion._kmsInPcMyr / 1000.0 -
1.1 * diskfzs / 2.7)
plt.quiver(
diskrs,
1.1 * np.ones_like(diskrs),
diskfrs,
diskfzs,
color="w",
zorder=10,
scale=scale,
)
# Labels
bovy_plot.bovy_text(5.8, 16.0, r"$\mathbf{Pal\ 5}$", color="w", size=17.0)
bovy_plot.bovy_text(12.5, 7.0, r"$\mathbf{GD-1}$", color="w", size=17.0)
bovy_plot.bovy_text(8.65,
0.5,
r"$\mathbf{disk\ stars}$",
color="w",
size=17.0)
if savefig is not None:
fig.savefig(savefig)
return fig
def test_vcirc():
#test the circular velocity of the KuzminKutuzovStaeckelPotential
#using parameters from Batsleer & Dejonghe 1994, fig. 1-3
#and their formula eq. (10)
#_____model parameters______
#surface ratios of disk and halo:
ac_Ds = [50., 50., 50., 50., 50., 50., 40., 40., 40.]
ac_Hs = [1.005, 1.005, 1.005, 1.01, 1.01, 1.01, 1.02, 1.02, 1.02]
#disk contribution to total mass:
ks = [0.05, 0.06, 0.07, 0.07, 0.1, 0.125, 0.1, 0.12, 0.125]
#focal distance:
Delta = 1.
for ii in range(9):
ac_D = ac_Ds[ii]
ac_H = ac_Hs[ii]
k = ks[ii]
#_____setup potential_____
#first try, not normalized:
V_D = KuzminKutuzovStaeckelPotential(amp=k,
ac=ac_D,
Delta=Delta,
normalize=False)
V_H = KuzminKutuzovStaeckelPotential(amp=(1. - k),
ac=ac_H,
Delta=Delta,
normalize=False)
pot = [V_D, V_H]
#normalization according to Batsleer & Dejonghe 1994:
V00 = evaluatePotentials(pot, 0., 0.)
#second try, normalized:
V_D = KuzminKutuzovStaeckelPotential(amp=k / (-V00),
ac=ac_D,
Delta=Delta,
normalize=False)
V_H = KuzminKutuzovStaeckelPotential(amp=(1. - k) / (-V00),
ac=ac_H,
Delta=Delta,
normalize=False)
pot = [V_D, V_H]
#_____calculate rotation curve_____
Rs = numpy.linspace(0., 20., 100)
z = 0.
vcirc_calc = numpy.sqrt(-Rs * evaluateRforces(pot, Rs, z))
#_____vcirc by Batsleer & Dejonghe eq. (10) (with proper Jacobian)_____
def vc2w(R):
g_D = Delta**2 / (1. - ac_D**2)
a_D = g_D - Delta**2
g_H = Delta**2 / (1. - ac_H**2)
a_H = g_H - Delta**2
l = R**2 - a_D
q = a_H - a_D
termD = numpy.sqrt(l) * (numpy.sqrt(l) + numpy.sqrt(-g_D))**2
termH = numpy.sqrt(l - q) * (numpy.sqrt(l - q) +
numpy.sqrt(-g_D - q))**2
return R**2 * (k / termD + (1. - k) / termH)
vcirc_formula = numpy.sqrt(vc2w(Rs) / (-V00))
assert numpy.all(numpy.fabs(vcirc_calc - vcirc_formula) < 10**-8.), \
'Calculated circular velocity for KuzminKutuzovStaeckelPotential '+ \
'does not agree with eq. (10) (corrected by proper Jacobian) '+ \
'by Batsleer & Dejonghe (1994)'
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:
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")
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 RK():
r=8e3*pc_in_km #位置の初期値 km
phi=0.0
z = 0.0*pc_in_km #km
vr=vr0 #km/s #速度の初期値
vphi=vphi0
vz=vz0
tmax=46e+8*year_in_sec #繰り返し最大回数
dt=1e+3*year_in_sec #刻み幅
t_list,r_list,phi_list,z_list,vr_list,vphi_list,vz_list,E_list,Lz_list =\
[],[],[],[],[],[],[],[],[]
for i in range(0,int(tmax/dt)):
t = dt*i
k10=dt*fr(r,phi,z,vr,vphi,vz)
k11=dt*fphi(r,phi,z,vr,vphi,vz)
k12=dt*fz(r,phi,z,vr,vphi,vz)
k13=dt*fvr(r,phi,z,vr,vphi,vz)
k14=dt*fvphi(r,phi,z,vr,vphi,vz)
k15=dt*fvz(r,phi,z,vr,vphi,vz)
k20=dt/2.*fr(r+k10/2.,phi+k11/2.,
z+k12/2.,vr+k13/2.,
vphi+k14/2.,vz+k15/2.)
k21=dt/2.*fphi(r+k10/2.,phi+k11/2.,
z+k12/2.,vr+k13/2.,
vphi+k14/2.,vz+k15/2.)
k22=dt/2.*fz(r+k10/2.,phi+k11/2.,
z+k12/2.,vr+k13/2.,
vphi+k14/2.,vz+k15/2.)
k23=dt/2.*fvr(r+k10/2.,phi+k11/2.,
z+k12/2.,vr+k13/2.,
vphi+k14/2.,vz+k15/2.)
k24=dt/2.*fvphi(r+k10/2.,phi+k11/2.,
z+k12/2.,vr+k13/2.,
vphi+k14/2.,vz+k15/2.)
k25=dt/2.*fvz(r+k10/2.,phi+k11/2.,
z+k12/2.,vr+k13/2.,
vphi+k14/2.,vz+k15/2.)
k30=dt/2.*fr(r+k20/2.,phi+k21/2.,
z+k22/2.,vr+k23/2.,
vphi+k24/2.,vz+k25/2.)
k31=dt/2.*fphi(r+k20/2.,phi+k21/2.,
z+k22/2.,vr+k23/2.,
vphi+k24/2.,vz+k25/2.)
k32=dt/2.*fz(r+k20/2.,phi+k21/2.,
z+k22/2.,vr+k23/2.,
vphi+k24/2.,vz+k25/2.)
k33=dt/2.*fvr(r+k20/2.,phi+k21/2.,
z+k22/2.,vr+k23/2.,
vphi+k24/2.,vz+k25/2.)
k34=dt/2.*fvphi(r+k20/2.,phi+k21/2.,
z+k22/2.,vr+k23/2.,
vphi+k24/2.,vz+k25/2.)
k35=dt/2.*fvz(r+k20/2.,phi+k21/2.,
z+k22/2.,vr+k23/2.,
vphi+k24/2.,vz+k25/2.)
k40=dt*fr(r+k30,phi+k31,
z+k32,vr+k33,
vphi+k34,vz+k35)
k41=dt*fphi(r+k30,phi+k31,
z+k32,vr+k33,
vphi+k34,vz+k35)
k42=dt*fz(r+k30,phi+k31,
z+k32,vr+k33,
vphi+k34,vz+k35)
k43=dt*fvr(r+k30,phi+k31,
z+k32,vr+k33,
vphi+k34,vz+k35)
k44=dt*fvphi(r+k30,phi+k31,
z+k32,vr+k33,
vphi+k34,vz+k35)
k45=dt*fvz(r+k30,phi+k31,
z+k32,vr+k33,
vphi+k34,vz+k35)
r+=(k10+2.*k20+2.*k30+k40)/6.0
phi+=(k11+2.*k21+2.*k31+k41)/6.0
z+=(k12+2.*k22+2.*k32+k42)/6.0
vr+=(k13+2.*k23+2.*k33+k43)/6.0
vphi+=(k14+2.*k24+2.*k34+k44)/6.0
vz+=(k15+2.*k25+2.*k35+k45)/6.0
E = 0.5*(vr**2. + vphi**2. + vz**2.) + evaluatePotentials(MWPotential2014,R=r,z=z)
Lz = r*vphi
t_list.append(t)
r_list.append(r)
phi_list.append(phi)
z_list.append(z)
vr_list.append(vr)
vphi_list.append(vphi)
vz_list.append(vz)
E_list.append(E)
Lz_list.append(Lz)
t_array = np.array([t_list])[0,:]
r_array = np.array([r_list])[0,:]
phi_array = np.array([phi_list])[0,:]
z_array = np.array([z_list])[0,:]
vr_array = np.array([vr_list])[0,:]
vphi_array = np.array([vphi_list])[0,:]
vz_array = np.array([vz_list])[0,:]
E_array = np.array([E_list])[0,:]
Lz_array = np.array([Lz_list])[0,:]
data = np.array([t_array/(1e6*year_in_sec),r_array/1e3/pc_in_km,phi_array,\
z_array/pc_in_km,vr_array,vphi_array,vz_array,E_array,Lz_array]).T
np.savetxt('output_MW14.dat', data)