Calculation of SCF expansion based on an N-body representation of the density

galpy/potential/SCFPotential.py

@@ -1,8 +1,8 @@
 import hashlib
 import numpy
 from numpy.polynomial.legendre import leggauss
-from scipy.special import lpmn
-from scipy.special import gammaln
+from scipy.special import lpmn,lpmv
+from scipy.special import gammaln, gamma, gegenbauer
 from ..util import coords, conversion
 from .Potential import Potential
 
@@ -506,7 +506,10 @@ def OmegaP(self):
 
 
 def _xiToR(xi, a =1):
-    return a*numpy.divide((1. + xi),(1. - xi))    
+    return a*numpy.divide((1. + xi),(1. - xi))  
+
+def _RToxi(r, a=1):
+    return numpy.divide((r/a-1.),(r/a+1.))
 
 
 def _C(xi, N,L, alpha = lambda x: 2*x + 3./2):
@@ -546,7 +549,33 @@ def _dC(xi, N, L):
     CC[0, :] = 0
     CC *= 2*(2*l + 3./2)
     return CC
-
+
+def scf_compute_coeffs_spherical_nbody(pos,m,N,a=1.):
+
+        '''
+        INPUT:
+            pos - position of particles in your nbody snapshot
+
+            m - masses of particles
+
+            N - size of expansion coefficients                                                                         
+
+            a - parameter used to shift the basis functions 
+
+        '''               
+        Acos = numpy.zeros((N,1,1), float)
+        Asin = None
+
+        r= numpy.sqrt(pos[0]**2+pos[1]**2+pos[2]**2)
+        Cs= numpy.array([_C(_RToxi(ri,a=a),N,1)[:,0] for ri in r])
+        RhoSum= numpy.sum((m/(4.*numpy.pi)/(r/a+1))[:,None]*Cs,axis=0)
+        n = numpy.arange(0,N)
+        K = 16*numpy.pi*(n + 3./2)/((n + 2)*(n + 1)*(1 + n*(n + 3.)/2.))
+        Acos[n,0,0] = 2*K*RhoSum
+
+        return Acos, Asin 
+
+
 def scf_compute_coeffs_spherical(dens, N, a=1., radial_order=None):
         """
         NAME:
@@ -685,6 +714,108 @@ def integrand(xi, costheta):
         Acos[:,:,0] = 2*I**-1 * integrated*constants
 
         return Acos, Asin
+
+def scf_compute_coeffs_nbody(pos,mass,N,L,a=1., radial_order=None, costheta_order=None, phi_order=None):
+
+        """        
+        NAME:
+
+           scf_compute_coeffs
+
+        PURPOSE:
+
+           Numerically compute the expansion coefficients for a given triaxial density
+
+        INPUT:
+
+           pos - Positions of particles
+
+           m - mass of particles
+
+           N - size of the Nth dimension of the expansion coefficients
+
+           L - size of the Lth and Mth dimension of the expansion coefficients
+
+           a - parameter used to shift the basis functions
+
+           radial_order - Number of sample points of the radial integral. If None, radial_order=max(20, N + 3/2L + 1)
+
+           costheta_order - Number of sample points of the costheta integral. If None, If costheta_order=max(20, L + 1)
+
+           phi_order - Number of sample points of the phi integral. If None, If costheta_order=max(20, L + 1)
+
+        OUTPUT:
+
+           (Acos,Asin) - Expansion coefficients for density dens that can be given to SCFPotential.__init__
+
+        HISTORY:
+
+           2020-11-18 - Written - Morgan Bennett
+
+        """ 
+
+        n = numpy.arange(0,N)
+        l = numpy.arange(0,L)
+        m = numpy.arange(0,L)
+
+        r = numpy.sqrt(pos[0]**2+pos[1]**2+pos[2]**2)
+        phi = numpy.arctan2(pos[1],pos[0])
+        costheta = pos[2]/r
+        '''
+        Plm= numpy.zeros([len(costheta),1,L,L])
+        for i,ll in enumerate(l):
+            for j,mm in enumerate(m):
+                Plm[:,0,j,i]= lpmv(ll,mm,costheta)
+
+        cosmphi= numpy.cos(phi[:,None]*m[None,:])[:,None,None,:]
+        sinmphi= numpy.sin(phi[:,None]*m[None,:])[:,None,None,:]
+
+        Ylm= (numpy.sqrt((2.*l[:,None]+1)*gamma(l[:,None]-m[None,:]+1)/gamma(l[:,None]+m[None,:]+1))[None,None,:,:]*Plm)[None,:,:,:,:]*numpy.array([cosmphi,sinmphi])
+        Ylm= numpy.nan_to_num(Ylm)
+
+        C= [[gegenbauer(nn,2.*ll+1.5) for ll in l] for nn in n] 
+        Cn= numpy.zeros((1,len(r),N,L,1))
+        for i in range(N):
+            for j in range(L):
+                Cn[0,:,i,j,0]= C[i][j]((r/a-1)/(r/a+1))
+
+        rl= ((r[:,None]/a)**l[None,:])[None,:,None,:,None]
+        r12l1= ((1+(r[:,None]/a))**(2.*l[None,:]+1))[None,:,None,:,None]
+
+        phinlm= -rl/r12l1*Cn*Ylm
+
+        Sum= numpy.sum(mass[None,:,None,None,None]*phinlm,axis=1)
+        Knl= 0.5*n[:,None]*(n[:,None]+4.*l[None,:]+3.)+(l[None,:]+1)*(2.*l[None,:]+1.)
+
+        Inl= (-Knl*4.*numpy.pi/2.**(8.*l[None,:]+6.)*gamma(n[:,None]+4.*l[None,:]+3.)/gamma(n[:,None]+1)/(n[:,None]+2.*l[None,:]+1.5)/gamma(2.*l[None,:]+1.5)**2)[None,:,:,None]
+
+        Anlm= Inl**(-1)*Sum'''
+        Anlm= numpy.zeros([2,L,L,L])
+        for i,nn in enumerate(n):
+            for j,ll in enumerate(l):
+                for k,mm in enumerate(m[:j+1]):
+
+                    Plm= lpmv(mm,ll,costheta)
+
+                    cosmphi= numpy.cos(phi*mm)
+                    sinmphi= numpy.sin(phi*mm)
+
+                    Ylm= (numpy.sqrt((2.*ll+1)*gamma(ll-mm+1)/gamma(ll+mm+1))*Plm)[None,:]*numpy.array([cosmphi,sinmphi])
+                    Ylm= numpy.nan_to_num(Ylm)
+
+                    C= gegenbauer(nn,2.*ll+1.5)
+                    Cn= C((r/a-1)/(r/a+1))
+
+                    phinlm= (-(r/a)**ll/(r/a+1)**(2.*ll+1)*Cn)[None,:]*Ylm
+
+                    Sum= numpy.sum(mass[None,:]*phinlm,axis=1)
+
+                    Knl= 0.5*nn*(nn+4.*ll+3.)+(ll+1)*(2.*ll+1.)
+                    Inl= (-Knl*4.*numpy.pi/2.**(8.*ll+6.)*gamma(nn+4.*ll+3.)/gamma(nn+1)/(nn+2.*ll+1.5)/gamma(2.*ll+1.5)**2)
+
+                    Anlm[:,i,j,k]= Inl**(-1)*Sum
+
+        return 2.*Anlm
 
 def scf_compute_coeffs(dens, N, L, a=1., radial_order=None, costheta_order=None, phi_order=None):
         """        

galpy/potential/__init__.py

@@ -104,6 +104,8 @@
 turn_physical_on= Potential.turn_physical_on
 _dim= Potential._dim
 _isNonAxi= Potential._isNonAxi
+scf_compute_coeffs_spherical_nbody = SCFPotential.scf_compute_coeffs_spherical_nbody
+scf_compute_coeffs_nbody = SCFPotential.scf_compute_coeffs_nbody
 scf_compute_coeffs_spherical = SCFPotential.scf_compute_coeffs_spherical
 scf_compute_coeffs_axi = SCFPotential.scf_compute_coeffs_axi
 scf_compute_coeffs = SCFPotential.scf_compute_coeffs

tests/test_scf.py

@@ -4,9 +4,10 @@
 import sys
 import numpy
 from galpy import potential
-from galpy.potential import SCFPotential
+from galpy.potential import SCFPotential,scf_compute_coeffs,scf_compute_coeffs_nbody
 from galpy.util import coords
 from galpy.orbit import Orbit
+from galpy import df
 _TRAVIS= bool(os.getenv('TRAVIS'))
 
 EPS = 1e-13 ## default epsilon
@@ -365,6 +366,42 @@ def test_phiforceMatches_nfw():
     scf = SCFPotential(amp=1, Acos=Acos, Asin=Asin)
     assertmsg = "Comparing the azimuth force of NFW Potential with SCF fails at R={0}, Z={1}, phi={2}"
     compareFunctions(nfw.phiforce,scf.phiforce, assertmsg)
+
+##############Test nbody coefficiennt calculation###############   
+
+def test_spherical():
+    #Test that the SCF coefficient estimators make a spherical potential
+
+    N= int(1e6)
+    Mh= 10.
+    ah= 50./8.
+    m= Mh/N
+
+    pothern= potential.HernquistPotential(amp=2*Mh,a=ah)
+    hdf= df.isotropicHernquistdf(pothern)
+    samp= hdf.sample(n=N)
+
+    positions= numpy.array([samp.x(),samp.y(),samp.z()])
+
+    Acdens, Asdens= scf_compute_coeffs(pothern.dens,10,10,a=50./8.)
+    Acnbdy, Asnbdy= scf_compute_coeffs_nbody(positions,m*numpy.ones(N),10,10,a=50./8.)
+
+    potdens= SCFPotential(Acos=Acdens,Asin=Asdens,a=50./8.)
+    potnbdy= SCFPotential(Acos=Acnbdy,Asin=Asnbdy,a=50./8.)
+
+    potdens.turn_physical_off()
+    potnbdy.turn_physical_off()
+    pothern.turn_physical_off()
+
+    r= numpy.linspace(0,1,100)
+    z= numpy.linspace(-0.5,0.5,100)
+
+    phern= pothern.dens(r[:,None],z[None,:])
+    pdens= potdens.dens(r[:,None],z[None,:])
+    pnbdy= potnbdy.dens(r[:,None],z[None,:])
+
+    assert numpy.max(numpy.abs((phern-pdens)/phern))<1.e-3, 'scf_compute_coeffs did not meet the required tolerance of 1e-3'
+    assert numpy.max(numpy.abs((phern-pnbdy)/phern))<0.2, 'scf_compute_coeffs_nbody did not meet the required tolerance of 0.1'
 
 ##############GENERIC FUNCTIONS BELOW###############
 
@@ -484,3 +521,4 @@ def density1(R, z=0, phi=0.):
     return h.dens(R, z, phi)*(1 + numpy.cos(theta) + numpy.cos(theta)**2.)*(1 + numpy.cos(phi) + numpy.sin(phi))
 
 
+