Source code for galpy.df.diskdf

###############################################################################
#   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.ars import ars
from ..util import save_pickles, conversion
from ..util.conversion import physical_conversion, \
    potential_physical_input, _APY_UNITS, _APY_LOADED, surfdens_in_msolpc2
from ..potential import PowerSphericalPotential
from ..actionAngle import actionAngleAdiabatic, actionAngleAxi
from .df import df
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=\
                    (conversion.parse_length(profileParams[0],ro=self._ro),
                     conversion.parse_length(profileParams[1],ro=self._ro),
                     conversion.parse_velocity(profileParams[2],vo=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= conversion.parse_angle(l) elif deg: lrad= l*_DEGTORAD else: lrad= l d= conversion.parse_length(d,ro=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= conversion.parse_angle(l) elif deg: lrad= l*_DEGTORAD else: lrad= l d= conversion.parse_length(d,ro=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 l= conversion.parse_angle(l) maxd= conversion.parse_length(maxd,ro=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= conversion.parse_angle(los) 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 ro= conversion.parse_length_kpc(ro) vo= kwargs.pop('vo',None) if vo is None and hasattr(self,'_voSet') and self._voSet: vo= self._vo vo= conversion.parse_velocity_kms(vo) 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() E= conversion.parse_energy(E,vo=self._vo) L= conversion.parse_angmom(L,ro=self._ro,vo=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(ars([0.,0.],[True,False],[0.05,2.],_ars_hx, _ars_hpx,nsamples=n, hxparams=(self._surfaceSigmaProfile, self._corr))) else: xE= numpy.array(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: rrange[0]= conversion.parse_length(rrange[0],ro=self._ro) rrange[1]= conversion.parse_length(rrange[1],ro=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) """ E= conversion.parse_energy(E,vo=self._vo) L= conversion.parse_angmom(L,ro=self._ro,vo=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(ars([0.,0.],[True,False],[0.05,2.],_ars_hx, _ars_hpx,nsamples=n, hxparams=(self._surfaceSigmaProfile, self._corr))) else: xL= numpy.array(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: rrange[0]= conversion.parse_length(rrange[0],ro=self._ro) rrange[1]= conversion.parse_length(rrange[1],ro=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))