<commit_before>###############################################################################
#   diskdf.py: module that interprets (E,Lz) pairs in terms of a 
#              distribution function
#
#   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
###############################################################################
_EPSREL=10.**-14.
_NSIGMA= 4.
_INTERPDEGREE= 3
_RMIN=10.**-10.
_MAXD_REJECTLOS= 4.
import copy
import os, os.path
import cPickle as pickle
import math as m
import numpy as nu
import scipy as sc
import scipy.integrate as integrate
import scipy.interpolate as interpolate
from scipy import linalg
from scipy import stats
from scipy import optimize
from surfaceSigmaProfile import *
from galpy.orbit import Orbit
from galpy.util.bovy_ars import bovy_ars
from galpy.potential import PowerSphericalPotential
from galpy.actionAngle import actionAngleFlat, actionAnglePower
from galpy.actionAngle_src.actionAngleFlat import calcRapRperiFromELFlat #HACK
from galpy.actionAngle_src.actionAnglePower import \
    calcRapRperiFromELPower #HACK
_CORRECTIONSDIR=os.path.join(os.path.dirname(os.path.realpath(__file__)),'data')
_DEGTORAD= m.pi/180.
class diskdf:
    """Class that represents a disk DF"""
    def __init__(self,dftype='dehnen',
                 surfaceSigma=expSurfaceSigmaProfile,
                 profileParams=(1./3.,1.0,0.2),
                 correct=False,
                 beta=0.,**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)
        """
        self._dftype= dftype
        if isinstance(surfaceSigma,surfaceSigmaProfile):
            self._surfaceSigmaProfile= surfaceSigma
        else:
            self._surfaceSigmaProfile= surfaceSigma(profileParams)
        self._beta= beta
        if correct or kwargs.has_key('corrections') or kwargs.has_key('rmax') or kwargs.has_key('niter') or kwargs.has_key('npoints'):
            self._correct= True
            #Load corrections
            self._corr= DFcorrection(dftype=self.__class__,
                                     surfaceSigmaProfile=self._surfaceSigmaProfile,
                                     beta=beta,**kwargs)
        else:
            self._correct= False
        return None
    
    def __call__(self,*args):
        """
        NAME:
           __call__
        PURPOSE:
           evaluate the distribution function
        INPUT:
           either an orbit instance or E,Lz

           1) Orbit instance:
              a) Orbit instance alone: use vxvv member
              b) Orbit instance + t: call the Orbit instance

           2)
              E - energy (/vo^2)
              L - angular momentun (/ro/vo)
        OUTPUT:
           DF(orbit/E,L)
        HISTORY:
           2010-07-10 - Written - Bovy (NYU)
        """
        if isinstance(args[0],Orbit):
            if len(args) == 1:
                return sc.real(self.eval(*vRvTRToEL(args[0].vxvv[1],
                                                    args[0].vxvv[2],
                                                    args[0].vxvv[0],
                                                    self._beta)))
            else:
                vxvv= args[0](args[1])
                return sc.real(self.eval(*vRvTRToEL(vxvv[1],
                                                    vxvv[2],
                                                    vxvv[0],
                                                    self._beta)))
        else:
            return sc.real(self.eval(*args))

    def targetSigma2(self,R,log=False):
        """
        NAME:
           targetSigma2
        PURPOSE:
           evaluate the target Sigma_R^2(R)
        INPUT:
           R - radius at which to evaluate (/ro)
        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)
    
    def targetSurfacemass(self,R,log=False):
        """
        NAME:
           targetSurfacemass
        PURPOSE:
           evaluate the target surface mass at R
        INPUT:
           R - radius at which to evaluate
           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)
        
    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
           l - Galactic longitude (in deg, unless deg=False)
           deg= if False, l is in radians
           log - if True, return the log (default: False)
        OUTPUT:
           Sigma(d,l)
        HISTORY:
           2011-03-23 - Written - Bovy (NYU)
        """
        #Calculate R and phi
        if deg:
            lrad= l*_DEGTORAD
        else:
            lrad= l
        R, phi= _dlToRphi(d,lrad)
        #Evaluate Jacobian
        jac= _jacobian_rphi_dl(d,lrad,R=R,phi=phi)
        if log:
            return self._surfaceSigmaProfile.surfacemass(R,log=log)\
                +m.log(m.fabs(jac))+m.log(R)
            pass
        else:
            return self._surfaceSigmaProfile.surfacemass(R,log=log)\
                *m.fabs(jac)*R
        
    def surfacemassLOS(self,d,l,deg=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
           l - Galactic longitude (in deg, unless deg=False)
        OPTIONAL INPUT:
           nsigma - number of sigma to integrate the velocities over
        KEYWORDS:
           romberg - if True, use a romberg integrator (default: False)
           deg= if False, l is in radians
        OUTPUT:
           Sigma(d,l)
        HISTORY:
           2011-03-24 - Written - Bovy (NYU)
        """
        #Calculate R and phi
        if deg:
            lrad= l*_DEGTORAD
        else:
            lrad= l
        R, phi= _dlToRphi(d,lrad)
        #Evaluate Jacobian
        jac= _jacobian_rphi_dl(d,lrad,R=R,phi=phi)
        return self.surfacemass(R,romberg=romberg,nsigma=nsigma,
                                relative=relative)\
                                *m.fabs(jac)*R

    def sampledSurfacemassLOS(self,l,n=1,maxd=None,target=True):
        """sample a distance along the line-of-sight
        BOVY: SHOULD USE CORRECTIONS?"""
        #First calculate where the maximum is
        if l == 0.:
            maxSM= self.targetSurfacemass(0.)
        elif l >= m.pi/2. and l <= 3.*m.pi/2.:
            maxSM= self.targetSurfacemass(1.)
        elif l < m.pi/2. or l > 3.*m.pi/2.:
            if target:
                minR= optimize.fmin_bfgs(lambda x: \
                                             -x*self.targetSurfacemassLOS(x,l,
                                                                          deg=False),
                                         0.,disp=False)[0]
                maxSM= self.targetSurfacemassLOS(minR,l,deg=False)*minR
            else:
                minR= optimize.fmin_bfgs(lambda x: \
                                             -x*self.surfacemassLOS(x,l,
                                                                    deg=False),
                                         0.,disp=False)[0]
                maxSM= self.surfacemassLOS(minR,l,deg=False)*minR
        #Now rejection-sample
        if maxd is None:
            maxd= _MAXD_REJECTLOS
        out= []
        while len(out) < n:
            #sample
            prop= nu.random.random()*maxd
            if target:
                surfmassatprop= self.targetSurfacemassLOS(prop,l,deg=False)*prop
            else:
                surfmassatprop= self.surfacemassLOS(prop,l,deg=False)*prop
            if surfmassatprop/maxSM > nu.random.random(): #accept
                out.append(prop)
        return nu.array(out)

    def sampleVRVTVdistdR(self,R,n=1,nsigma=None,target=True):
        """sample a radial and tangential velocity at R
        BOVY: USE THE FACT THAT THIS CAN BE SEPARATED INTO vR and vT"""
        #Determine where the max of the v-distribution is using asymmetric drift
        maxVR= 0.
        gamma= sc.sqrt(2./(1.+self._beta))           
        maxVT= optimize.brentq(_vtmaxEq,0.,1.,(R,self))
        maxVD= self(Orbit([R,maxVR,maxVT]))
        #Now rejection-sample
        if nsigma == None:
            nsigma= _NSIGMA
        out= []
        if target:
            sigma= m.sqrt(self.targetSigma2(R))
        else:
            sigma= m.sqrt(self.sigma2(R))
        while len(out) < n:
            #sample
            propvR= nu.random.normal()*_NSIGMA*sigma
            propvT= nu.random.normal()*_NSIGMA*sigma/gamma+1.
            VDatprop= self(Orbit([R,propvR,propvT]))
            if VDatprop/maxVD > nu.random.random(): #accept
                out.append(sc.array([propvR,propvT]))
        return nu.array(out)

    def sampleLOS(self,los,n=1,deg=True,maxd=None,nsigma=None,
                  target=True):
        """
        NAME:
           sampleLOS
        PURPOSE:
           sample along a given LOS
        INPUT:
           los - line of sight (in deg, unless deg=False)
           n= number of desired samples
           deg= los in degrees? (default=True)
           target= if True, use target surface mass and sigma2 profiles
                   (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 deg:
            l= los*_DEGTORAD
        else:
            l= los
        out= []
        #sample distances
        ds= self.sampledSurfacemassLOS(l,n=n,maxd=maxd,target=target)
        for ii in range(int(n)):
            #Calculate R and phi
            thisR,thisphi= _dlToRphi(ds[ii],l)
            #sample velocities
            vv= self.sampleVRVTVdistdR(thisR,n=1,nsigma=nsigma,target=target)[0]
            out.append(Orbit([thisR,vv[0],vv[1],thisphi]))
        return out

    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 (/ro)
        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)
        sigmaR2= self.targetSigma2(R)
        sigmaR1= sc.sqrt(sigmaR2)
        logsigmaR2= sc.log(sigmaR2)
        gamma= sc.sqrt(2./(1.+self._beta))
        if relative:
            norm= 1.
        else:
            norm= sc.exp(logSigmaR)
        if romberg:
            return bovy_dblquad(_surfaceIntegrand,
                                gamma*R**self._beta/sigmaR1-nsigma,
                                gamma*R**self._beta/sigmaR1+nsigma,
                                lambda x: 0., lambda x: nsigma,
                                [R,self,logSigmaR,logsigmaR2,sigmaR1,gamma],
                                tol=10.**-8)/sc.pi*norm
        else:
            return integrate.dblquad(_surfaceIntegrand,
                                     gamma*R**self._beta/sigmaR1-nsigma,
                                     gamma*R**self._beta/sigmaR1+nsigma,
                                     lambda x: 0., lambda x: nsigma,
                                     (R,self,logSigmaR,logsigmaR2,sigmaR1,
                                      gamma),
                                     epsrel=_EPSREL)[0]/sc.pi*norm

    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 (/ro)
        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)
        sigmaR2= self.targetSigma2(R)
        sigmaR1= sc.sqrt(sigmaR2)
        logsigmaR2= sc.log(sigmaR2)
        gamma= sc.sqrt(2./(1.+self._beta))
        if relative:
            norm= 1.
        else:
            norm= sc.exp(logSigmaR+logsigmaR2)
        if romberg:
            return bovy_dblquad(_sigma2surfaceIntegrand,
                                gamma*R**self._beta/sigmaR1-nsigma,
                                gamma*R**self._beta/sigmaR1+nsigma,
                                lambda x: 0., lambda x: nsigma,
                                [R,self,logSigmaR,logsigmaR2,sigmaR1,gamma],
                                tol=10.**-8)/sc.pi*norm
        else:
            return integrate.dblquad(_sigma2surfaceIntegrand,
                                     gamma*R**self._beta/sigmaR1-nsigma,
                                     gamma*R**self._beta/sigmaR1+nsigma,
                                     lambda x: 0., lambda x: nsigma,
                                     (R,self,logSigmaR,logsigmaR2,sigmaR1,
                                      gamma),
                                     epsrel=_EPSREL)[0]/sc.pi*norm

    def sigma2(self,R,romberg=False,nsigma=None):
        """
        NAME:
           sigma2
        PURPOSE:
           calculate sigma_R^2 at R by marginalizing over velocity
        INPUT:
           R - radius at which to calculate sigma_R^2 density (/ro)
        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)/self.surfacemass(R,romberg,nsigma)

    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.:
            rperi, rap= calcRapRperiFromELFlat(E,L,vc=1.,ro=1.)
            aA= actionAngleFlat(rperi,0.,L/rperi)
        else:
            rperi, rap= calcRapRperiFromELPower(E,L,vc=1.,ro=1.)
            aA= actionAnglePower(rperi,0.,L/rperi,beta=self._beta)
        TR= aA.TR()[0]
        return (2.*m.pi/TR, rap, rperi)

    def sample(self,n=1,rrange=None,returnROrbit=True,returnOrbit=False,
               nphi=1.):
        """
        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
        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 AttributeError("'sample' method for this disk df is not implemented")

class dehnendf(diskdf):
    """Dehnen's 'new' df"""
    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 / 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 - if True, correct the DF

           + DFcorrection kwargs (except for those already specified)
        OUTPUT:
        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 (/vo^2)
           L - angular momentun (/ro/vo)
       OUTPUT:
           DF(E,L)
        HISTORY:
           2010-03-10 - Written - Bovy (NYU)
           2010-03-28 - Moved to dehnenDF - Bovy (NYU)
        """
        #Calculate Re,LE, OmegaE
        if self._beta == 0.:
            xE= sc.exp(E-.5)
            logOLLE= sc.log(L/xE-1.)
        else: #non-flat rotation curve
            xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
            logOLLE= self._beta*sc.log(xE)+sc.log(L/xE-xE**self._beta)
        if self._correct: 
            correction= self._corr.correct(xE,log=True)
        else:
            correction= sc.zeros(2)
        SRE2= self.targetSigma2(xE,log=True)+correction[1]
        return sc.exp(logsigmaR2-SRE2+self.targetSurfacemass(xE,log=True)-logSigmaR+sc.exp(logOLLE-SRE2)+correction[0])

    def sample(self,n=1,rrange=None,returnROrbit=True,returnOrbit=False,
               nphi=1.,los=None,losdeg=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
           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)
        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)
        """
        #First sample xE
        if self._correct:
            xE= sc.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
                                  _ars_hpx,nsamples=n,
                                  hxparams=(self._surfaceSigmaProfile,
                                            self._corr)))
        else:
            xE= sc.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= sc.log(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)*sc.log(stats.uniform.rvs(size=n))/OR
        if self._correct:
            for ii in range(len(xE)):
                Lz[ii]*= self._corr.correct(xE[ii],log=False)[1]
        Lz+= LCE
        if not returnROrbit and not returnOrbit:
            out= [[e,l] for e,l in zip(E,Lz)]
        else:
            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.*m.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(sc.array([0.,tr]),self._psp)
                if returnOrbit:
                    vxvv= thisOrbit(tr).vxvv
                    if not los is None: #Sample along a given line of sight
                        if losdeg: l= los*_DEGTORAD
                        else: l= los
                        if l > (2.*m.pi): l-= 2.*m.pi
                        if l < 0: l+= 2.*m.pi
                        sinphil= 1./vxvv[0]*m.sin(l)
                        if m.fabs(sinphil) > 1.: continue
                        if stats.uniform.rvs() < 0.5:
                            phil= m.asin(sinphil)
                        else:
                            phil= m.pi-m.asin(sinphil)
                        phi= phil-l
                        if phi > (2.*m.pi): phi-= 2.*m.pi
                        if phi < 0: phi+= 2.*m.pi
                        #make sure this is on the right side of the los
                        if l >= 0. and l <= m.pi/2. and phi > m.pi: continue
                        elif l >= m.pi/2. and l <= m.pi and phi > m.pi/2.: \
                                continue
                        elif l >= m.pi and l <= 3.*m.pi/2. and phi < 3.*m.pi/2.: continue
                        elif l >= 3.*m.pi/2. and phi < m.pi: continue
                            
                        thisOrbit= Orbit(vxvv=sc.array([vxvv[0],vxvv[1],
                                                        vxvv[2],
                                                        phi]).reshape(4))
                    else:
                        thisOrbit= Orbit(vxvv=sc.array([vxvv[0],vxvv[1],vxvv[2],
                                                        stats.uniform.rvs()\
                                                            *m.pi*2.])\
                                             .reshape(4))
                else:
                    thisOrbit= Orbit(thisOrbit(tr))
                kappa= _kappa(thisOrbit.vxvv[0],self._beta)
                if not rrange == None:
                    if thisOrbit.vxvv[0] < rrange[0] \
                            or thisOrbit.vxvv[0] > rrange[1]:
                        continue
                if los is None:
                    mult= sc.ceil(kappa/wR*nphi)-1.
                    kappawR= kappa/wR*nphi-mult
                    while mult > 0:
                        if returnOrbit:
                            out.append(Orbit(vxvv=sc.array([vxvv[0],vxvv[1],
                                                            vxvv[2],
                                                            stats.uniform.rvs()*m.pi*2.]).reshape(4)))
                        else:
                            out.append(thisOrbit)
                        mult-= 1
                else:
                    if losdeg: l= los*_DEGTORAD
                    else: l= los
                    if l > (2.*m.pi): l-= 2.*m.pi
                    if l < 0: l+= 2.*m.pi
                    sinphil= 1./vxvv[0]*m.sin(l)
                    if m.fabs(sinphil) > 1.: continue
                    if stats.uniform.rvs() < 0.5:
                        phil= m.asin(sinphil)
                    else:
                        phil= m.pi-m.asin(sinphil)
                    phi= phil-l
                    if phi > (2.*m.pi): phi-= 2.*m.pi
                    if phi < 0: phi+= 2.*m.pi
                                #make sure this is on the right side of the los
                    if l >= 0. and l <= m.pi/2. and phi > m.pi: continue
                    elif l >= m.pi/2. and l <= m.pi and phi > m.pi/2.: \
                            continue
                    elif l >= m.pi and l <= 3.*m.pi/2. and phi < 3.*m.pi/2.: \
                            continue
                    elif l >= 3.*m.pi/2. and phi < m.pi: continue
                    #Calcualte dphidl
                    dphidl= m.fabs(1./vxvv[0]*m.cos(l)/m.cos(l+phi)-1.)
                    mult= sc.ceil(dphidl*kappa/wR*nphi)-1.
                    kappawR= kappa/wR*nphi-mult
                    while mult > 0:
                        out.append(Orbit(vxvv=sc.array([vxvv[0],vxvv[1],
                                                        vxvv[2],
                                                        phi]).reshape(4)))
                        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)]
        return out

class shudf(diskdf):
    """Shu's df (1969)"""
    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 / 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 - if True, correct the DF

           + DFcorrection kwargs (except for those already specified)
        OUTPUT:
        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)
        """
        #Calculate RL,LL, OmegaL
        if self._beta == 0.:
            xL= L
            logECLE= sc.log(-sc.log(xL)-0.5+E)
        else: #non-flat rotation curve
            xL= L**(1./(self._beta+1.))
            logECLE= sc.log(-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= sc.zeros(2)
        SRE2= self.targetSigma2(xL,log=True)+correction[1]
        return sc.exp(logsigmaR2-SRE2+self.targetSurfacemass(xL,log=True)-logSigmaR-sc.exp(logECLE-SRE2)+correction[0])

    def sample(self,n=1,rrange=None,returnROrbit=True,returnOrbit=False,
               nphi=1.,los=None,losdeg=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
           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)
           losdeg= if False, los is in radians (default=True)
        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)
        """
        #First sample xL
        if self._correct:
            xL= sc.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
                                  _ars_hpx,nsamples=n,
                                  hxparams=(self._surfaceSigmaProfile,
                                            self._corr)))
        else:
            xL= sc.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.:<commit_msg>clean up los sampling in diskdf<commit_after>###############################################################################
#   diskdf.py: module that interprets (E,Lz) pairs in terms of a 
#              distribution function
#
#   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
###############################################################################
_EPSREL=10.**-14.
_NSIGMA= 4.
_INTERPDEGREE= 3
_RMIN=10.**-10.
_MAXD_REJECTLOS= 4.
import copy
import os, os.path
import cPickle as pickle
import math as m
import numpy as nu
import scipy as sc
import scipy.integrate as integrate
import scipy.interpolate as interpolate
from scipy import linalg
from scipy import stats
from scipy import optimize
from surfaceSigmaProfile import *
from galpy.orbit import Orbit
from galpy.util.bovy_ars import bovy_ars
from galpy.potential import PowerSphericalPotential
from galpy.actionAngle import actionAngleFlat, actionAnglePower
from galpy.actionAngle_src.actionAngleFlat import calcRapRperiFromELFlat #HACK
from galpy.actionAngle_src.actionAnglePower import \
    calcRapRperiFromELPower #HACK
_CORRECTIONSDIR=os.path.join(os.path.dirname(os.path.realpath(__file__)),'data')
_DEGTORAD= m.pi/180.
class diskdf:
    """Class that represents a disk DF"""
    def __init__(self,dftype='dehnen',
                 surfaceSigma=expSurfaceSigmaProfile,
                 profileParams=(1./3.,1.0,0.2),
                 correct=False,
                 beta=0.,**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)
        """
        self._dftype= dftype
        if isinstance(surfaceSigma,surfaceSigmaProfile):
            self._surfaceSigmaProfile= surfaceSigma
        else:
            self._surfaceSigmaProfile= surfaceSigma(profileParams)
        self._beta= beta
        if correct or kwargs.has_key('corrections') or kwargs.has_key('rmax') or kwargs.has_key('niter') or kwargs.has_key('npoints'):
            self._correct= True
            #Load corrections
            self._corr= DFcorrection(dftype=self.__class__,
                                     surfaceSigmaProfile=self._surfaceSigmaProfile,
                                     beta=beta,**kwargs)
        else:
            self._correct= False
        return None
    
    def __call__(self,*args):
        """
        NAME:
           __call__
        PURPOSE:
           evaluate the distribution function
        INPUT:
           either an orbit instance or E,Lz

           1) Orbit instance:
              a) Orbit instance alone: use vxvv member
              b) Orbit instance + t: call the Orbit instance

           2)
              E - energy (/vo^2)
              L - angular momentun (/ro/vo)
        OUTPUT:
           DF(orbit/E,L)
        HISTORY:
           2010-07-10 - Written - Bovy (NYU)
        """
        if isinstance(args[0],Orbit):
            if len(args) == 1:
                return sc.real(self.eval(*vRvTRToEL(args[0].vxvv[1],
                                                    args[0].vxvv[2],
                                                    args[0].vxvv[0],
                                                    self._beta)))
            else:
                vxvv= args[0](args[1])
                return sc.real(self.eval(*vRvTRToEL(vxvv[1],
                                                    vxvv[2],
                                                    vxvv[0],
                                                    self._beta)))
        else:
            return sc.real(self.eval(*args))

    def targetSigma2(self,R,log=False):
        """
        NAME:
           targetSigma2
        PURPOSE:
           evaluate the target Sigma_R^2(R)
        INPUT:
           R - radius at which to evaluate (/ro)
        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)
    
    def targetSurfacemass(self,R,log=False):
        """
        NAME:
           targetSurfacemass
        PURPOSE:
           evaluate the target surface mass at R
        INPUT:
           R - radius at which to evaluate
           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)
        
    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
           l - Galactic longitude (in deg, unless deg=False)
           deg= if False, l is in radians
           log - if True, return the log (default: False)
        OUTPUT:
           Sigma(d,l)
        HISTORY:
           2011-03-23 - Written - Bovy (NYU)
        """
        #Calculate R and phi
        if deg:
            lrad= l*_DEGTORAD
        else:
            lrad= l
        R, phi= _dlToRphi(d,lrad)
        #Evaluate Jacobian
        jac= _jacobian_rphi_dl(d,lrad,R=R,phi=phi)
        if log:
            return self._surfaceSigmaProfile.surfacemass(R,log=log)\
                +m.log(m.fabs(jac))+m.log(R)
            pass
        else:
            return self._surfaceSigmaProfile.surfacemass(R,log=log)\
                *m.fabs(jac)*R
        
    def surfacemassLOS(self,d,l,deg=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
           l - Galactic longitude (in deg, unless deg=False)
        OPTIONAL INPUT:
           nsigma - number of sigma to integrate the velocities over
        KEYWORDS:
           romberg - if True, use a romberg integrator (default: False)
           deg= if False, l is in radians
        OUTPUT:
           Sigma(d,l)
        HISTORY:
           2011-03-24 - Written - Bovy (NYU)
        """
        #Calculate R and phi
        if deg:
            lrad= l*_DEGTORAD
        else:
            lrad= l
        R, phi= _dlToRphi(d,lrad)
        #Evaluate Jacobian
        jac= _jacobian_rphi_dl(d,lrad,R=R,phi=phi)
        return self.surfacemass(R,romberg=romberg,nsigma=nsigma,
                                relative=relative)\
                                *m.fabs(jac)*R

    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)
           n= number of distances to sample
           maxd= maximum distance to consider (for the rejection sampling)
           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 l == 0.:
            maxSM= self.targetSurfacemass(0.)
        elif l >= m.pi/2. and l <= 3.*m.pi/2.:
            maxSM= self.targetSurfacemass(1.)
        elif l < m.pi/2. or l > 3.*m.pi/2.:
            if target:
                minR= optimize.fmin_bfgs(lambda x: \
                                             -x*self.targetSurfacemassLOS(x,l,
                                                                          deg=False),
                                         0.,disp=False)[0]
                maxSM= self.targetSurfacemassLOS(minR,l,deg=False)*minR
            else:
                minR= optimize.fmin_bfgs(lambda x: \
                                             -x*self.surfacemassLOS(x,l,
                                                                    deg=False),
                                         0.,disp=False)[0]
                maxSM= self.surfacemassLOS(minR,l,deg=False)*minR
        #Now rejection-sample
        if maxd is None:
            maxd= _MAXD_REJECTLOS
        out= []
        while len(out) < n:
            #sample
            prop= nu.random.random()*maxd
            if target:
                surfmassatprop= self.targetSurfacemassLOS(prop,l,deg=False)*prop
            else:
                surfmassatprop= self.surfacemassLOS(prop,l,deg=False)*prop
            if surfmassatprop/maxSM > nu.random.random(): #accept
                out.append(prop)
        return nu.array(out)

    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
           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.
        gamma= sc.sqrt(2./(1.+self._beta))           
        maxVT= optimize.brentq(_vtmaxEq,0.,1.,(R,self))
        maxVD= self(Orbit([R,maxVR,maxVT]))
        #Now rejection-sample
        if nsigma == None:
            nsigma= _NSIGMA
        out= []
        if target:
            sigma= m.sqrt(self.targetSigma2(R))
        else:
            sigma= m.sqrt(self.sigma2(R))
        while len(out) < n:
            #sample
            propvR= nu.random.normal()*_NSIGMA*sigma
            propvT= nu.random.normal()*_NSIGMA*sigma/gamma+1.
            VDatprop= self(Orbit([R,propvR,propvT]))
            if VDatprop/maxVD > nu.random.random(): #accept
                out.append(sc.array([propvR,propvT]))
        return nu.array(out)

    def sampleLOS(self,los,n=1,deg=True,maxd=None,nsigma=None,
                  target=True):
        """
        NAME:
           sampleLOS
        PURPOSE:
           sample along a given LOS
        INPUT:
           los - line of sight (in deg, unless deg=False)
           n= number of desired samples
           deg= los in degrees? (default=True)
           target= if True, use target surface mass and sigma2 profiles
                   (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 deg:
            l= los*_DEGTORAD
        else:
            l= los
        out= []
        #sample distances
        ds= self.sampledSurfacemassLOS(l,n=n,maxd=maxd,target=target)
        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=target)[0]
            out.append(Orbit([thisR,vv[0],vv[1],thisphi]))
        return out

    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 (/ro)
        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)
        sigmaR2= self.targetSigma2(R)
        sigmaR1= sc.sqrt(sigmaR2)
        logsigmaR2= sc.log(sigmaR2)
        gamma= sc.sqrt(2./(1.+self._beta))
        if relative:
            norm= 1.
        else:
            norm= sc.exp(logSigmaR)
        if romberg:
            return bovy_dblquad(_surfaceIntegrand,
                                gamma*R**self._beta/sigmaR1-nsigma,
                                gamma*R**self._beta/sigmaR1+nsigma,
                                lambda x: 0., lambda x: nsigma,
                                [R,self,logSigmaR,logsigmaR2,sigmaR1,gamma],
                                tol=10.**-8)/sc.pi*norm
        else:
            return integrate.dblquad(_surfaceIntegrand,
                                     gamma*R**self._beta/sigmaR1-nsigma,
                                     gamma*R**self._beta/sigmaR1+nsigma,
                                     lambda x: 0., lambda x: nsigma,
                                     (R,self,logSigmaR,logsigmaR2,sigmaR1,
                                      gamma),
                                     epsrel=_EPSREL)[0]/sc.pi*norm

    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 (/ro)
        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)
        sigmaR2= self.targetSigma2(R)
        sigmaR1= sc.sqrt(sigmaR2)
        logsigmaR2= sc.log(sigmaR2)
        gamma= sc.sqrt(2./(1.+self._beta))
        if relative:
            norm= 1.
        else:
            norm= sc.exp(logSigmaR+logsigmaR2)
        if romberg:
            return bovy_dblquad(_sigma2surfaceIntegrand,
                                gamma*R**self._beta/sigmaR1-nsigma,
                                gamma*R**self._beta/sigmaR1+nsigma,
                                lambda x: 0., lambda x: nsigma,
                                [R,self,logSigmaR,logsigmaR2,sigmaR1,gamma],
                                tol=10.**-8)/sc.pi*norm
        else:
            return integrate.dblquad(_sigma2surfaceIntegrand,
                                     gamma*R**self._beta/sigmaR1-nsigma,
                                     gamma*R**self._beta/sigmaR1+nsigma,
                                     lambda x: 0., lambda x: nsigma,
                                     (R,self,logSigmaR,logsigmaR2,sigmaR1,
                                      gamma),
                                     epsrel=_EPSREL)[0]/sc.pi*norm

    def sigma2(self,R,romberg=False,nsigma=None):
        """
        NAME:
           sigma2
        PURPOSE:
           calculate sigma_R^2 at R by marginalizing over velocity
        INPUT:
           R - radius at which to calculate sigma_R^2 density (/ro)
        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)/self.surfacemass(R,romberg,nsigma)

    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.:
            rperi, rap= calcRapRperiFromELFlat(E,L,vc=1.,ro=1.)
            aA= actionAngleFlat(rperi,0.,L/rperi)
        else:
            rperi, rap= calcRapRperiFromELPower(E,L,vc=1.,ro=1.)
            aA= actionAnglePower(rperi,0.,L/rperi,beta=self._beta)
        TR= aA.TR()[0]
        return (2.*m.pi/TR, rap, rperi)

    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
           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
           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 AttributeError("'sample' method for this disk df is not implemented")

class dehnendf(diskdf):
    """Dehnen's 'new' df"""
    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 / 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 - if True, correct the DF

           + DFcorrection kwargs (except for those already specified)
        OUTPUT:
        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 (/vo^2)
           L - angular momentun (/ro/vo)
       OUTPUT:
           DF(E,L)
        HISTORY:
           2010-03-10 - Written - Bovy (NYU)
           2010-03-28 - Moved to dehnenDF - Bovy (NYU)
        """
        #Calculate Re,LE, OmegaE
        if self._beta == 0.:
            xE= sc.exp(E-.5)
            logOLLE= sc.log(L/xE-1.)
        else: #non-flat rotation curve
            xE= (2.*E/(1.+1./self._beta))**(1./2./self._beta)
            logOLLE= self._beta*sc.log(xE)+sc.log(L/xE-xE**self._beta)
        if self._correct: 
            correction= self._corr.correct(xE,log=True)
        else:
            correction= sc.zeros(2)
        SRE2= self.targetSigma2(xE,log=True)+correction[1]
        return sc.exp(logsigmaR2-SRE2+self.targetSurfacemass(xE,log=True)-logSigmaR+sc.exp(logOLLE-SRE2)+correction[0])

    def sample(self,n=1,rrange=None,returnROrbit=True,returnOrbit=False,
               nphi=1.,los=None,losdeg=True,nsigma=None,target=True,
               maxd=None):
        """
        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)
           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)
        """
        if not los is None:
            return self.sampleLOS(los,deg=losdeg,n=n,maxd=maxd,
                                  nsigma=nsigma,target=target)
        #First sample xE
        if self._correct:
            xE= sc.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
                                  _ars_hpx,nsamples=n,
                                  hxparams=(self._surfaceSigmaProfile,
                                            self._corr)))
        else:
            xE= sc.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= sc.log(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)*sc.log(stats.uniform.rvs(size=n))/OR
        if self._correct:
            for ii in range(len(xE)):
                Lz[ii]*= self._corr.correct(xE[ii],log=False)[1]
        Lz+= LCE
        if not returnROrbit and not returnOrbit:
            out= [[e,l] for e,l in zip(E,Lz)]
        else:
            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.*m.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(sc.array([0.,tr]),self._psp)
                if returnOrbit:
                    vxvv= thisOrbit(tr).vxvv
                    if not los is None: #Sample along a given line of sight
                        if losdeg: l= los*_DEGTORAD
                        else: l= los
                        if l > (2.*m.pi): l-= 2.*m.pi
                        if l < 0: l+= 2.*m.pi
                        sinphil= 1./vxvv[0]*m.sin(l)
                        if m.fabs(sinphil) > 1.: continue
                        if stats.uniform.rvs() < 0.5:
                            phil= m.asin(sinphil)
                        else:
                            phil= m.pi-m.asin(sinphil)
                        phi= phil-l
                        if phi > (2.*m.pi): phi-= 2.*m.pi
                        if phi < 0: phi+= 2.*m.pi
                        #make sure this is on the right side of the los
                        if l >= 0. and l <= m.pi/2. and phi > m.pi: continue
                        elif l >= m.pi/2. and l <= m.pi and phi > m.pi/2.: \
                                continue
                        elif l >= m.pi and l <= 3.*m.pi/2. and phi < 3.*m.pi/2.: continue
                        elif l >= 3.*m.pi/2. and phi < m.pi: continue
                            
                        thisOrbit= Orbit(vxvv=sc.array([vxvv[0],vxvv[1],
                                                        vxvv[2],
                                                        phi]).reshape(4))
                    else:
                        thisOrbit= Orbit(vxvv=sc.array([vxvv[0],vxvv[1],vxvv[2],
                                                        stats.uniform.rvs()\
                                                            *m.pi*2.])\
                                             .reshape(4))
                else:
                    thisOrbit= Orbit(thisOrbit(tr))
                kappa= _kappa(thisOrbit.vxvv[0],self._beta)
                if not rrange == None:
                    if thisOrbit.vxvv[0] < rrange[0] \
                            or thisOrbit.vxvv[0] > rrange[1]:
                        continue
                if los is None:
                    mult= sc.ceil(kappa/wR*nphi)-1.
                    kappawR= kappa/wR*nphi-mult
                    while mult > 0:
                        if returnOrbit:
                            out.append(Orbit(vxvv=sc.array([vxvv[0],vxvv[1],
                                                            vxvv[2],
                                                            stats.uniform.rvs()*m.pi*2.]).reshape(4)))
                        else:
                            out.append(thisOrbit)
                        mult-= 1
                else:
                    if losdeg: l= los*_DEGTORAD
                    else: l= los
                    if l > (2.*m.pi): l-= 2.*m.pi
                    if l < 0: l+= 2.*m.pi
                    sinphil= 1./vxvv[0]*m.sin(l)
                    if m.fabs(sinphil) > 1.: continue
                    if stats.uniform.rvs() < 0.5:
                        phil= m.asin(sinphil)
                    else:
                        phil= m.pi-m.asin(sinphil)
                    phi= phil-l
                    if phi > (2.*m.pi): phi-= 2.*m.pi
                    if phi < 0: phi+= 2.*m.pi
                                #make sure this is on the right side of the los
                    if l >= 0. and l <= m.pi/2. and phi > m.pi: continue
                    elif l >= m.pi/2. and l <= m.pi and phi > m.pi/2.: \
                            continue
                    elif l >= m.pi and l <= 3.*m.pi/2. and phi < 3.*m.pi/2.: \
                            continue
                    elif l >= 3.*m.pi/2. and phi < m.pi: continue
                    #Calcualte dphidl
                    dphidl= m.fabs(1./vxvv[0]*m.cos(l)/m.cos(l+phi)-1.)
                    mult= sc.ceil(dphidl*kappa/wR*nphi)-1.
                    kappawR= kappa/wR*nphi-mult
                    while mult > 0:
                        out.append(Orbit(vxvv=sc.array([vxvv[0],vxvv[1],
                                                        vxvv[2],
                                                        phi]).reshape(4)))
                        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)]
        return out

class shudf(diskdf):
    """Shu's df (1969)"""
    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 / 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 - if True, correct the DF

           + DFcorrection kwargs (except for those already specified)
        OUTPUT:
        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)
        """
        #Calculate RL,LL, OmegaL
        if self._beta == 0.:
            xL= L
            logECLE= sc.log(-sc.log(xL)-0.5+E)
        else: #non-flat rotation curve
            xL= L**(1./(self._beta+1.))
            logECLE= sc.log(-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= sc.zeros(2)
        SRE2= self.targetSigma2(xL,log=True)+correction[1]
        return sc.exp(logsigmaR2-SRE2+self.targetSurfacemass(xL,log=True)-logSigmaR-sc.exp(logECLE-SRE2)+correction[0])

    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
           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)
           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)
        """
        if not los is None:
            return self.sampleLOS(los,n=n,maxd=maxd,nsigma=nsigma,target=target)
        #First sample xL
        if self._correct:
            xL= sc.array(bovy_ars([0.,0.],[True,False],[0.05,2.],_ars_hx,
                                  _ars_hpx,nsamples=n,
                                  hxparams=(self._surfaceSigmaProfile,
                                            self._corr)))
        else:
            xL= sc.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.:
{'fraction_non_alphanumeric': 0.0978859346936372, 'fraction_numerical': 0.014412743279451167, 'mean_word_length': 1.3031566812463975, 'pattern_counts': {'":': 0, '<': 49, '<?xml version=': 0, '>': 59, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 153, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4771547', 'n_tokens_mistral': 20420, 'n_tokens_neox': 19021, 'n_words': 4789}