###############################################################################
# diskdf.py: module that interprets (E,Lz) pairs in terms of a
# distribution function (following Dehnen 1999)
#
# This module contains the following classes:
#
# diskdf - top-level class that represents a distribution function
# dehnendf - inherits from diskdf, implements Dehnen's 'new' DF
# shudf - inherits from diskdf, implements Shu's DF
# DFcorrection - class that represents corrections to the input Sigma(R)
# and sigma_R(R) to get closer to the targets
###############################################################################
from __future__ import print_function
_EPSREL=10.**-14.
_NSIGMA= 4.
_INTERPDEGREE= 3
_RMIN=10.**-10.
_MAXD_REJECTLOS= 4.
_PROFILE= False
import copy
import os, os.path
import pickle
import numpy
import scipy
numpylog= numpy.lib.scimath.log # somehow, this code produces log(negative), which scipy (now numpy.lib.scimath.log) implements as log(|negative|) + i pi while numpy gives NaN and we want the scipy behavior; not sure where the log(negative) comes from though! I think it's for sigma=0 DFs (this test fails with numpy.log) where the DF eval has a log(~zero) that can be slightly negative because of numerical precision issues
from scipy import integrate, interpolate, stats, optimize
from .surfaceSigmaProfile import surfaceSigmaProfile, expSurfaceSigmaProfile
from ..orbit import Orbit
from ..util.bovy_ars import bovy_ars
from ..util import save_pickles
from ..util.bovy_conversion import physical_conversion, \
potential_physical_input, _APY_UNITS, surfdens_in_msolpc2
from ..potential import PowerSphericalPotential
from ..actionAngle import actionAngleAdiabatic, actionAngleAxi
from .df import df, _APY_LOADED
if _APY_LOADED:
from astropy import units
#scipy version
from pkg_resources import parse_version
_SCIPY_VERSION= parse_version(scipy.__version__)
_SCIPY_VERSION_BREAK= parse_version('0.9')
_CORRECTIONSDIR=os.path.join(os.path.dirname(os.path.realpath(__file__)),'data')
_DEGTORAD= numpy.pi/180.
class diskdf(df):
"""Class that represents a disk DF"""
def __init__(self,dftype='dehnen',
surfaceSigma=expSurfaceSigmaProfile,
profileParams=(1./3.,1.0,0.2),
correct=False,
beta=0.,ro=None,vo=None,**kwargs):
"""
NAME:
__init__
PURPOSE:
Initialize a DF
INPUT:
dftype= 'dehnen' or 'corrected-dehnen', 'shu' or 'corrected-shu'
surfaceSigma - instance or class name of the target
surface density and sigma_R profile
(default: both exponential)
profileParams - parameters of the surface and sigma_R profile:
(xD,xS,Sro) where
xD - disk surface mass scalelength / Ro
xS - disk velocity dispersion scalelength / Ro
Sro - disk velocity dispersion at Ro (/vo)
Directly given to the 'surfaceSigmaProfile class, so
could be anything that class takes
beta - power-law index of the rotation curve
correct - correct the DF (i.e., DFcorrection kwargs are also given)
+ DFcorrection kwargs (except for those already specified)
OUTPUT:
HISTORY:
2010-03-10 - Written - Bovy (NYU)
"""
df.__init__(self,ro=ro,vo=vo)
self._dftype= dftype
if isinstance(surfaceSigma,surfaceSigmaProfile):
self._surfaceSigmaProfile= surfaceSigma
else:
if _APY_LOADED and isinstance(profileParams[0],units.Quantity):
newprofileParams=\
(profileParams[0].to(units.kpc).value/self._ro,
profileParams[1].to(units.kpc).value/self._ro,
profileParams[2].to(units.km/units.s).value/self._vo)
self._roSet= True
self._voSet= True
profileParams= newprofileParams
self._surfaceSigmaProfile= surfaceSigma(profileParams)
self._beta= beta
self._gamma= numpy.sqrt(2./(1.+self._beta))
if correct or 'corrections' in kwargs or 'rmax' in kwargs \
or 'niter' in kwargs or 'npoints' in kwargs:
self._correct= True
#Load corrections
self._corr= DFcorrection(dftype=self.__class__,
surfaceSigmaProfile=self._surfaceSigmaProfile,
beta=beta,**kwargs)
else:
self._correct= False
#Setup aA objects for frequency and rap,rperi calculation
self._aA= actionAngleAdiabatic(pot=PowerSphericalPotential(normalize=1.,
alpha=2.-2.*self._beta),gamma=0.)
return None
[docs] @physical_conversion('phasespacedensity2d',pop=True)
def __call__(self,*args,**kwargs):
"""
NAME:
__call__
PURPOSE:
evaluate the distribution function
INPUT:
either an orbit instance, a list of such instances, or E,Lz
1) Orbit instance or list:
a) Orbit instance alone: use initial condition
b) Orbit instance + t: call the Orbit instance (for list, each instance is called at t)
2)
E - energy (/vo^2; or can be Quantity)
L - angular momentun (/ro/vo; or can be Quantity)
3) array vxvv [3/4,nt] [must be in natural units /vo,/ro; use Orbit interface for physical-unit input)
KWARGS:
marginalizeVperp - marginalize over perpendicular velocity (only supported with 1a) for single orbits above)
marginalizeVlos - marginalize over line-of-sight velocity (only supported with 1a) for single orbits above)
nsigma= number of sigma to integrate over when marginalizing
+scipy.integrate.quad keywords
OUTPUT:
DF(orbit/E,L)
HISTORY:
2010-07-10 - Written - Bovy (NYU)
"""
if isinstance(args[0],Orbit):
if len(args[0]) > 1:
raise RuntimeError('Only single-object Orbit instances can be passed to DF instances at this point') #pragma: no cover
if len(args) == 1:
if kwargs.pop('marginalizeVperp',False):
return self._call_marginalizevperp(args[0],**kwargs)
elif kwargs.pop('marginalizeVlos',False):
return self._call_marginalizevlos(args[0],**kwargs)
else:
return numpy.real(self.eval(*vRvTRToEL(\
args[0].vR(use_physical=False),
args[0].vT(use_physical=False),
args[0].R(use_physical=False),
self._beta,self._dftype)))
else:
no= args[0](args[1])
return numpy.real(self.eval(*vRvTRToEL(no.vR(use_physical=False),
no.vT(use_physical=False),
no.R(use_physical=False),
self._beta,
self._dftype)))
elif isinstance(args[0],list) \
and isinstance(args[0][0],Orbit):
if numpy.any([len(no) > 1 for no in args[0]]):
raise RuntimeError('Only single-object Orbit instances can be passed to DF instances at this point') #pragma: no cover
#Grab all of the vR, vT, and R
vR= numpy.array([o.vR(use_physical=False) for o in args[0]])
vT= numpy.array([o.vT(use_physical=False) for o in args[0]])
R= numpy.array([o.R(use_physical=False) for o in args[0]])
return numpy.real(self.eval(*vRvTRToEL(vR,vT,R,self._beta,
self._dftype)))
elif isinstance(args[0],numpy.ndarray) and \
not (hasattr(args[0],'isscalar') and args[0].isscalar):
#Grab all of the vR, vT, and R
vR= args[0][1]
vT= args[0][2]
R= args[0][0]
return numpy.real(self.eval(*vRvTRToEL(vR,vT,R,self._beta,
self._dftype)))
else:
return numpy.real(self.eval(*args))
def _call_marginalizevperp(self,o,**kwargs):
"""Call the DF, marginalizing over perpendicular velocity"""
#Get l, vlos
l= o.ll(obs=[1.,0.,0.],ro=1.)*_DEGTORAD
vlos= o.vlos(ro=1.,vo=1.,obs=[1.,0.,0.,0.,0.,0.])
R= o.R(use_physical=False)
phi= o.phi(use_physical=False)
#Get local circular velocity, projected onto the los
vcirc= R**self._beta
vcirclos= vcirc*numpy.sin(phi+l)
#Marginalize
alphalos= phi+l
if not 'nsigma' in kwargs or ('nsigma' in kwargs and \
kwargs['nsigma'] is None):
nsigma= _NSIGMA
else:
nsigma= kwargs['nsigma']
kwargs.pop('nsigma',None)
sigmaR2= self.targetSigma2(R,use_physical=False)
sigmaR1= numpy.sqrt(sigmaR2)
#Use the asymmetric drift equation to estimate va
va= sigmaR2/2./R**self._beta*(1./self._gamma**2.-1.
-R*self._surfaceSigmaProfile.surfacemassDerivative(R,log=True)
-R*self._surfaceSigmaProfile.sigma2Derivative(R,log=True))
if numpy.fabs(va) > sigmaR1: va = 0. #To avoid craziness near the center
if numpy.fabs(numpy.sin(alphalos)) < numpy.sqrt(1./2.):
cosalphalos= numpy.cos(alphalos)
tanalphalos= numpy.tan(alphalos)
return integrate.quad(_marginalizeVperpIntegrandSinAlphaSmall,
-self._gamma*va/sigmaR1-nsigma,
-self._gamma*va/sigmaR1+nsigma,
args=(self,R,cosalphalos,tanalphalos,
vlos-vcirclos,vcirc,
sigmaR1/self._gamma),
**kwargs)[0]/numpy.fabs(cosalphalos)\
*sigmaR1/self._gamma
else:
sinalphalos= numpy.sin(alphalos)
cotalphalos= 1./numpy.tan(alphalos)
return integrate.quad(_marginalizeVperpIntegrandSinAlphaLarge,
-nsigma,nsigma,
args=(self,R,sinalphalos,cotalphalos,
vlos-vcirclos,vcirc,sigmaR1),
**kwargs)[0]/numpy.fabs(sinalphalos)*sigmaR1
def _call_marginalizevlos(self,o,**kwargs):
"""Call the DF, marginalizing over line-of-sight velocity"""
#Get d, l, vperp
l= o.ll(obs=[1.,0.,0.],ro=1.)*_DEGTORAD
vperp= o.vll(ro=1.,vo=1.,obs=[1.,0.,0.,0.,0.,0.])
R= o.R(use_physical=False)
phi= o.phi(use_physical=False)
#Get local circular velocity, projected onto the perpendicular
#direction
vcirc= R**self._beta
vcircperp= vcirc*numpy.cos(phi+l)
#Marginalize
alphaperp= numpy.pi/2.+phi+l
if not 'nsigma' in kwargs or ('nsigma' in kwargs and \
kwargs['nsigma'] is None):
nsigma= _NSIGMA
else:
nsigma= kwargs['nsigma']
kwargs.pop('nsigma',None)
sigmaR2= self.targetSigma2(R,use_physical=False)
sigmaR1= numpy.sqrt(sigmaR2)
#Use the asymmetric drift equation to estimate va
va= sigmaR2/2./R**self._beta*(1./self._gamma**2.-1.
-R*self._surfaceSigmaProfile.surfacemassDerivative(R,log=True)
-R*self._surfaceSigmaProfile.sigma2Derivative(R,log=True))
if numpy.fabs(va) > sigmaR1: va = 0. #To avoid craziness near the center
if numpy.fabs(numpy.sin(alphaperp)) < numpy.sqrt(1./2.):
cosalphaperp= numpy.cos(alphaperp)
tanalphaperp= numpy.tan(alphaperp)
#we can reuse the VperpIntegrand, since it is just another angle
return integrate.quad(_marginalizeVperpIntegrandSinAlphaSmall,
-self._gamma*va/sigmaR1-nsigma,
-self._gamma*va/sigmaR1+nsigma,
args=(self,R,cosalphaperp,tanalphaperp,
vperp-vcircperp,vcirc,
sigmaR1/self._gamma),
**kwargs)[0]/numpy.fabs(cosalphaperp)\
*sigmaR1/self._gamma
else:
sinalphaperp= numpy.sin(alphaperp)
cotalphaperp= 1./numpy.tan(alphaperp)
#we can reuse the VperpIntegrand, since it is just another angle
return integrate.quad(_marginalizeVperpIntegrandSinAlphaLarge,
-nsigma,nsigma,
args=(self,R,sinalphaperp,cotalphaperp,
vperp-vcircperp,vcirc,sigmaR1),
**kwargs)[0]/numpy.fabs(sinalphaperp)*sigmaR1
[docs] @potential_physical_input
@physical_conversion('velocity2',pop=True)
def targetSigma2(self,R,log=False):
"""
NAME:
targetSigma2
PURPOSE:
evaluate the target Sigma_R^2(R)
INPUT:
R - radius at which to evaluate (can be Quantity)
OUTPUT:
target Sigma_R^2(R)
log - if True, return the log (default: False)
HISTORY:
2010-03-28 - Written - Bovy (NYU)
"""
return self._surfaceSigmaProfile.sigma2(R,log=log)
[docs] @potential_physical_input
@physical_conversion('surfacedensity',pop=True)
def targetSurfacemass(self,R,log=False):
"""
NAME:
targetSurfacemass
PURPOSE:
evaluate the target surface mass at R
INPUT:
R - radius at which to evaluate (can be Quantity)
log - if True, return the log (default: False)
OUTPUT:
Sigma(R)
HISTORY:
2010-03-28 - Written - Bovy (NYU)
"""
return self._surfaceSigmaProfile.surfacemass(R,log=log)
[docs] @physical_conversion('surfacedensitydistance',pop=True)
def targetSurfacemassLOS(self,d,l,log=False,deg=True):
"""
NAME:
targetSurfacemassLOS
PURPOSE:
evaluate the target surface mass along the LOS given l and d
INPUT:
d - distance along the line of sight (can be Quantity)
l - Galactic longitude (in deg, unless deg=False; can be Quantity)
deg= if False, l is in radians
log - if True, return the log (default: False)
OUTPUT:
Sigma(d,l) x d
HISTORY:
2011-03-23 - Written - Bovy (NYU)
"""
#Calculate R and phi
if _APY_LOADED and isinstance(l,units.Quantity):
lrad= l.to(units.rad).value
elif deg:
lrad= l*_DEGTORAD
else:
lrad= l
if _APY_LOADED and isinstance(d,units.Quantity):
d= d.to(units.kpc).value/self._ro
R, phi= _dlToRphi(d,lrad)
if log:
return self._surfaceSigmaProfile.surfacemass(R,log=log)\
+numpylog(d)
else:
return self._surfaceSigmaProfile.surfacemass(R,log=log)\
*d
[docs] @physical_conversion('surfacedensitydistance',pop=True)
def surfacemassLOS(self,d,l,deg=True,target=True,
romberg=False,nsigma=None,relative=None):
"""
NAME:
surfacemassLOS
PURPOSE:
evaluate the surface mass along the LOS given l and d
INPUT:
d - distance along the line of sight (can be Quantity)
l - Galactic longitude (in deg, unless deg=False; can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
target= if True, use target surfacemass (default)
romberg - if True, use a romberg integrator (default: False)
deg= if False, l is in radians
OUTPUT:
Sigma(d,l) x d
HISTORY:
2011-03-24 - Written - Bovy (NYU)
"""
#Calculate R and phi
if _APY_LOADED and isinstance(l,units.Quantity):
lrad= l.to(units.rad).value
elif deg:
lrad= l*_DEGTORAD
else:
lrad= l
if _APY_LOADED and isinstance(d,units.Quantity):
d= d.to(units.kpc).value/self._ro
R, phi= _dlToRphi(d,lrad)
if target:
if relative: return d
else: return self.targetSurfacemass(R,use_physical=False)*d
else:
return self.surfacemass(R,romberg=romberg,nsigma=nsigma,
relative=relative,use_physical=False)\
*d
[docs] @physical_conversion('position',pop=True)
def sampledSurfacemassLOS(self,l,n=1,maxd=None,target=True):
"""
NAME:
sampledSurfacemassLOS
PURPOSE:
sample a distance along the line of sight
INPUT:
l - Galactic longitude (in rad; can be Quantity)
n= number of distances to sample
maxd= maximum distance to consider (for the rejection sampling) (can be Quantity)
target= if True, sample from the 'target' surface mass density, rather than the actual surface mass density (default=True)
OUTPUT:
list of samples
HISTORY:
2011-03-24 - Written - Bovy (NYU)
"""
#First calculate where the maximum is
if target:
minR= optimize.fmin_bfgs(lambda x: \
-self.targetSurfacemassLOS(x,l,
use_physical=False,
deg=False),
0.,disp=False)[0]
maxSM= self.targetSurfacemassLOS(minR,l,deg=False,
use_physical=False)
else:
minR= optimize.fmin_bfgs(lambda x: \
-self.surfacemassLOS(x,l,
deg=False,
use_physical=False),
0.,disp=False)[0]
maxSM= self.surfacemassLOS(minR,l,deg=False,use_physical=False)
#Now rejection-sample
if _APY_LOADED and isinstance(l,units.Quantity):
l= l.to(units.rad).value
if _APY_LOADED and isinstance(maxd,units.Quantity):
maxd= maxd.to(units.kpc).value/self._ro
if maxd is None:
maxd= _MAXD_REJECTLOS
out= []
while len(out) < n:
#sample
prop= numpy.random.random()*maxd
if target:
surfmassatprop= self.targetSurfacemassLOS(prop,l,deg=False,
use_physical=False)
else:
surfmassatprop= self.surfacemassLOS(prop,l,deg=False,
use_physical=False)
if surfmassatprop/maxSM > numpy.random.random(): #accept
out.append(prop)
return numpy.array(out)
[docs] @potential_physical_input
@physical_conversion('velocity',pop=True)
def sampleVRVT(self,R,n=1,nsigma=None,target=True):
"""
NAME:
sampleVRVT
PURPOSE:
sample a radial and azimuthal velocity at R
INPUT:
R - Galactocentric distance (can be Quantity)
n= number of distances to sample
nsigma= number of sigma to rejection-sample on
target= if True, sample using the 'target' sigma_R rather than the actual sigma_R (default=True)
OUTPUT:
list of samples
BUGS:
should use the fact that vR and vT separate
HISTORY:
2011-03-24 - Written - Bovy (NYU)
"""
#Determine where the max of the v-distribution is using asymmetric drift
maxVR= 0.
maxVT= optimize.brentq(_vtmaxEq,0.,R**self._beta+0.2,(R,self))
maxVD= self(Orbit([R,maxVR,maxVT]))
#Now rejection-sample
if nsigma == None:
nsigma= _NSIGMA
out= []
if target:
sigma= numpy.sqrt(self.targetSigma2(R,use_physical=False))
else:
sigma= numpy.sqrt(self.sigma2(R,use_physical=False))
while len(out) < n:
#sample
vrg, vtg= numpy.random.normal(), numpy.random.normal()
propvR= vrg*nsigma*sigma
propvT= vtg*nsigma*sigma/self._gamma+maxVT
VDatprop= self(Orbit([R,propvR,propvT]))
if VDatprop/maxVD > numpy.random.uniform()*numpy.exp(-0.5*(vrg**2.+vtg**2.)): #accept
out.append(numpy.array([propvR,propvT]))
return numpy.array(out)
[docs] def sampleLOS(self,los,n=1,deg=True,maxd=None,nsigma=None,
targetSurfmass=True,targetSigma2=True):
"""
NAME:
sampleLOS
PURPOSE:
sample along a given LOS
INPUT:
los - line of sight (in deg, unless deg=False; can be Quantity)
n= number of desired samples
deg= los in degrees? (default=True)
targetSurfmass, targetSigma2= if True, use target surface mass and sigma2 profiles, respectively (there is not much point to doing the latter)
(default=True)
OUTPUT:
returns list of Orbits
BUGS:
target=False uses target distribution for derivatives (this is a detail)
HISTORY:
2011-03-24 - Started - Bovy (NYU)
"""
if _APY_LOADED and isinstance(los,units.Quantity):
l= los.to(units.rad).value
elif deg:
l= los*_DEGTORAD
else:
l= los
out= []
#sample distances
ds= self.sampledSurfacemassLOS(l,n=n,maxd=maxd,target=targetSurfmass,
use_physical=False)
for ii in range(int(n)):
#Calculate R and phi
thisR,thisphi= _dlToRphi(ds[ii],l)
#sample velocities
vv= self.sampleVRVT(thisR,n=1,nsigma=nsigma,target=targetSigma2,
use_physical=False)[0]
if self._roSet and self._voSet:
out.append(Orbit([thisR,vv[0],vv[1],thisphi],ro=self._ro,
vo=self._vo))
else:
out.append(Orbit([thisR,vv[0],vv[1],thisphi]))
return out
[docs] @potential_physical_input
@physical_conversion('velocity',pop=True)
def asymmetricdrift(self,R):
"""
NAME:
asymmetricdrift
PURPOSE:
estimate the asymmetric drift (vc-mean-vphi) from an approximation to the Jeans equation
INPUT:
R - radius at which to calculate the asymmetric drift (can be Quantity)
OUTPUT:
asymmetric drift at R
HISTORY:
2011-04-02 - Written - Bovy (NYU)
"""
sigmaR2= self.targetSigma2(R,use_physical=False)
return sigmaR2/2./R**self._beta*(1./self._gamma**2.-1.
-R*self._surfaceSigmaProfile.surfacemassDerivative(R,log=True)
-R*self._surfaceSigmaProfile.sigma2Derivative(R,log=True))
[docs] @potential_physical_input
@physical_conversion('surfacedensity',pop=True)
def surfacemass(self,R,romberg=False,nsigma=None,relative=False):
"""
NAME:
surfacemass
PURPOSE:
calculate the surface-mass at R by marginalizing over velocity
INPUT:
R - radius at which to calculate the surfacemass density (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
surface mass at R
HISTORY:
2010-03-XX - Written - Bovy (NYU)
"""
if nsigma == None:
nsigma= _NSIGMA
logSigmaR= self.targetSurfacemass(R,log=True,use_physical=False)
sigmaR2= self.targetSigma2(R,use_physical=False)
sigmaR1= numpy.sqrt(sigmaR2)
logsigmaR2= numpylog(sigmaR2)
if relative:
norm= 1.
else:
norm= numpy.exp(logSigmaR)
#Use the asymmetric drift equation to estimate va
va= sigmaR2/2./R**self._beta*(1./self._gamma**2.-1.
-R*self._surfaceSigmaProfile.surfacemassDerivative(R,log=True)
-R*self._surfaceSigmaProfile.sigma2Derivative(R,log=True))
if numpy.fabs(va) > sigmaR1: va = 0.#To avoid craziness near the center
if romberg:
return numpy.real(bovy_dblquad(_surfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: 0., lambda x: nsigma,
[R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma],
tol=10.**-8)/numpy.pi*norm)
else:
return integrate.dblquad(_surfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: 0., lambda x: nsigma,
(R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma),
epsrel=_EPSREL)[0]/numpy.pi*norm
[docs] @potential_physical_input
@physical_conversion('velocity2surfacedensity',pop=True)
def sigma2surfacemass(self,R,romberg=False,nsigma=None,
relative=False):
"""
NAME:
sigma2surfacemass
PURPOSE:
calculate the product sigma_R^2 x surface-mass at R by marginalizing over velocity
INPUT:
R - radius at which to calculate the sigma_R^2 x surfacemass density (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
sigma_R^2 x surface-mass at R
HISTORY:
2010-03-XX - Written - Bovy (NYU)
"""
if nsigma == None:
nsigma= _NSIGMA
logSigmaR= self.targetSurfacemass(R,log=True,use_physical=False)
sigmaR2= self.targetSigma2(R,use_physical=False)
sigmaR1= numpy.sqrt(sigmaR2)
logsigmaR2= numpylog(sigmaR2)
if relative:
norm= 1.
else:
norm= numpy.exp(logSigmaR+logsigmaR2)
#Use the asymmetric drift equation to estimate va
va= sigmaR2/2./R**self._beta*(1./self._gamma**2.-1.
-R*self._surfaceSigmaProfile.surfacemassDerivative(R,log=True)
-R*self._surfaceSigmaProfile.sigma2Derivative(R,log=True))
if numpy.fabs(va) > sigmaR1: va = 0. #To avoid craziness near the center
if romberg:
return numpy.real(bovy_dblquad(_sigma2surfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: 0., lambda x: nsigma,
[R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma],
tol=10.**-8)/numpy.pi*norm)
else:
return integrate.dblquad(_sigma2surfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: 0., lambda x: nsigma,
(R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma),
epsrel=_EPSREL)[0]/numpy.pi*norm
def vmomentsurfacemass(self,*args,**kwargs):
"""
NAME:
vmomentsurfacemass
PURPOSE:
calculate the an arbitrary moment of the velocity distribution
at R times the surfacmass
INPUT:
R - radius at which to calculate the moment (in natural units)
n - vR^n
m - vT^m
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
deriv= None, 'R', or 'phi': calculates derivative of the moment wrt R or phi
OUTPUT:
<vR^n vT^m x surface-mass> at R (no support for units)
HISTORY:
2011-03-30 - Written - Bovy (NYU)
"""
use_physical= kwargs.pop('use_physical',True)
ro= kwargs.pop('ro',None)
if ro is None and hasattr(self,'_roSet') and self._roSet:
ro= self._ro
if _APY_LOADED and isinstance(ro,units.Quantity):
ro= ro.to(units.kpc).value
vo= kwargs.pop('vo',None)
if vo is None and hasattr(self,'_voSet') and self._voSet:
vo= self._vo
if _APY_LOADED and isinstance(vo,units.Quantity):
vo= vo.to(units.km/units.s).value
if use_physical and not vo is None and not ro is None:
fac= surfdens_in_msolpc2(vo,ro)*vo**(args[1]+args[2])
if _APY_UNITS:
u= units.Msun/units.pc**2*(units.km/units.s)**(args[1]+args[2])
out= self._vmomentsurfacemass(*args,**kwargs)
if _APY_UNITS:
return units.Quantity(out*fac,unit=u)
else:
return out*fac
else:
return self._vmomentsurfacemass(*args,**kwargs)
[docs] def _vmomentsurfacemass(self,R,n,m,romberg=False,nsigma=None,
relative=False,phi=0.,deriv=None):
"""Non-physical version of vmomentsurfacemass, otherwise the same"""
#odd moments of vR are zero
if isinstance(n,int) and n%2 == 1:
return 0.
if nsigma == None:
nsigma= _NSIGMA
logSigmaR= self.targetSurfacemass(R,log=True,use_physical=False)
sigmaR2= self.targetSigma2(R,use_physical=False)
sigmaR1= numpy.sqrt(sigmaR2)
logsigmaR2= numpylog(sigmaR2)
if relative:
norm= 1.
else:
norm= numpy.exp(logSigmaR+logsigmaR2*(n+m)/2.)/self._gamma**m
#Use the asymmetric drift equation to estimate va
va= sigmaR2/2./R**self._beta*(1./self._gamma**2.-1.
-R*self._surfaceSigmaProfile.surfacemassDerivative(R,log=True)
-R*self._surfaceSigmaProfile.sigma2Derivative(R,log=True))
if numpy.fabs(va) > sigmaR1: va = 0. #To avoid craziness near the center
if deriv is None:
if romberg:
return numpy.real(bovy_dblquad(_vmomentsurfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: -nsigma, lambda x: nsigma,
[R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma,n,m],
tol=10.**-8)/numpy.pi*norm/2.)
else:
return integrate.dblquad(_vmomentsurfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: -nsigma, lambda x: nsigma,
(R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma,n,m),
epsrel=_EPSREL)[0]/numpy.pi*norm/2.
else:
if romberg:
return numpy.real(bovy_dblquad(_vmomentderivsurfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: -nsigma, lambda x: nsigma,
[R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma,n,m,deriv],
tol=10.**-8)/numpy.pi*norm/2.)
else:
return integrate.dblquad(_vmomentderivsurfaceIntegrand,
self._gamma*(R**self._beta-va)/sigmaR1-nsigma,
self._gamma*(R**self._beta-va)/sigmaR1+nsigma,
lambda x: -nsigma, lambda x: nsigma,
(R,self,logSigmaR,logsigmaR2,sigmaR1,
self._gamma,n,m,deriv),
epsrel=_EPSREL)[0]/numpy.pi*norm/2.
[docs] @potential_physical_input
@physical_conversion('frequency-kmskpc',pop=True)
def oortA(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
oortA
PURPOSE:
calculate the Oort function A
INPUT:
R - radius at which to calculate A (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
Oort A at R
HISTORY:
2011-04-19 - Written - Bovy (NYU)
BUGS:
could be made more efficient, e.g., surfacemass is calculated multiple times
"""
#2A= meanvphi/R-dmeanvR/R/dphi-dmeanvphi/dR
meanvphi= self.meanvT(R,romberg=romberg,nsigma=nsigma,phi=phi,
use_physical=False)
dmeanvRRdphi= 0. #We know this, since the DF does not depend on phi
surfmass= self._vmomentsurfacemass(R,0,0,phi=phi,romberg=romberg,nsigma=nsigma)
dmeanvphidR= self._vmomentsurfacemass(R,0,1,deriv='R',phi=phi,romberg=romberg,nsigma=nsigma)/\
surfmass\
-self._vmomentsurfacemass(R,0,1,phi=phi,romberg=romberg,nsigma=nsigma)\
/surfmass**2.\
*self._vmomentsurfacemass(R,0,0,deriv='R',phi=phi,romberg=romberg,nsigma=nsigma)
return 0.5*(meanvphi/R-dmeanvRRdphi/R-dmeanvphidR)
[docs] @potential_physical_input
@physical_conversion('frequency-kmskpc',pop=True)
def oortB(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
oortB
PURPOSE:
calculate the Oort function B
INPUT:
R - radius at which to calculate B (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
Oort B at R
HISTORY:
2011-04-19 - Written - Bovy (NYU)
BUGS:
could be made more efficient, e.g., surfacemass is calculated multiple times
"""
#2B= -meanvphi/R+dmeanvR/R/dphi-dmeanvphi/dR
meanvphi= self.meanvT(R,romberg=romberg,nsigma=nsigma,phi=phi,
use_physical=False)
dmeanvRRdphi= 0. #We know this, since the DF does not depend on phi
surfmass= self._vmomentsurfacemass(R,0,0,phi=phi,romberg=romberg,nsigma=nsigma)
dmeanvphidR= self._vmomentsurfacemass(R,0,1,deriv='R',phi=phi,romberg=romberg,nsigma=nsigma)/\
surfmass\
-self._vmomentsurfacemass(R,0,1,phi=phi,romberg=romberg,nsigma=nsigma)\
/surfmass**2.\
*self._vmomentsurfacemass(R,0,0,deriv='R',phi=phi,romberg=romberg,nsigma=nsigma)
return 0.5*(-meanvphi/R+dmeanvRRdphi/R-dmeanvphidR)
[docs] @potential_physical_input
@physical_conversion('frequency-kmskpc',pop=True)
def oortC(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
oortC
PURPOSE:
calculate the Oort function C
INPUT:
R - radius at which to calculate C (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
Oort C at R
HISTORY:
2011-04-19 - Written - Bovy (NYU)
BUGS:
could be made more efficient, e.g., surfacemass is calculated multiple times
we know this is zero, but it is calculated anyway (bug or feature?)
"""
#2C= -meanvR/R-dmeanvphi/R/dphi+dmeanvR/dR
meanvr= self.meanvR(R,romberg=romberg,nsigma=nsigma,phi=phi,
use_physical=False)
dmeanvphiRdphi= 0. #We know this, since the DF does not depend on phi
surfmass= self._vmomentsurfacemass(R,0,0,phi=phi,romberg=romberg,nsigma=nsigma)
dmeanvRdR= self._vmomentsurfacemass(R,1,0,deriv='R',phi=phi,romberg=romberg,nsigma=nsigma)/\
surfmass #other terms is zero because f is even in vR
return 0.5*(-meanvr/R-dmeanvphiRdphi/R+dmeanvRdR)
[docs] @potential_physical_input
@physical_conversion('frequency-kmskpc',pop=True)
def oortK(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
oortK
PURPOSE:
calculate the Oort function K
INPUT:
R - radius at which to calculate K (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
Oort K at R
HISTORY:
2011-04-19 - Written - Bovy (NYU)
BUGS:
could be made more efficient, e.g., surfacemass is calculated multiple times
we know this is zero, but it is calculated anyway (bug or feature?)
"""
#2K= meanvR/R+dmeanvphi/R/dphi+dmeanvR/dR
meanvr= self.meanvR(R,romberg=romberg,nsigma=nsigma,phi=phi,
use_physical=False)
dmeanvphiRdphi= 0. #We know this, since the DF does not depend on phi
surfmass= self._vmomentsurfacemass(R,0,0,phi=phi,romberg=romberg,nsigma=nsigma)
dmeanvRdR= self._vmomentsurfacemass(R,1,0,deriv='R',phi=phi,romberg=romberg,nsigma=nsigma)/\
surfmass #other terms is zero because f is even in vR
return 0.5*(+meanvr/R+dmeanvphiRdphi/R+dmeanvRdR)
[docs] @potential_physical_input
@physical_conversion('velocity2',pop=True)
def sigma2(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
sigma2
PURPOSE:
calculate sigma_R^2 at R by marginalizing over velocity
INPUT:
R - radius at which to calculate sigma_R^2 density (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
sigma_R^2 at R
HISTORY:
2010-03-XX - Written - Bovy (NYU)
"""
return self.sigma2surfacemass(R,romberg,nsigma,use_physical=False)\
/self.surfacemass(R,romberg,nsigma,use_physical=False)
[docs] @potential_physical_input
@physical_conversion('velocity2',pop=True)
def sigmaT2(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
sigmaT2
PURPOSE:
calculate sigma_T^2 at R by marginalizing over velocity
INPUT:
R - radius at which to calculate sigma_T^2 (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
sigma_T^2 at R
HISTORY:
2011-03-30 - Written - Bovy (NYU)
"""
surfmass= self.surfacemass(R,romberg=romberg,nsigma=nsigma,
use_physical=False)
return (self._vmomentsurfacemass(R,0,2,romberg=romberg,nsigma=nsigma)
-self._vmomentsurfacemass(R,0,1,romberg=romberg,nsigma=nsigma)\
**2.\
/surfmass)/surfmass
[docs] @potential_physical_input
@physical_conversion('velocity2',pop=True)
def sigmaR2(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
sigmaR2 (duplicate of sigma2 for consistency)
PURPOSE:
calculate sigma_R^2 at R by marginalizing over velocity
INPUT:
R - radius at which to calculate sigma_R^2 (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
sigma_R^2 at R
HISTORY:
2011-03-30 - Written - Bovy (NYU)
"""
return self.sigma2(R,romberg=romberg,nsigma=nsigma,use_physical=False)
[docs] @potential_physical_input
@physical_conversion('velocity',pop=True)
def meanvT(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
meanvT
PURPOSE:
calculate <vT> at R by marginalizing over velocity
INPUT:
R - radius at which to calculate <vT> (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
<vT> at R
HISTORY:
2011-03-30 - Written - Bovy (NYU)
"""
return self._vmomentsurfacemass(R,0,1,romberg=romberg,nsigma=nsigma)\
/self.surfacemass(R,romberg=romberg,nsigma=nsigma,
use_physical=False)
[docs] @potential_physical_input
@physical_conversion('velocity',pop=True)
def meanvR(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
meanvR
PURPOSE:
calculate <vR> at R by marginalizing over velocity
INPUT:
R - radius at which to calculate <vR> (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
<vR> at R
HISTORY:
2011-03-30 - Written - Bovy (NYU)
"""
return self._vmomentsurfacemass(R,1,0,romberg=romberg,nsigma=nsigma)\
/self.surfacemass(R,romberg=romberg,nsigma=nsigma,
use_physical=False)
[docs] @potential_physical_input
def skewvT(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
skewvT
PURPOSE:
calculate skew in vT at R by marginalizing over velocity
INPUT:
R - radius at which to calculate <vR> (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
skewvT
HISTORY:
2011-12-07 - Written - Bovy (NYU)
"""
surfmass= self.surfacemass(R,romberg=romberg,nsigma=nsigma,
use_physical=False)
vt= self._vmomentsurfacemass(R,0,1,romberg=romberg,nsigma=nsigma)\
/surfmass
vt2= self._vmomentsurfacemass(R,0,2,romberg=romberg,nsigma=nsigma)\
/surfmass
vt3= self._vmomentsurfacemass(R,0,3,romberg=romberg,nsigma=nsigma)\
/surfmass
s2= vt2-vt**2.
return (vt3-3.*vt*vt2+2.*vt**3.)*s2**(-1.5)
[docs] @potential_physical_input
def skewvR(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
skewvR
PURPOSE:
calculate skew in vR at R by marginalizing over velocity
INPUT:
R - radius at which to calculate <vR> (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
skewvR
HISTORY:
2011-12-07 - Written - Bovy (NYU)
"""
surfmass= self.surfacemass(R,romberg=romberg,nsigma=nsigma,
use_physical=False)
vr= self._vmomentsurfacemass(R,1,0,romberg=romberg,nsigma=nsigma)\
/surfmass
vr2= self._vmomentsurfacemass(R,2,0,romberg=romberg,nsigma=nsigma)\
/surfmass
vr3= self._vmomentsurfacemass(R,3,0,romberg=romberg,nsigma=nsigma)\
/surfmass
s2= vr2-vr**2.
return (vr3-3.*vr*vr2+2.*vr**3.)*s2**(-1.5)
[docs] @potential_physical_input
def kurtosisvT(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
kurtosisvT
PURPOSE:
calculate excess kurtosis in vT at R by marginalizing over velocity
INPUT:
R - radius at which to calculate <vR> (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
kurtosisvT
HISTORY:
2011-12-07 - Written - Bovy (NYU)
"""
surfmass= self.surfacemass(R,romberg=romberg,nsigma=nsigma,
use_physical=False)
vt= self._vmomentsurfacemass(R,0,1,romberg=romberg,nsigma=nsigma)\
/surfmass
vt2= self._vmomentsurfacemass(R,0,2,romberg=romberg,nsigma=nsigma)\
/surfmass
vt3= self._vmomentsurfacemass(R,0,3,romberg=romberg,nsigma=nsigma)\
/surfmass
vt4= self._vmomentsurfacemass(R,0,4,romberg=romberg,nsigma=nsigma)\
/surfmass
s2= vt2-vt**2.
return (vt4-4.*vt*vt3+6.*vt**2.*vt2-3.*vt**4.)*s2**(-2.)-3.
[docs] @potential_physical_input
def kurtosisvR(self,R,romberg=False,nsigma=None,phi=0.):
"""
NAME:
kurtosisvR
PURPOSE:
calculate excess kurtosis in vR at R by marginalizing over velocity
INPUT:
R - radius at which to calculate <vR> (can be Quantity)
OPTIONAL INPUT:
nsigma - number of sigma to integrate the velocities over
KEYWORDS:
romberg - if True, use a romberg integrator (default: False)
OUTPUT:
kurtosisvR
HISTORY:
2011-12-07 - Written - Bovy (NYU)
"""
surfmass= self.surfacemass(R,romberg=romberg,nsigma=nsigma,
use_physical=False)
vr= self._vmomentsurfacemass(R,1,0,romberg=romberg,nsigma=nsigma)\
/surfmass
vr2= self._vmomentsurfacemass(R,2,0,romberg=romberg,nsigma=nsigma)\
/surfmass
vr3= self._vmomentsurfacemass(R,3,0,romberg=romberg,nsigma=nsigma)\
/surfmass
vr4= self._vmomentsurfacemass(R,4,0,romberg=romberg,nsigma=nsigma)\
/surfmass
s2= vr2-vr**2.
return (vr4-4.*vr*vr3+6.*vr**2.*vr2-3.*vr**4.)*s2**(-2.)-3.
def _ELtowRRapRperi(self,E,L):
"""
NAME:
_ELtowRRapRperi
PURPOSE:
calculate the radial frequency based on E,L, also return rap and
rperi
INPUT:
E - energy
L - angular momentum
OUTPUT:
wR(E.L)
HISTORY:
2010-07-11 - Written - Bovy (NYU)
"""
if self._beta == 0.:
xE= numpy.exp(E-.5)
else: #non-flat rotation curve
xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
rperi,rap= self._aA.calcRapRperi(xE,0.,L/xE,0.,0.)
#Replace the above w/
aA= actionAngleAxi(xE,0.,L/xE,
pot=PowerSphericalPotential(normalize=1.,
alpha=2.-2.*self._beta).toPlanar())
TR= aA.TR()
return (2.*numpy.pi/TR,rap,rperi)
[docs] def sample(self,n=1,rrange=None,returnROrbit=True,returnOrbit=False,
nphi=1.,los=None,losdeg=True,nsigma=None,maxd=None,target=True):
"""
NAME:
sample
PURPOSE:
sample n*nphi points from this DF
INPUT:
n - number of desired sample (specifying this rather than calling this routine n times is more efficient)
rrange - if you only want samples in this rrange, set this keyword (only works when asking for an (RZ)Orbit) (can be Quantity)
returnROrbit - if True, return a planarROrbit instance:
[R,vR,vT] (default)
returnOrbit - if True, return a planarOrbit instance (including phi)
nphi - number of azimuths to sample for each E,L
los= line of sight sampling along this line of sight (can be Quantity)
losdeg= los in degrees? (default=True)
target= if True, use target surface mass and sigma2 profiles (default=True)
nsigma= number of sigma to rejection-sample on
maxd= maximum distance to consider (for the rejection sampling)
OUTPUT:
n*nphi list of [[E,Lz],...] or list of planar(R)Orbits
CAUTION: lists of EL need to be post-processed to account for the
\kappa/\omega_R discrepancy
HISTORY:
2010-07-10 - Started - Bovy (NYU)
"""
raise NotImplementedError("'sample' method for this disk df is not implemented")
def _estimatemeanvR(self,R,phi=0.,log=False):
"""
NAME:
_estimatemeanvR
PURPOSE:
quickly estimate meanvR (useful in evolveddiskdf where we
need an estimate of this but we do not want to spend too
much time on it)
INPUT:
R - radius at which to evaluate (/ro)
phi= azimuth (not used)
OUTPUT:
target Sigma_R^2(R)
log - if True, return the log (default: False)
HISTORY:
2010-03-28 - Written - Bovy (NYU)
"""
return 0.
def _estimatemeanvT(self,R,phi=0.,log=False):
"""
NAME:
_estimatemeanvT
PURPOSE:
quickly estimate meanvR (useful in evolveddiskdf where we
need an estimate of this but we do not want to spend too
much time on it)
INPUT:
R - radius at which to evaluate (/ro)
phi= azimuth (not used)
OUTPUT:
target Sigma_R^2(R)
HISTORY:
2010-03-28 - Written - Bovy (NYU)
"""
return R**self._beta-self.asymmetricdrift(R,use_physical=False)
def _estimateSigmaR2(self,R,phi=0.,log=False):
"""
NAME:
_estimateSigmaR2
PURPOSE:
quickly estimate SigmaR2 (useful in evolveddiskdf where we
need an estimate of this but we do not want to spend too
much time on it)
INPUT:
R - radius at which to evaluate (/ro)
phi= azimuth (not used)
OUTPUT:
target Sigma_R^2(R)
log - if True, return the log (default: False)
HISTORY:
2010-03-28 - Written - Bovy (NYU)
"""
return self.targetSigma2(R,log=log,use_physical=False)
def _estimateSigmaT2(self,R,phi=0.,log=False):
"""
NAME:
_estimateSigmaT2
PURPOSE:
quickly estimate SigmaT2 (useful in evolveddiskdf where we
need an estimate of this but we do not want to spend too
much time on it)
INPUT:
R - radius at which to evaluate (/ro)
phi= azimuth (not used)
OUTPUT:
target Sigma_R^2(R)
log - if True, return the log (default: False)
HISTORY:
2010-03-28 - Written - Bovy (NYU)
"""
if log:
return self.targetSigma2(R,log=log,use_physical=False)\
-2.*numpylog(self._gamma)
else:
return self.targetSigma2(R,log=log,use_physical=False)\
/self._gamma**2.
[docs]class dehnendf(diskdf):
"""Dehnen's 'new' df"""
[docs] def __init__(self,surfaceSigma=expSurfaceSigmaProfile,
profileParams=(1./3.,1.0,0.2),
correct=False,
beta=0.,**kwargs):
"""
NAME:
__init__
PURPOSE:
Initialize a Dehnen 'new' DF
INPUT:
surfaceSigma - instance or class name of the target
surface density and sigma_R profile
(default: both exponential)
profileParams - parameters of the surface and sigma_R profile:
(xD,xS,Sro) where
xD - disk surface mass scalelength (can be Quantity)
xS - disk velocity dispersion scalelength (can be Quantity)
Sro - disk velocity dispersion at Ro (can be Quantity)
Directly given to the 'surfaceSigmaProfile class, so
could be anything that class takes
beta - power-law index of the rotation curve
correct - if True, correct the DF
ro= distance from vantage point to GC (kpc; can be Quantity)
vo= circular velocity at ro (km/s; can be Quantity)
+DFcorrection kwargs (except for those already specified)
OUTPUT:
instance
HISTORY:
2010-03-10 - Written - Bovy (NYU)
"""
return diskdf.__init__(self,surfaceSigma=surfaceSigma,
profileParams=profileParams,
correct=correct,dftype='dehnen',
beta=beta,**kwargs)
def eval(self,E,L,logSigmaR=0.,logsigmaR2=0.):
"""
NAME:
eval
PURPOSE:
evaluate the distribution function
INPUT:
E - energy (can be Quantity)
L - angular momentum (can be Quantity)
OUTPUT:
DF(E,L)
HISTORY:
2010-03-10 - Written - Bovy (NYU)
2010-03-28 - Moved to dehnenDF - Bovy (NYU)
"""
if _PROFILE: #pragma: no cover
import time
start= time.time()
if _APY_LOADED and isinstance(E,units.Quantity):
E= E.to(units.km**2/units.s**2).value/self._vo**2.
if _APY_LOADED and isinstance(L,units.Quantity):
L= L.to(units.kpc*units.km/units.s).value/self._ro/self._vo
#Calculate Re,LE, OmegaE
if self._beta == 0.:
xE= numpy.exp(E-.5)
logOLLE= numpylog(L/xE-1.)
else: #non-flat rotation curve
xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
logOLLE= self._beta*numpylog(xE)+numpylog(L/xE-xE**self._beta)
if _PROFILE: #pragma: no cover
one_time= (time.time()-start)
start= time.time()
if self._correct:
correction= self._corr.correct(xE,log=True)
else:
correction= numpy.zeros(2)
if _PROFILE: #pragma: no cover
corr_time= (time.time()-start)
start= time.time()
SRE2= self.targetSigma2(xE,log=True,use_physical=False)+correction[1]
if _PROFILE: #pragma: no cover
targSigma_time= (time.time()-start)
start= time.time()
out= self._gamma*numpy.exp(logsigmaR2-SRE2+self.targetSurfacemass(xE,log=True,use_physical=False)-logSigmaR+numpy.exp(logOLLE-SRE2)+correction[0])/2./numpy.pi
out_time= (time.time()-start)
tot_time= one_time+corr_time+targSigma_time+out_time
print(one_time/tot_time, corr_time/tot_time, targSigma_time/tot_time, out_time/tot_time, tot_time)
return out
else:
return self._gamma*numpy.exp(logsigmaR2-SRE2+self.targetSurfacemass(xE,log=True,use_physical=False)-logSigmaR+numpy.exp(logOLLE-SRE2)+correction[0])/2./numpy.pi
def sample(self,n=1,rrange=None,returnROrbit=True,returnOrbit=False,
nphi=1.,los=None,losdeg=True,nsigma=None,targetSurfmass=True,
targetSigma2=True,
maxd=None,**kwargs):
"""
NAME:
sample
PURPOSE:
sample n*nphi points from this DF
INPUT:
n - number of desired sample (specifying this rather than calling
this routine n times is more efficient)
rrange - if you only want samples in this rrange, set this keyword
(only works when asking for an (RZ)Orbit
returnROrbit - if True, return a planarROrbit instance:
[R,vR,vT] (default)
returnOrbit - if True, return a planarOrbit instance (including phi)
nphi - number of azimuths to sample for each E,L
los= if set, sample along this line of sight (deg) (assumes that the Sun is located at R=1,phi=0)
losdeg= if False, los is in radians (default=True)
targetSurfmass, targetSigma2= if True, use target surface mass and sigma2 profiles, respectively (there is not much point to doing the latter)
(default=True)
nsigma= number of sigma to rejection-sample on
maxd= maximum distance to consider (for the rejection sampling)
OUTPUT:
n*nphi list of [[E,Lz],...] or list of planar(R)Orbits
CAUTION: lists of EL need to be post-processed to account for the
\kappa/\omega_R discrepancy; EL not returned in physical units
HISTORY:
2010-07-10 - Started - Bovy (NYU)
"""
if not los is None:
return self.sampleLOS(los,deg=losdeg,n=n,maxd=maxd,
nsigma=nsigma,targetSurfmass=targetSurfmass,
targetSigma2=targetSigma2)
#First sample xE
if self._correct:
xE= numpy.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
_ars_hpx,nsamples=n,
hxparams=(self._surfaceSigmaProfile,
self._corr)))
else:
xE= numpy.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
_ars_hpx,nsamples=n,
hxparams=(self._surfaceSigmaProfile,
None)))
#Calculate E
if self._beta == 0.:
E= numpylog(xE)+0.5
else: #non-flat rotation curve
E= .5*xE**(2.*self._beta)*(1.+1./self._beta)
#Then sample Lz
LCE= xE**(self._beta+1.)
OR= xE**(self._beta-1.)
Lz= self._surfaceSigmaProfile.sigma2(xE)*numpylog(stats.uniform.rvs(size=n))/OR
if self._correct:
Lz*= self._corr.correct(xE,log=False)[1,:]
Lz+= LCE
if not returnROrbit and not returnOrbit:
out= [[e,l] for e,l in zip(E,Lz)]
else:
if not rrange is None \
and _APY_LOADED and isinstance(rrange[0],units.Quantity):
rrange[0]= rrange[0].to(units.kpc).value/self._ro
rrange[1]= rrange[1].to(units.kpc).value/self._ro
if not hasattr(self,'_psp'):
self._psp= PowerSphericalPotential(alpha=2.-self._beta,normalize=True).toPlanar()
out= []
for ii in range(int(n)):
try:
wR, rap, rperi= self._ELtowRRapRperi(E[ii],Lz[ii])
except ValueError:
continue
TR= 2.*numpy.pi/wR
tr= stats.uniform.rvs()*TR
if tr > TR/2.:
tr-= TR/2.
thisOrbit= Orbit([rperi,0.,Lz[ii]/rperi])
else:
thisOrbit= Orbit([rap,0.,Lz[ii]/rap])
thisOrbit.integrate(numpy.array([0.,tr]),self._psp)
if returnOrbit:
vxvv= thisOrbit(tr).vxvv[0]
thisOrbit= Orbit(vxvv=numpy.array([vxvv[0],vxvv[1],vxvv[2],
stats.uniform.rvs()\
*numpy.pi*2.])\
.reshape(4))
else:
thisOrbit= thisOrbit(tr)
kappa= _kappa(thisOrbit.vxvv[0,0],self._beta)
if not rrange == None:
if thisOrbit.vxvv[0,0] < rrange[0] \
or thisOrbit.vxvv[0,0] > rrange[1]:
continue
mult= numpy.ceil(kappa/wR*nphi)-1.
kappawR= kappa/wR*nphi-mult
while mult > 0:
if returnOrbit:
out.append(Orbit(vxvv=numpy.array([vxvv[0],vxvv[1],
vxvv[2],
stats.uniform.rvs()*numpy.pi*2.]).reshape(4)))
else:
out.append(thisOrbit)
mult-= 1
if stats.uniform.rvs() > kappawR:
continue
out.append(thisOrbit)
#Recurse to get enough
if len(out) < n*nphi:
out.extend(self.sample(n=int(n-len(out)/nphi),rrange=rrange,
returnROrbit=returnROrbit,
returnOrbit=returnOrbit,nphi=int(nphi),
los=los,losdeg=losdeg))
if len(out) > n*nphi:
print(n, nphi, n*nphi)
out= out[0:int(n*nphi)]
if kwargs.get('use_physical',True) and \
self._roSet and self._voSet:
if isinstance(out[0],Orbit):
dum= [o.turn_physical_on(ro=self._ro,vo=self._vo) for o in out]
return out
def _dlnfdR(self,R,vR,vT):
#Calculate a bunch of stuff that we need
if self._beta == 0.:
E= vR**2./2.+vT**2./2.+numpylog(R)
xE= numpy.exp(E-.5)
OE= xE**-1.
LCE= xE
dRedR= xE/R
else: #non-flat rotation curve
E= vR**2./2.+vT**2./2.+1./2./self._beta*R**(2.*self._beta)
xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
OE= xE**(self._beta-1.)
LCE= xE**(self._beta+1.)
dRedR= xE/2./self._beta/E*R**(2.*self._beta-1.)
return self._dlnfdRe(R,vR,vT,E=E,xE=xE,OE=OE,LCE=LCE)*dRedR\
+self._dlnfdl(R,vR,vT,E=E,xE=xE,OE=OE)*vT
def _dlnfdvR(self,R,vR,vT):
#Calculate a bunch of stuff that we need
if self._beta == 0.:
E= vR**2./2.+vT**2./2.+numpylog(R)
xE= numpy.exp(E-.5)
OE= xE**-1.
LCE= xE
dRedvR= xE*vR
else: #non-flat rotation curve
E= vR**2./2.+vT**2./2.+1./2./self._beta*R**(2.*self._beta)
xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
OE= xE**(self._beta-1.)
LCE= xE**(self._beta+1.)
dRedvR= xE/2./self._beta/E*vR
return self._dlnfdRe(R,vR,vT,E=E,xE=xE,OE=OE,LCE=LCE)*dRedvR
def _dlnfdvT(self,R,vR,vT):
#Calculate a bunch of stuff that we need
if self._beta == 0.:
E= vR**2./2.+vT**2./2.+numpylog(R)
xE= numpy.exp(E-.5)
OE= xE**-1.
LCE= xE
dRedvT= xE*vT
else: #non-flat rotation curve
E= vR**2./2.+vT**2./2.+1./2./self._beta*R**(2.*self._beta)
xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
OE= xE**(self._beta-1.)
LCE= xE**(self._beta+1.)
dRedvT= xE/2./self._beta/E*vT
return self._dlnfdRe(R,vR,vT,E=E,xE=xE,OE=OE,LCE=LCE)*dRedvT\
+self._dlnfdl(R,vR,vT,E=E,xE=xE,OE=OE)*R
def _dlnfdRe(self,R,vR,vT,E=None,xE=None,OE=None,LCE=None):
"""d ln f(x,v) / d R_e"""
#Calculate a bunch of stuff that we need
if E is None or xE is None or OE is None or LCE is None:
if self._beta == 0.:
E= vR**2./2.+vT**2./2.+numpylog(R)
xE= numpy.exp(E-.5)
OE= xE**-1.
LCE= xE
else: #non-flat rotation curve
E= vR**2./2.+vT**2./2.+1./2./self._beta*R**(2.*self._beta)
xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
OE= xE**(self._beta-1.)
LCE= xE**(self._beta+1.)
L= R*vT
sigma2xE= self._surfaceSigmaProfile.sigma2(xE,log=False)
return (self._surfaceSigmaProfile.surfacemassDerivative(xE,log=True)\
-(1.+OE*(L-LCE)/sigma2xE)*self._surfaceSigmaProfile.sigma2Derivative(xE,log=True)\
+(L-LCE)/sigma2xE*(self._beta-1.)*xE**(self._beta-2.)\
-OE*(self._beta+1.)/sigma2xE*xE**self._beta)
def _dlnfdl(self,R,vR,vT,E=None,xE=None,OE=None):
#Calculate a bunch of stuff that we need
if E is None or xE is None or OE is None:
if self._beta == 0.:
E= vR**2./2.+vT**2./2.+numpylog(R)
xE= numpy.exp(E-.5)
OE= xE**-1.
else: #non-flat rotation curve
E= vR**2./2.+vT**2./2.+1./2./self._beta*R**(2.*self._beta)
xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
OE= xE**(self._beta-1.)
sigma2xE= self._surfaceSigmaProfile.sigma2(xE,log=False)
return OE/sigma2xE
[docs]class shudf(diskdf):
"""Shu's df (1969)"""
[docs] def __init__(self,surfaceSigma=expSurfaceSigmaProfile,
profileParams=(1./3.,1.0,0.2),
correct=False,
beta=0.,**kwargs):
"""
NAME:
__init__
PURPOSE:
Initialize a Shu DF
INPUT:
surfaceSigma - instance or class name of the target
surface density and sigma_R profile
(default: both exponential)
profileParams - parameters of the surface and sigma_R profile:
(xD,xS,Sro) where
xD - disk surface mass scalelength (can be Quantity)
xS - disk velocity dispersion scalelength (can be Quantity)
Sro - disk velocity dispersion at Ro (can be Quantity)
Directly given to the 'surfaceSigmaProfile class, so
could be anything that class takes
beta - power-law index of the rotation curve
correct - if True, correct the DF
ro= distance from vantage point to GC (kpc; can be Quantity)
vo= circular velocity at ro (km/s; can be Quantity)
+DFcorrection kwargs (except for those already specified)
OUTPUT:
instance
HISTORY:
2010-05-09 - Written - Bovy (NYU)
"""
return diskdf.__init__(self,surfaceSigma=surfaceSigma,
profileParams=profileParams,
correct=correct,dftype='shu',
beta=beta,**kwargs)
def eval(self,E,L,logSigmaR=0.,logsigmaR2=0.):
"""
NAME:
eval
PURPOSE:
evaluate the distribution function
INPUT:
E - energy (/vo^2)
L - angular momentun (/ro/vo)
OUTPUT:
DF(E,L)
HISTORY:
2010-05-09 - Written - Bovy (NYU)
"""
if _APY_LOADED and isinstance(E,units.Quantity):
E= E.to(units.km**2/units.s**2).value/self._vo**2.
if _APY_LOADED and isinstance(L,units.Quantity):
L= L.to(units.kpc*units.km/units.s).value/self._ro/self._vo
#Calculate RL,LL, OmegaL
if self._beta == 0.:
xL= L
logECLE= numpylog(-numpylog(xL)-0.5+E)
else: #non-flat rotation curve
xL= L**(1./(self._beta+1.))
logECLE= numpylog(-0.5*(1./self._beta+1.)*xL**(2.*self._beta)+E)
if xL < 0.: #We must remove counter-rotating mass
return 0.
if self._correct:
correction= self._corr.correct(xL,log=True)
else:
correction= numpy.zeros(2)
SRE2= self.targetSigma2(xL,log=True,use_physical=False)+correction[1]
return self._gamma*numpy.exp(logsigmaR2-SRE2+self.targetSurfacemass(xL,log=True,use_physical=False)-logSigmaR-numpy.exp(logECLE-SRE2)+correction[0])/2./numpy.pi
def sample(self,n=1,rrange=None,returnROrbit=True,returnOrbit=False,
nphi=1.,los=None,losdeg=True,nsigma=None,maxd=None,
targetSurfmass=True,targetSigma2=True,**kwargs):
"""
NAME:
sample
PURPOSE:
sample n*nphi points from this DF
INPUT:
n - number of desired sample (specifying this rather than calling
this routine n times is more efficient)
rrange - if you only want samples in this rrange, set this keyword
(only works when asking for an (RZ)Orbit
returnROrbit - if True, return a planarROrbit instance:
[R,vR,vT] (default)
returnOrbit - if True, return a planarOrbit instance (including phi)
nphi - number of azimuths to sample for each E,L
los= if set, sample along this line of sight (deg) (assumes that the Sun is located at R=1,phi=0)
losdeg= if False, los is in radians (default=True)
targetSurfmass, targetSigma2= if True, use target surface mass and sigma2 profiles, respectively (there is not much point to doing the latter)
(default=True)
nsigma= number of sigma to rejection-sample on
maxd= maximum distance to consider (for the rejection sampling)
OUTPUT:
n*nphi list of [[E,Lz],...] or list of planar(R)Orbits
CAUTION: lists of EL need to be post-processed to account for the
\kappa/\omega_R discrepancy
HISTORY:
2010-07-10 - Started - Bovy (NYU)
"""
if not los is None:
return self.sampleLOS(los,n=n,maxd=maxd,
nsigma=nsigma,targetSurfmass=targetSurfmass,
targetSigma2=targetSigma2)
#First sample xL
if self._correct:
xL= numpy.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
_ars_hpx,nsamples=n,
hxparams=(self._surfaceSigmaProfile,
self._corr)))
else:
xL= numpy.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
_ars_hpx,nsamples=n,
hxparams=(self._surfaceSigmaProfile,
None)))
#Calculate Lz
Lz= xL**(self._beta+1.)
#Then sample E
if self._beta == 0.:
ECL= numpylog(xL)+0.5
else:
ECL= 0.5*(1./self._beta+1.)*xL**(2.*self._beta)
E= -self._surfaceSigmaProfile.sigma2(xL)*numpylog(stats.uniform.rvs(size=n))
if self._correct:
E*= self._corr.correct(xL,log=False)[1,:]
E+= ECL
if not returnROrbit and not returnOrbit:
out= [[e,l] for e,l in zip(E,Lz)]
else:
if not rrange is None \
and _APY_LOADED and isinstance(rrange[0],units.Quantity):
rrange[0]= rrange[0].to(units.kpc).value/self._ro
rrange[1]= rrange[1].to(units.kpc).value/self._ro
if not hasattr(self,'_psp'):
self._psp= PowerSphericalPotential(alpha=2.-self._beta,normalize=True).toPlanar()
out= []
for ii in range(n):
try:
wR, rap, rperi= self._ELtowRRapRperi(E[ii],Lz[ii])
except ValueError: #pragma: no cover
continue
TR= 2.*numpy.pi/wR
tr= stats.uniform.rvs()*TR
if tr > TR/2.:
tr-= TR/2.
thisOrbit= Orbit([rperi,0.,Lz[ii]/rperi])
else:
thisOrbit= Orbit([rap,0.,Lz[ii]/rap])
thisOrbit.integrate(numpy.array([0.,tr]),self._psp)
if returnOrbit:
vxvv= thisOrbit(tr).vxvv[0]
thisOrbit= Orbit(vxvv=numpy.array([vxvv[0],vxvv[1],vxvv[2],
stats.uniform.rvs()*numpy.pi*2.]).reshape(4))
else:
thisOrbit= thisOrbit(tr)
kappa= _kappa(thisOrbit.vxvv[0,0],self._beta)
if not rrange == None:
if thisOrbit.vxvv[0,0] < rrange[0] \
or thisOrbit.vxvv[0,0] > rrange[1]:
continue
mult= numpy.ceil(kappa/wR*nphi)-1.
kappawR= kappa/wR*nphi-mult
while mult > 0:
if returnOrbit:
out.append(Orbit(vxvv=numpy.array([vxvv[0],vxvv[1],
vxvv[2],
stats.uniform.rvs()*numpy.pi*2.]).reshape(4)))
else:
out.append(thisOrbit)
mult-= 1
if stats.uniform.rvs() > kappawR:
continue
out.append(thisOrbit)
#Recurse to get enough
if len(out) < n*nphi:
out.extend(self.sample(n=int(n-len(out)/nphi),rrange=rrange,
returnROrbit=returnROrbit,
returnOrbit=returnOrbit,nphi=nphi))
if len(out) > n*nphi:
out= out[0:int(n*nphi)]
if kwargs.get('use_physical',True) and \
self._roSet and self._voSet:
if isinstance(out[0],Orbit):
dum= [o.turn_physical_on(ro=self._ro,vo=self._vo) for o in out]
return out
def _dlnfdR(self,R,vR,vT):
#Calculate a bunch of stuff that we need
E, L= vRvTRToEL(vR,vT,R,self._beta,self._dftype)
if self._beta == 0.:
xL= L
dRldR= vT
ECL= numpylog(xL)+0.5
dECLEdR= 0.
else: #non-flat rotation curve
xL= L**(1./(self._beta+1.))
dRldR= L**(1./(self._beta+1.))/R/(self._beta+1.)
ECL= 0.5*(1./self._beta+1.)*xL**(2.*self._beta)
dECLdRl= (1.+self._beta)*xL**(2.*self._beta-1)
dEdR= R**(2.*self._beta-1.)
dECLEdR= dECLdRl*dRldR-dEdR
sigma2xL= self._surfaceSigmaProfile.sigma2(xL,log=False)
return (self._surfaceSigmaProfile.surfacemassDerivative(xL,log=True)\
-(1.+(ECL-E)/sigma2xL)*self._surfaceSigmaProfile.sigma2Derivative(xL,log=True))*dRldR\
+dECLEdR/sigma2xL
def _dlnfdvR(self,R,vR,vT):
#Calculate a bunch of stuff that we need
E, L= vRvTRToEL(vR,vT,R,self._beta,self._dftype)
if self._beta == 0.:
xL= L
else: #non-flat rotation curve
xL= L**(1./(self._beta+1.))
sigma2xL= self._surfaceSigmaProfile.sigma2(xL,log=False)
return -vR/sigma2xL
def _dlnfdvT(self,R,vR,vT):
#Calculate a bunch of stuff that we need
E, L= vRvTRToEL(vR,vT,R,self._beta,self._dftype)
if self._beta == 0.:
xL= L
dRldvT= R
ECL= numpylog(xL)+0.5
dECLEdvT= 1./vT-vT
else: #non-flat rotation curve
xL= L**(1./(self._beta+1.))
dRldvT= L**(1./(self._beta+1.))/vT/(self._beta+1.)
ECL= 0.5*(1./self._beta+1.)*xL**(2.*self._beta)
dECLdRl= (1.+self._beta)*xL**(2.*self._beta-1)
dEdvT= vT
dECLEdvT= dECLdRl*dRldvT-dEdvT
sigma2xL= self._surfaceSigmaProfile.sigma2(xL,log=False)
return (self._surfaceSigmaProfile.surfacemassDerivative(xL,log=True)\
-(1.+(ECL-E)/sigma2xL)*self._surfaceSigmaProfile.sigma2Derivative(xL,log=True))*dRldvT\
+dECLEdvT/sigma2xL
[docs]class schwarzschilddf(shudf):
"""Schwarzschild's df"""
[docs] def __init__(self,surfaceSigma=expSurfaceSigmaProfile,
profileParams=(1./3.,1.0,0.2),
correct=False,
beta=0.,**kwargs):
"""
NAME:
__init__
PURPOSE:
Initialize a Schwarzschild DF
INPUT:
surfaceSigma - instance or class name of the target
surface density and sigma_R profile
(default: both exponential)
profileParams - parameters of the surface and sigma_R profile:
(xD,xS,Sro) where
xD - disk surface mass scalelength (can be Quantity)
xS - disk velocity dispersion scalelength (can be Quantity)
Sro - disk velocity dispersion at Ro (can be Quantity)
Directly given to the 'surfaceSigmaProfile class, so
could be anything that class takes
beta - power-law index of the rotation curve
correct - if True, correct the DF
ro= distance from vantage point to GC (kpc; can be Quantity)
vo= circular velocity at ro (km/s; can be Quantity)
+DFcorrection kwargs (except for those already specified)
OUTPUT:
instance
HISTORY:
2017-09-17 - Written - Bovy (UofT)
"""
# Schwarzschild == Shu w/ energy computed in epicycle approx.
# so all functions are the same as in Shu, only thing different is
# how E is computed
return diskdf.__init__(self,surfaceSigma=surfaceSigma,
profileParams=profileParams,
correct=correct,dftype='schwarzschild',
beta=beta,**kwargs)
def _surfaceIntegrand(vR,vT,R,df,logSigmaR,logsigmaR2,sigmaR1,gamma):
"""Internal function that is the integrand for the surface mass integration"""
E,L= _vRpvTpRToEL(vR,vT,R,df._beta,sigmaR1,gamma,df._dftype)
return df.eval(E,L,logSigmaR,logsigmaR2)*2.*numpy.pi/df._gamma #correct
def _sigma2surfaceIntegrand(vR,vT,R,df,logSigmaR,logsigmaR2,sigmaR1,gamma):
"""Internal function that is the integrand for the sigma-squared times
surface mass integration"""
E,L= _vRpvTpRToEL(vR,vT,R,df._beta,sigmaR1,gamma,df._dftype)
return vR**2.*df.eval(E,L,logSigmaR,logsigmaR2)*2.*numpy.pi/df._gamma #correct
def _vmomentsurfaceIntegrand(vR,vT,R,df,logSigmaR,logsigmaR2,sigmaR1,gamma,
n,m):
"""Internal function that is the integrand for the velocity moment times
surface mass integration"""
E,L= _vRpvTpRToEL(vR,vT,R,df._beta,sigmaR1,gamma,df._dftype)
return vR**n*vT**m*df.eval(E,L,logSigmaR,logsigmaR2)*2.*numpy.pi/df._gamma #correct
def _vmomentderivsurfaceIntegrand(vR,vT,R,df,logSigmaR,logsigmaR2,sigmaR1,
gamma,n,m,deriv):
"""Internal function that is the integrand for the derivative of velocity
moment times surface mass integration"""
E,L= _vRpvTpRToEL(vR,vT,R,df._beta,sigmaR1,gamma,df._dftype)
if deriv.lower() == 'r':
return vR**n*vT**m*df.eval(E,L,logSigmaR,logsigmaR2)*2.*numpy.pi/df._gamma*df._dlnfdR(R,vR*sigmaR1,vT*sigmaR1/gamma) #correct
else:
return 0.
def _vRpvTpRToEL(vR,vT,R,beta,sigmaR1,gamma,dftype='dehnen'):
"""Internal function that calculates E and L given velocities normalized by the velocity dispersion"""
vR*= sigmaR1
vT*= sigmaR1/gamma
return vRvTRToEL(vR,vT,R,beta,dftype)
def _oned_intFunc(x,twodfunc,gfun,hfun,tol,args):
"""Internal function for bovy_dblquad"""
thisargs= copy.deepcopy(args)
thisargs.insert(0,x)
return integrate.romberg(twodfunc,gfun(x),hfun(x),args=thisargs,tol=tol)
def bovy_dblquad(func, a, b, gfun, hfun, args=(), tol=1.48e-08):
"""
NAME:
bovy_dblquad
PURPOSE:
like scipy.integrate's dblquad, but using Romberg integration for the one-d integrals and using tol
INPUT:
same as scipy.integrate.dblquad except for tol and epsrel,epsabs
OUTPUT:
value
HISTORY:
2010-03-11 - Written - Bpvy (NYU)
"""
return integrate.romberg(_oned_intFunc,a,b,args=(func,gfun,hfun,tol,args),tol=tol)
class DFcorrection(object):
"""Class that contains the corrections necessary to reach
exponential profiles"""
def __init__(self,**kwargs):
"""
NAME:
__init__
PURPOSE:
initialize the corrections: set them, load them, or calculate
and save them
OPTIONAL INPUTS:
corrections - if Set, these are the corrections and they should
be used as such
npoints - number of points from 0 to Rmax
rmax - correct up to this radius (/ro) (default: 5)
savedir - save the corrections in this directory
surfaceSigmaProfile - target surfacemass and sigma_R^2 instance
beta - power-law index of the rotation curve (when calculating)
dftype - classname of the DF
niter - number of iterations to perform to calculate the corrections
interp_k - 'k' keyword to give to InterpolatedUnivariateSpline
OUTPUT:
HISTORY:
2010-03-10 - Written - Bovy (NYU)
"""
if not 'surfaceSigmaProfile' in kwargs:
raise DFcorrectionError("surfaceSigmaProfile not given")
else:
self._surfaceSigmaProfile= kwargs['surfaceSigmaProfile']
self._rmax= kwargs.get('rmax',5.)
self._niter= kwargs.get('niter',20)
if not 'npoints' in kwargs:
if 'corrections' in kwargs:
self._npoints= kwargs['corrections'].shape[0]
else: #pragma: no cover
self._npoints= 151 #would take too long to cover
else:
self._npoints= kwargs['npoints']
self._dftype= kwargs.get('dftype',dehnendf)
self._beta= kwargs.get('beta',0.)
self._rs= numpy.linspace(_RMIN,self._rmax,self._npoints)
self._interp_k= kwargs.get('interp_k',_INTERPDEGREE)
if 'corrections' in kwargs:
self._corrections= kwargs['corrections']
if not len(self._corrections) == self._npoints:
raise DFcorrectionError("Number of corrections has to be equal to the number of points npoints")
else:
self._savedir= kwargs.get('savedir',_CORRECTIONSDIR)
self._savefilename= self._createSavefilename(self._niter)
if os.path.exists(self._savefilename):
savefile= open(self._savefilename,'rb')
self._corrections= numpy.array(pickle.load(savefile))
savefile.close()
else: #Calculate the corrections
self._corrections= self._calc_corrections()
#Interpolation; smoothly go to zero
interpRs= numpy.append(self._rs,2.*self._rmax)
self._surfaceInterpolate= interpolate.InterpolatedUnivariateSpline(interpRs,
numpylog(numpy.append(self._corrections[:,0],1.)),
k=self._interp_k)
self._sigma2Interpolate= interpolate.InterpolatedUnivariateSpline(interpRs,
numpylog(numpy.append(self._corrections[:,1],1.)),
k=self._interp_k)
#Interpolation for R < _RMIN
surfaceInterpolateSmallR= interpolate.UnivariateSpline(interpRs[0:_INTERPDEGREE+2],numpylog(self._corrections[0:_INTERPDEGREE+2,0]),k=_INTERPDEGREE)
self._surfaceDerivSmallR= surfaceInterpolateSmallR.derivatives(interpRs[0])[1]
sigma2InterpolateSmallR= interpolate.UnivariateSpline(interpRs[0:_INTERPDEGREE+2],numpylog(self._corrections[0:_INTERPDEGREE+2,1]),k=_INTERPDEGREE)
self._sigma2DerivSmallR= sigma2InterpolateSmallR.derivatives(interpRs[0])[1]
return None
def _createSavefilename(self,niter):
#Form surfaceSigmaProfile string
sspFormat= self._surfaceSigmaProfile.formatStringParams()
sspString= ''
for format in sspFormat:
sspString+= format+'_'
return os.path.join(self._savedir,'dfcorrection_'+
self._dftype.__name__+'_'+
self._surfaceSigmaProfile.__class__.__name__+'_'+
sspString % self._surfaceSigmaProfile.outputParams()+
'%6.4f_%i_%6.4f_%i.sav'
% (self._beta,self._npoints,self._rmax,niter))
def correct(self,R,log=False):
"""
NAME:
correct
PURPOSE:
calculate the correction in Sigma and sigma2 at R
INPUT:
R - Galactocentric radius(/ro)
log - if True, return the log of the correction
OUTPUT:
[Sigma correction, sigma2 correction]
HISTORY:
2010-03-10 - Written - Bovy (NYU)
"""
if isinstance(R,numpy.ndarray):
out= numpy.empty((2,len(R)))
#R < _RMIN
rmin_indx= (R < _RMIN)
if numpy.sum(rmin_indx) > 0:
out[0,rmin_indx]= numpylog(self._corrections[0,0])\
+self._surfaceDerivSmallR*(R[rmin_indx]-_RMIN)
out[1,rmin_indx]= numpylog(self._corrections[0,1])\
+self._sigma2DerivSmallR*(R[rmin_indx]-_RMIN)
#R > 2rmax
rmax_indx= (R > (2.*self._rmax))
if numpy.sum(rmax_indx) > 0:
out[:,rmax_indx]= 0.
#'normal' R
r_indx= (R >= _RMIN)*(R <= (2.*self._rmax))
if numpy.sum(r_indx) > 0:
out[0,r_indx]= self._surfaceInterpolate(R[r_indx])
out[1,r_indx]= self._sigma2Interpolate(R[r_indx])
if log: return out
else: return numpy.exp(out)
if R < _RMIN:
out= numpy.array([numpylog(self._corrections[0,0])+self._surfaceDerivSmallR*(R-_RMIN),
numpylog(self._corrections[0,1])+self._sigma2DerivSmallR*(R-_RMIN)])
elif R > (2.*self._rmax):
out= numpy.array([0.,0.])
else:
if _SCIPY_VERSION >= _SCIPY_VERSION_BREAK:
out= numpy.array([self._surfaceInterpolate(R),
self._sigma2Interpolate(R)])
else: #pragma: no cover
out= numpy.array([self._surfaceInterpolate(R)[0],
self._sigma2Interpolate(R)[0]])
if log:
return out
else:
return numpy.exp(out)
def derivLogcorrect(self,R):
"""
NAME:
derivLogcorrect
PURPOSE:
calculate the derivative of the log of the correction in Sigma
and sigma2 at R
INPUT:
R - Galactocentric radius(/ro)
OUTPUT:
[d log(Sigma correction)/dR, d log(sigma2 correction)/dR]
HISTORY:
2010-03-10 - Written - Bovy (NYU)
"""
if R < _RMIN:
out= numpy.array([self._surfaceDerivSmallR,
self._sigma2DerivSmallR])
elif R > (2.*self._rmax):
out= numpy.array([0.,0.])
else:
if _SCIPY_VERSION >= _SCIPY_VERSION_BREAK:
out= numpy.array([self._surfaceInterpolate(R,nu=1),
self._sigma2Interpolate(R,nu=1)])
else: #pragma: no cover
out= numpy.array([self._surfaceInterpolate(R,nu=1)[0],
self._sigma2Interpolate(R,nu=1)[0]])
return out
def _calc_corrections(self):
"""Internal function that calculates the corrections"""
searchIter= self._niter-1
while searchIter > 0:
trySavefilename= self._createSavefilename(searchIter)
if os.path.exists(trySavefilename):
trySavefile= open(trySavefilename,'rb')
corrections= numpy.array(pickle.load(trySavefile))
trySavefile.close()
break
else:
searchIter-= 1
if searchIter == 0:
corrections= numpy.ones((self._npoints,2))
for ii in range(searchIter,self._niter):
if ii == 0:
currentDF= self._dftype(surfaceSigma=self._surfaceSigmaProfile,
beta=self._beta)
else:
currentDF= self._dftype(surfaceSigma=self._surfaceSigmaProfile,
beta=self._beta,
corrections=corrections,
npoints=self._npoints,
rmax=self._rmax,
savedir=self._savedir,
interp_k=self._interp_k)
newcorrections= numpy.zeros((self._npoints,2))
for jj in range(self._npoints):
thisSurface= currentDF.surfacemass(self._rs[jj],
use_physical=False)
newcorrections[jj,0]= currentDF.targetSurfacemass(self._rs[jj],use_physical=False)/thisSurface
newcorrections[jj,1]= currentDF.targetSigma2(self._rs[jj],use_physical=False)*thisSurface\
/currentDF.sigma2surfacemass(self._rs[jj],
use_physical=False)
#print(jj, newcorrections[jj,:])
corrections*= newcorrections
#Save
picklethis= []
for arr in list(corrections):
picklethis.append([float(a) for a in arr])
save_pickles(self._savefilename,picklethis) #We pickle a list for platform-independence)
return corrections
class DFcorrectionError(Exception):
def __init__(self, value):
self.value = value
def __str__(self):
return repr(self.value)
def vRvTRToEL(vR,vT,R,beta,dftype='dehnen'):
"""
NAME:
vRvTRToEL
PURPOSE:
calculate the energy and angular momentum
INPUT:
vR - radial velocity
vT - tangential velocity
R - Galactocentric radius
OUTPUT:
HISTORY:
2010-03-10 - Written - Bovy (NYU)
"""
if dftype == 'schwarzschild':
# Compute E in the epicycle approximation
gamma= numpy.sqrt(2./(1.+beta))
L= R*vT
if beta == 0.:
xL= L
else: #non-flat rotation curve
xL= L**(1./(beta+1.))
return (0.5*vR**2.+0.5*gamma**2.*(vT-R**beta)**2.
+xL**(2.*beta)/2.+axipotential(xL,beta=beta),
L)
else:
return (axipotential(R,beta)+0.5*vR**2.+0.5*vT**2.,vT*R)
def axipotential(R,beta=0.):
"""
NAME:
axipotential
PURPOSE:
return the axisymmetric potential at R/Ro
INPUT:
R - Galactocentric radius
beta - rotation curve power-law
OUTPUT:
Pot(R)/vo**2.
HISTORY:
2010-03-01 - Written - Bovy (NYU)
"""
if beta == 0.:
if numpy.any(R == 0.):
out= numpy.empty(R.shape)
out[R == 0.]= numpylog(_RMIN)
out[R != 0.]= numpylog(R[R != 0.])
return out
else:
return numpylog(R)
else: #non-flat rotation curve
return R**(2.*beta)/2./beta
def _ars_hx(x,args):
"""
NAME:
_ars_hx
PURPOSE:
h(x) for ARS sampling of the input surfacemass profile
INPUT:
x - R(/ro)
args= (surfaceSigma, dfcorr)
surfaceSigma - surfaceSigmaProfile instance
dfcorr - DFcorrection instance
OUTPUT:
log(x)+log surface(x) + log(correction)
HISTORY:
2010-07-11 - Written - Bovy (NYU)
"""
surfaceSigma, dfcorr= args
if dfcorr is None:
return numpylog(x)+surfaceSigma.surfacemass(x,log=True)
else:
return numpylog(x)+surfaceSigma.surfacemass(x,log=True)+dfcorr.correct(x)[0]
def _ars_hpx(x,args):
"""
NAME:
_ars_hpx
PURPOSE:
h'(x) for ARS sampling of the input surfacemass profile
INPUT:
x - R(/ro)
args= (surfaceSigma, dfcorr)
surfaceSigma - surfaceSigmaProfile instance
dfcorr - DFcorrection instance
OUTPUT:
derivative of log(x)+log surface(x) + log(correction) wrt x
HISTORY:
2010-07-11 - Written - Bovy (NYU)
"""
surfaceSigma, dfcorr= args
if dfcorr is None:
return 1./x+surfaceSigma.surfacemassDerivative(x,log=True)
else:
return 1./x+surfaceSigma.surfacemassDerivative(x,log=True)+dfcorr.derivLogcorrect(x)[0]
def _kappa(R,beta):
"""Internal function to give kappa(r)"""
return numpy.sqrt(2.*(1.+beta))*R**(beta-1)
def _dlToRphi(d,l):
"""Convert d and l to R and phi, l is in radians"""
R= numpy.sqrt(1.+d**2.-2.*d*numpy.cos(l))
if R == 0.:
R+= 0.0001
d+= 0.0001
if 1./numpy.cos(l) < d and numpy.cos(l) > 0.:
theta= numpy.pi-numpy.arcsin(d/R*numpy.sin(l))
else:
theta= numpy.arcsin(d/R*numpy.sin(l))
return (R,theta)
def _vtmaxEq(vT,R,diskdf):
"""Equation to solve to find the max vT at R"""
#Calculate a bunch of stuff that we need
if diskdf._beta == 0.:
E= vT**2./2.+numpylog(R)
xE= numpy.exp(E-.5)
OE= xE**-1.
LCE= xE
dxEdvT= xE*vT
else: #non-flat rotation curve
E= vT**2./2.+1./2./diskdf._beta*R**(2.*diskdf._beta)
xE= (2.*E/(1.+1./diskdf._beta))**(1./2./diskdf._beta)
OE= xE**(diskdf._beta-1.)
LCE= xE**(diskdf._beta+1.)
dxEdvT= xE/2./diskdf._beta/E*vT
L= R*vT
sigma2xE= diskdf._surfaceSigmaProfile.sigma2(xE,log=False)
return OE*R/sigma2xE+\
(diskdf._surfaceSigmaProfile.surfacemassDerivative(xE,log=True)\
-(1.+OE*(L-LCE)/sigma2xE)*diskdf._surfaceSigmaProfile.sigma2Derivative(xE,log=True)\
+(L-LCE)/sigma2xE*(diskdf._beta-1.)*xE**(diskdf._beta-2.)\
-OE*(diskdf._beta+1.)/sigma2xE*xE**diskdf._beta)\
*dxEdvT
def _marginalizeVperpIntegrandSinAlphaLarge(vR,df,R,sinalpha,cotalpha,
vlos,vcirc,sigma):
return df(*vRvTRToEL(vR*sigma,cotalpha*vR*sigma+vlos/sinalpha+vcirc,
R,df._beta,df._dftype))
def _marginalizeVperpIntegrandSinAlphaSmall(vT,df,R,cosalpha,tanalpha,
vlos,vcirc,sigma):
return df(*vRvTRToEL(tanalpha*vT*sigma-vlos/cosalpha,vT*sigma+vcirc,
R,df._beta,df._dftype))