scf Nbody compute: make mass optional and allow single value

galpy/potential/SCFPotential.py

@@ -1,8 +1,8 @@
 import hashlib
 import numpy
 from numpy.polynomial.legendre import leggauss
-from scipy.special import lpmn,lpmv
-from scipy.special import gammaln, gamma, gegenbauer
+from scipy.special import lpmn
+from scipy.special import gammaln, gamma
 from ..util import coords, conversion
 from .Potential import Potential
 
@@ -572,7 +572,7 @@ def _dC(xi, N, L):
     CC *= 2*(2*l + 3./2)
     return CC
 
-def scf_compute_coeffs_spherical_nbody(pos,mass,N,a=1.):
+def scf_compute_coeffs_spherical_nbody(pos,N,mass=1.,a=1.):
     """        
     NAME:
 
@@ -586,10 +586,10 @@ def scf_compute_coeffs_spherical_nbody(pos,mass,N,a=1.):
 
        pos - positions of particles in rectangular coordinates with shape [3,n]
            
-       mass - mass of particles
-
        N - size of the Nth dimension of the expansion coefficients
 
+       mass= (1.) mass of particles (scalar or array with size n)
+
        a= (1.) parameter used to scale the radius
 
     OUTPUT:
@@ -602,7 +602,7 @@ def scf_compute_coeffs_spherical_nbody(pos,mass,N,a=1.):
 
        2021-02-22 - Sped-up - Bovy (UofT)
 
-    """ 
+    """
     Acos = numpy.zeros((N,1,1), float)
     Asin = None
     r= numpy.sqrt(pos[0]**2+pos[1]**2+pos[2]**2)
@@ -673,7 +673,7 @@ def integrand(xi):
     Acos[n,0,0] = 2*K*integrated
     return Acos, Asin    
 
-def scf_compute_coeffs_axi_nbody(pos,mass,N,L,a=1.):
+def scf_compute_coeffs_axi_nbody(pos,N,L,mass=1.,a=1.):
     """        
     NAME:
 
@@ -685,14 +685,14 @@ def scf_compute_coeffs_axi_nbody(pos,mass,N,L,a=1.):
 
     INPUT:
 
-       pos - positions of particles with shape [3,N]
+       pos - positions of particles in rectangular coordinates with shape [3,n]
            
-       mass - mass of particles
-
        N - size of the Nth dimension of the expansion coefficients
 
        L - size of the Lth dimension of the expansion coefficients
 
+       mass= (1.) mass of particles (scalar or array with size n)
+
        a= (1.) parameter used to scale the radius
 
     OUTPUT:
@@ -703,9 +703,10 @@ def scf_compute_coeffs_axi_nbody(pos,mass,N,L,a=1.):
 
        2021-02-22 - Written based on general code - Bovy (UofT)
 
-    """ 
-    r = numpy.sqrt(pos[0]**2+pos[1]**2+pos[2]**2)
+    """
+    r= numpy.sqrt(pos[0]**2+pos[1]**2+pos[2]**2)
     costheta = pos[2]/r
+    mass= numpy.atleast_1d(mass)
     Acos, Asin= numpy.zeros([N,L,1]), None
     Pll= numpy.ones(len(r)) # Set up Assoc. Legendre recursion
     # (n,l) dependent constant
@@ -806,7 +807,7 @@ 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.):
+def scf_compute_coeffs_nbody(pos,N,L,mass=1.,a=1.):
     """        
     NAME:
 
@@ -818,14 +819,14 @@ def scf_compute_coeffs_nbody(pos,mass,N,L,a=1.):
 
     INPUT:
 
-       pos - positions of particles with shape [3,N]
+       pos - positions of particles in rectangular coordinates with shape [3,n]
            
-       mass - 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
 
+       mass= (1.) mass of particles (scalar or array with size n)
+
        a= (1.) parameter used to scale the radius
 
     OUTPUT:
@@ -837,10 +838,11 @@ def scf_compute_coeffs_nbody(pos,mass,N,L,a=1.):
        2020-11-18 - Written - Morgan Bennett (UofT)
 
     """ 
-    r = numpy.sqrt(pos[0]**2+pos[1]**2+pos[2]**2)
-    phi = numpy.arctan2(pos[1],pos[0])
-    costheta = pos[2]/r
+    r= numpy.sqrt(pos[0]**2+pos[1]**2+pos[2]**2)
+    phi= numpy.arctan2(pos[1],pos[0])
+    costheta= pos[2]/r
     sintheta= numpy.sqrt(1.-costheta**2.)
+    mass= numpy.atleast_1d(mass)
     Acos, Asin= numpy.zeros([N,L,L]), numpy.zeros([N,L,L])
     Pll= numpy.ones(len(r)) # Set up Assoc. Legendre recursion
     # (n,l) dependent constant

tests/test_scf.py

@@ -223,13 +223,23 @@ def test_scf_compute_spherical_nbody_hernquist():
     c= numpy.zeros((nsamp,Norder,1,1))
     s= numpy.zeros((nsamp,Norder,1,1))
     for i in range(nsamp):
-        c[i],s[i]= potential.scf_compute_coeffs_spherical_nbody(mass=m*numpy.ones(N),pos=positions[i],N=Norder,a=ah)
+        c[i],s[i]= potential.scf_compute_coeffs_spherical_nbody(\
+                                positions[i],Norder,mass=m*numpy.ones(N),a=ah)
 
-    cc,ss= potential.scf_compute_coeffs_spherical(N=Norder,a=ah,dens=hern.dens)
+    cc,ss= potential.scf_compute_coeffs_spherical(hern.dens,Norder,a=ah)
 
     # Check that the difference between the coefficients is within the standard deviation
     assert (cc-numpy.mean(c,axis=0)<numpy.std(c,axis=0)).all()
 
+    # Repeat test for single mass
+    c= numpy.zeros((nsamp,Norder,1,1))
+    s= numpy.zeros((nsamp,Norder,1,1))
+    for i in range(nsamp):
+        c[i],s[i]= potential.scf_compute_coeffs_spherical_nbody(\
+                                positions[i],Norder,mass=m,a=ah)
+    assert (cc-numpy.mean(c,axis=0)<numpy.std(c,axis=0)).all()
+    return None  
+
 ## Tests how nbody calculation compares to density calculation for scf_compute_coeff
 def test_scf_compute_axi_nbody_twopowertriaxial():
     N= int(1e5)
@@ -259,11 +269,20 @@ def test_scf_compute_axi_nbody_twopowertriaxial():
     cc, ss= potential.scf_compute_coeffs_axi(tptp.dens,Norder,Lorder,a=ah)
     c,s= numpy.zeros((2, nsamp, Norder, Lorder,1))
     for i,p in enumerate(positions):
-        c[i],s[i]= potential.scf_compute_coeffs_axi_nbody(p,m*numpy.ones(N),
-                                                          Norder,Lorder,a=ah)
+        c[i],s[i]= potential.scf_compute_coeffs_axi_nbody(p,Norder,Lorder,
+                                                          mass=m*numpy.ones(N),
+                                                          a=ah)
 
     # Check that the difference between the coefficients is within two standard deviations
     assert (cc-(numpy.mean(c,axis=0))<=(2.*numpy.std(c,axis=0))).all()
+
+    # Repeat test for single mass
+    c,s= numpy.zeros((2, nsamp, Norder, Lorder,1))
+    for i,p in enumerate(positions):
+        c[i],s[i]= potential.scf_compute_coeffs_axi_nbody(p,Norder,Lorder,
+                                                          mass=m,a=ah)
+    assert (cc-(numpy.mean(c,axis=0))<=(2.*numpy.std(c,axis=0))).all()
+    return None
 
 ## Tests how nbody calculation compares to density calculation for scf_compute_coeff
 def test_scf_compute_nbody_twopowertriaxial():
@@ -295,12 +314,21 @@ def test_scf_compute_nbody_twopowertriaxial():
     cc, ss= potential.scf_compute_coeffs(tptp.dens,Norder,Lorder,a=ah)
     c,s= numpy.zeros((2, nsamp, Norder, Lorder, Lorder))
     for i,p in enumerate(positions):
-        c[i],s[i]= potential.scf_compute_coeffs_nbody(p,m*numpy.ones(N),
-                                                      Norder,Lorder,a=ah)
+        c[i],s[i]= potential.scf_compute_coeffs_nbody(p,Norder,Lorder,
+                                                      mass=m*numpy.ones(N),
+                                                      a=ah)
 
     # Check that the difference between the coefficients is within two standard deviations
     assert (cc-(numpy.mean(c,axis=0))<=(2.*numpy.std(c,axis=0))).all()
 
+    # Repeat test for single mass
+    c,s= numpy.zeros((2, nsamp, Norder, Lorder, Lorder))
+    for i,p in enumerate(positions):
+        c[i],s[i]= potential.scf_compute_coeffs_nbody(p,Norder,Lorder,
+                                                  mass=m,a=ah)        
+    assert (cc-(numpy.mean(c,axis=0))<=(2.*numpy.std(c,axis=0))).all()
+    return None
+
 def test_scf_compute_nfw(): 
     Acos, Asin = potential.scf_compute_coeffs_spherical(rho_NFW, 10)
     spherical_coeffsTest(Acos, Asin)