Merge branch 'omhernquistdf'

HISTORY.txt

@@ -1,6 +1,15 @@
 v1.7 (XXXX-XX-XX)
 =================
 
+- Added a general framework for spherical distribution function as
+  well as implementations of a few well-known distribution functions,
+  including (a) Hernquist distribution functions that are isotropic,
+  have constant anisotropy, or have Osipkov-Merritt type anisotropy;
+  (b) an isotropic Plummer profile; (c) the King model (also added as
+  a potential as KingPotential). These distribution functions can be
+  evaluated and sampled, and their velocity moments can be
+  computed. Work started in #432 and continued on from there.
+
 - Added NFWPotential initialization method using rmax and vmax, the
   radius and alue at which the rotation curve peaks; also added rmax
   and vmax methods to NFWPotential to compute these quantities for any

doc/source/reference/df.rst

@@ -70,6 +70,7 @@ Anisotropic versions also exist:
    :maxdepth: 2
 
    Hernquist DF with constant anisotropy beta <dfhernquistconstantbeta.rst>
+   Hernquist DF with Osipkov-Merritt anisotropy <dfhernquistosipkov.rst>
 
 Two-dimensional, axisymmetric disk distribution functions
 ----------------------------------------------------------

doc/source/reference/dfhernquistosipkov.rst

@@ -0,0 +1,5 @@
+Anisotropic Hernquist DF of the Osipkov-Merritt type
+====================================================
+
+.. autoclass:: galpy.df.osipkovmerrittHernquistdf
+   :members: __init__

galpy/df/__init__.py

@@ -9,6 +9,7 @@
 from . import eddingtondf
 from . import isotropicHernquistdf
 from . import constantbetaHernquistdf
+from . import osipkovmerrittHernquistdf
 from . import kingdf
 from . import isotropicPlummerdf
 #
@@ -42,5 +43,6 @@
 eddingtondf= eddingtondf.eddingtondf
 isotropicHernquistdf= isotropicHernquistdf.isotropicHernquistdf
 constantbetaHernquistdf= constantbetaHernquistdf.constantbetaHernquistdf
+osipkovmerrittHernquistdf= osipkovmerrittHernquistdf.osipkovmerrittHernquistdf
 kingdf= kingdf.kingdf
 isotropicPlummerdf= isotropicPlummerdf.isotropicPlummerdf

galpy/df/constantbetadf.py

@@ -58,16 +58,20 @@ def _call_internal(self,*args):
         E, L, _= args
         return L**(-2*self._beta)*self.fE(E)
 
-    def _sample_eta(self,n=1):
+    def _sample_eta(self,r,n=1):
         """Sample the angle eta which defines radial vs tangential velocities"""
-        deta = 0.00005*numpy.pi
-        etas = (numpy.arange(0, numpy.pi, deta)+deta/2)
-        eta_pdf_cml = numpy.cumsum(numpy.power(numpy.sin(etas),1.-2.*self._beta))
-        eta_pdf_cml_norm = eta_pdf_cml / eta_pdf_cml[-1]
-        eta_icml_interp = scipy.interpolate.interp1d(eta_pdf_cml_norm, etas, 
-            bounds_error=False, fill_value='extrapolate')
-        eta_samples = eta_icml_interp(numpy.random.uniform(size=n))
-        return eta_samples
+        if not hasattr(self,'_coseta_icmf_interp'):
+            # Cumulative dist for cos(eta) =
+            # 0.5 + x 2F1(0.5,beta,1.5,x^2)/sqrt(pi)/Gamma(1-beta)*Gamma(1.5-beta)
+            cosetas= numpy.linspace(-1.,1.,20001)
+            coseta_cmf= cosetas*special.hyp2f1(0.5,self._beta,1.5,cosetas**2.)\
+                /numpy.sqrt(numpy.pi)/special.gamma(1.-self._beta)\
+                *special.gamma(1.5-self._beta)+0.5
+            self._coseta_icmf_interp= scipy.interpolate.interp1d(\
+                                coseta_cmf,cosetas,
+                                bounds_error=False,fill_value='extrapolate')
+        return numpy.arccos(self._coseta_icmf_interp(\
+                                                numpy.random.uniform(size=n)))
 
     def _p_v_at_r(self,v,r):
         return self.fE(evaluatePotentials(self._pot,r,0,use_physical=False)\

galpy/df/osipkovmerrittHernquistdf.py

@@ -0,0 +1,94 @@
+# Class that implements the anisotropic spherical Hernquist DF with radially
+# varying anisotropy of the Osipkov-Merritt type
+import numpy
+from ..util import conversion
+from ..potential import evaluatePotentials, HernquistPotential
+from .osipkovmerrittdf import osipkovmerrittdf
+
+class osipkovmerrittHernquistdf(osipkovmerrittdf):
+    """Class that implements the anisotropic spherical Hernquist DF with radially varying anisotropy of the Osipkov-Merritt type
+    
+    .. math::
+
+        \\beta(r) = \\frac{1}{1+r_a^2/r^2}
+
+    with :math:`r_a` the anistropy radius.
+
+    """
+    def __init__(self,pot=None,ra=1.4,ro=None,vo=None):
+        """
+        NAME:
+
+            __init__
+
+        PURPOSE:
+
+            Initialize a Hernquist DF with Osipkov-Merritt anisotropy
+
+        INPUT:
+
+            pot - Hernquist potential which determines the DF
+
+            ra - anisotropy radius (can be a Quantity)
+
+        OUTPUT:
+
+            None
+
+        HISTORY:
+
+            2020-11-12 - Written - Bovy (UofT)
+        """
+        assert isinstance(pot,HernquistPotential), \
+            'pot= must be potential.HernquistPotential'
+        osipkovmerrittdf.__init__(self,pot=pot,ra=ra,ro=ro,vo=vo)
+        self._psi0= -evaluatePotentials(self._pot,0,0,use_physical=False)
+        self._GMa = self._psi0*self._pot.a**2.
+        self._a2overra2= self._pot.a**2./self._ra2
+        self._fQnorm= 1./numpy.sqrt(2.)/(2*numpy.pi)**3/self._GMa**1.5
+
+    def fQ(self,Q):
+        """
+        NAME:
+
+            fQ
+
+        PURPOSE
+
+            Calculate the f(Q) portion of an Osipkov-Merritt Hernquist distribution function
+
+        INPUT:
+
+            Q - The Osipkov-Merritt 'energy' E-L^2/[2ra^2] (can be Quantity)
+
+        OUTPUT:
+
+            fQ - The value of the f(Q) portion of the DF
+
+        HISTORY:
+
+            2020-11-12 - Written - Bovy (UofT)
+
+        """
+        Qtilde= conversion.parse_energy(Q,vo=self._vo)/self._psi0
+        # Handle potential Q outside of bounds
+        Qtilde_out = numpy.where(numpy.logical_or(Qtilde<0,Qtilde>1))[0]
+        if len(Qtilde_out)>0:
+            # Dummy variable now and 0 later, prevents numerical issues
+            Qtilde[Qtilde_out]=0.5
+        sqrtQtilde= numpy.sqrt(Qtilde)
+        # The 'ergodic' part
+        fQ= sqrtQtilde/(1-Qtilde)**2.\
+            *((1.-2.*Qtilde)*(8.*Qtilde**2.-8.*Qtilde-3.)\
+              +((3.*numpy.arcsin(sqrtQtilde))\
+                /numpy.sqrt(Qtilde*(1.-Qtilde))))
+        # The other part
+        fQ+= 8.*self._a2overra2*sqrtQtilde*(1.-2.*Qtilde)
+        if len(Qtilde_out) > 0:
+            fQ[Qtilde_out]= 0.
+        return self._fQnorm*fQ
+
+    def _icmf(self,ms):
+        '''Analytic expression for the normalized inverse cumulative mass 
+        function. The argument ms is normalized mass fraction [0,1]'''
+        return self._pot.a*numpy.sqrt(ms)/(1-numpy.sqrt(ms))

galpy/df/osipkovmerrittdf.py

@@ -0,0 +1,115 @@
+# Class that implements anisotropic DFs of the Osipkov-Merritt type
+import numpy
+from scipy import integrate, special
+from ..util import conversion
+from ..potential import evaluatePotentials
+from .sphericaldf import anisotropicsphericaldf
+
+class osipkovmerrittdf(anisotropicsphericaldf):
+    """Class that implements anisotropic DFs of the Osipkov-Merritt type with radially varying anisotropy
+    
+    .. math::
+
+        \\beta(r) = \\frac{1}{1+r_a^2/r^2}
+
+    with :math:`r_a` the anistropy radius.
+
+    """
+    def __init__(self,pot=None,ra=1.4,scale=None,ro=None,vo=None):
+        """
+        NAME:
+
+            __init__
+
+        PURPOSE:
+
+            Initialize a DF with Osipkov-Merritt anisotropy
+
+        INPUT:
+
+            pot - Hernquist potential which determines the DF
+
+            ra - anisotropy radius (can be a Quantity)
+
+            scale - Characteristic scale radius to aid sampling calculations. 
+                Not necessary, and will also be overridden by value from pot 
+                if available.
+
+        OUTPUT:
+
+            None
+
+        HISTORY:
+
+            2020-11-12 - Written - Bovy (UofT)
+
+        """
+        anisotropicsphericaldf.__init__(self,pot=pot,scale=scale,ro=ro,vo=vo)
+        self._ra= -conversion.parse_length(ra,ro=self._ro)
+        self._ra2= self._ra**2.
+
+    def _call_internal(self,*args):
+        """
+        NAME:
+
+            _call_internal
+
+        PURPOSE:
+
+            Evaluate the DF for an Osipkov-Merritt-anisotropy DF
+
+        INPUT:
+
+            E - The energy
+
+            L - The angular momentum
+
+        OUTPUT:
+
+            fH - The value of the DF
+
+        HISTORY:
+
+            2020-11-12 - Written - Bovy (UofT)
+
+        """
+        E, L, _= args
+        return self.fQ(-E-0.5*L**2./self._ra2)
+
+    def _sample_eta(self,r,n=1):
+        """Sample the angle eta which defines radial vs tangential velocities"""
+        # cumulative distribution of x = cos eta satisfies
+        # x/(sqrt(A+1 -A* x^2)) = 2 b - 1 = c
+        # where b \in [0,1] and A = (r/ra)^2
+        # Solved by
+        # x = c sqrt(1+[r/ra]^2) / sqrt( [r/ra]^2 c^2 + 1 ) for c > 0 [b > 0.5]
+        # and symmetric wrt c
+        c= numpy.random.uniform(size=n)
+        x= c*numpy.sqrt(1+r**2./self._ra2)/numpy.sqrt(r**2./self._ra2*c**2.+1)
+        x*= numpy.random.choice([1.,-1.],size=n)
+        return numpy.arccos(x)
+
+    def _p_v_at_r(self,v,r):
+        """p( v*sqrt[1+r^2/ra^2*sin^2eta] | r) used in sampling """
+        return self.fQ(-evaluatePotentials(self._pot,r,0,use_physical=False)\
+                       -0.5*v**2.)*v**2.
+
+    def _sample_v(self,r,eta,n=1):
+        """Generate velocity samples"""
+        # Use super-class method to obtain v*[1+r^2/ra^2*sin^2eta]
+        out= super(osipkovmerrittdf,self)._sample_v(r,eta,n=n)
+        # Transform to v
+        return out/numpy.sqrt(1.+r**2./self._ra2*numpy.sin(eta)**2.)
+
+    def _vmomentdensity(self,r,n,m):
+         if m%2 == 1 or n%2 == 1:
+             return 0.
+         psir= -evaluatePotentials(self._pot,r,0,use_physical=False)
+         return 2.*numpy.pi*integrate.quad(lambda v: v**(2.+m+n)
+                                    *self.fQ(-evaluatePotentials(self._pot,r,0,
+                                                         use_physical=False)
+                                             -0.5*v**2.),
+                             0.,self._vmax_at_r(self._pot,r))[0]\
+            *special.gamma(m/2.+1.)*special.gamma((n+1)/2.)/\
+            special.gamma(0.5*(m+n+3.))/(1+r**2./self._ra2)**(m/2+1)
+

galpy/df/sphericaldf.py

@@ -13,7 +13,7 @@
 #     to implement a bunch of functions:
 #       * _call_internal(self,*args,**kwargs): which returns the DF as a
 #                                              function of (E,L,Lz)
-#       * _sample_eta(self,n=1): to sample the velocity angle
+#       * _sample_eta(self,r,n=1): to sample the velocity angle at r
 #       * _p_v_at_r(self,v,r): whcih returns p(v|r)
 #     constantbetadf is an example of this
 #
@@ -276,8 +276,8 @@ def beta(self,r):
          
             2020-09-04 - Written - Bovy (UofT)
         """
-        return 1.-self.sigmat(r,use_physical=False)**2./2.\
-            /self.sigmar(r,use_physical=False)**2.
+        r= conversion.parse_length(r,ro=self._ro)
+        return 1.-self._vmomentdensity(r,0,2)/2./self._vmomentdensity(r,2,0)
 
 ############################### SAMPLING THE DF################################
     def sample(self,R=None,z=None,phi=None,n=1,return_orbit=True):
@@ -345,8 +345,8 @@ def sample(self,R=None,z=None,phi=None,n=1,return_orbit=True):
                 phi= phi*numpy.ones(n) \
                     if not hasattr(phi,'__len__') or len(phi) < n \
                     else phi
-        v = self._sample_v(r,n=n)
-        eta,psi = self._sample_velocity_angles(n=n)
+        eta,psi = self._sample_velocity_angles(r,n=n)
+        v = self._sample_v(r,eta,n=n)
         vr = v*numpy.cos(eta)
         vtheta = v*numpy.sin(eta)*numpy.cos(psi)
         vT = v*numpy.sin(eta)*numpy.sin(psi)
@@ -409,19 +409,19 @@ def _sample_position_angles(self,n=1):
         theta_samples = numpy.arccos(1.-2*numpy.random.uniform(size=n))
         return phi_samples,theta_samples
 
-    def _sample_v(self,r,n=1):
-        """Generate velocity samples"""
+    def _sample_v(self,r,eta,n=1):
+        """Generate velocity samples: typically the total velocity, but not for OM"""
         if not hasattr(self,'_v_vesc_pvr_interpolator'):
             self._v_vesc_pvr_interpolator = self._make_pvr_interpolator()
         return self._v_vesc_pvr_interpolator(\
                     numpy.log10(r/self._scale),numpy.random.uniform(size=n),
                     grid=False)*self._vmax_at_r(self._pot,r)
 
-    def _sample_velocity_angles(self,n=1):
+    def _sample_velocity_angles(self,r,n=1):
         """Generate samples of angles that set radial vs tangential 
         velocities"""
-        eta_samples = self._sample_eta(n)
-        psi_samples = numpy.random.uniform(size=n)*2*numpy.pi
+        eta_samples= self._sample_eta(r,n)
+        psi_samples= numpy.random.uniform(size=n)*2*numpy.pi
         return eta_samples,psi_samples
 
     def _vmax_at_r(self,pot,r,**kwargs):
@@ -568,7 +568,7 @@ def _vmomentdensity(self,r,n,m):
             *special.gamma(m//2+1)*special.gamma(n//2+0.5)\
             /special.gamma(m//2+n//2+1.5)
 
-    def _sample_eta(self,n=1):
+    def _sample_eta(self,r,n=1):
         """Sample the angle eta which defines radial vs tangential velocities"""
         return numpy.arccos(1.-2.*numpy.random.uniform(size=n))