diff --git a/src/scf.py b/src/scf.py new file mode 100644 index 0000000..b3d416a --- /dev/null +++ b/src/scf.py @@ -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)