added a Python scf library

This commit is contained in:
Yohai 2019-11-02 15:12:57 -04:00
parent 1718e6287b
commit 6e45215356

44
src/scf.py Normal file
View file

@ -0,0 +1,44 @@
from scipy.special import eval_gegenbauer, sph_harm
import numpy as np
class Scf:
def __init__(self, nmax, lmax, mmax=None, length_scale=1, basis='hernquist'):
self.nmax = nmax
self.lmax = lmax
self.length_scale = length_scale
self.basis = basis
self.A = None
if not mmax is None: raise NotImplemented
if (basis.lower()=='hernquist'): self.basis_phi = self.basis_phi_hernquist
else: raise RuntimeError('Unknown basis set "%s".' % basis)
def set_coefficients(self, A):
if len(A) != (self.nmax+1)*(self.lmax+1)*(self.lmax+2)//2: raise ValueError('Coefficient list has unexpected length.')
self.A = A
def calc_gravity(self, x, y, z, potential=True, force=False):
r = np.sqrt(x**2 + y**2 + z**2)
theta = np.arccos(z/r)
phi = np.arctan2(y, x)
if force: raise NotImplemented
if self.A is None: raise RuntimeError('Coefficient list must be set before gravity could be calculated.')
pot = 0
for n in range(self.nmax+1):
for l in range(self.lmax+1):
for m in range(l+1):
i = n*(self.lmax+1)*(self.lmax+2)//2 + l*(l+1)//2 + m
pot_ = np.real(self.A[i]*self.basis_phi(n, l, m, r*self.length_scale, theta, phi))
if m == 0:
pot += pot_
else:
pot += 2*pot_
pot *= self.length_scale
return pot
def calc_coefficients(self, x, y, z):
raise NotImplemented
def basis_phi_hernquist(self, n, l, m, r, theta, phi):
xi = (r-1)/(r+1)
W_nl = eval_gegenbauer(n, 2*l+1.5, xi)
return - r**l*(1+r)**(-2*l-1)*W_nl*2*np.sqrt(np.pi)*sph_harm(m, l, phi, theta)