# Source code for biff.scf.core

# coding: utf-8

# Third-party
from astropy import log as logger
import numpy as np
import scipy.integrate as si

# Project
from ._computecoeff import Snlm_integrand, Tnlm_integrand, STnlm_discrete, STnlm_var_discrete

__all__ = ['compute_coeffs', 'compute_coeffs_discrete']

[docs]def compute_coeffs(density_func, nmax, lmax, M, r_s, args=(),
skip_odd=False, skip_even=False, skip_m=False,
"""
Compute the expansion coefficients for representing the input density function using a basis
function expansion.

Computing the coefficients involves computing triple integrals which are computationally
expensive. For an example of how to parallelize the computation of the coefficients, see
examples/parallel_compute_Anlm.py.

Parameters
----------
density_func : function, callable
A function or callable object that evaluates the density at a given position. The call
format must be of the form: density_func(x, y, z, M, r_s, args) where x,y,z are
cartesian coordinates, M is a scale mass, r_s a scale radius, and args is an
iterable containing any other arguments needed by the density function.
nmax : int
Maximum value of n for the radial expansion.
lmax : int
Maximum value of l for the spherical harmonics.
M : numeric
Scale mass.
r_s : numeric
args : iterable (optional)
A list or iterable of any other arguments needed by the density
function.
skip_odd : bool (optional)
Skip the odd terms in the angular portion of the expansion. For example, only
take :math:l=0,2,4,...
skip_even : bool (optional)
Skip the even terms in the angular portion of the expansion. For example, only
take :math:l=1,3,5,...
skip_m : bool (optional)
Ignore terms with :math:m > 0.
S_only : bool (optional)
Only compute the S coefficients.
Any additional keyword arguments are passed through to
~scipy.integrate.nquad as options, opts.

Returns
-------
Snlm : float, ~numpy.ndarray
The value of the cosine expansion coefficient.
Snlm_err : , ~numpy.ndarray
An estimate of the uncertainty in the coefficient value (from ~scipy.integrate.nquad).
Tnlm : , ~numpy.ndarray
The value of the sine expansion coefficient.
Tnlm_err : , ~numpy.ndarray
An estimate of the uncertainty in the coefficient value. (from ~scipy.integrate.nquad).

"""
lmin = 0
lstride = 1

if skip_odd or skip_even:
lstride = 2

if skip_even:
lmin = 1

Snlm = np.zeros((nmax+1, lmax+1, lmax+1))
Snlm_e = np.zeros((nmax+1, lmax+1, lmax+1))
Tnlm = np.zeros((nmax+1, lmax+1, lmax+1))
Tnlm_e = np.zeros((nmax+1, lmax+1, lmax+1))

limits = [[0,2*np.pi], # phi
[-1,1.], # X (cos(theta))
[-1,1.]] # xsi

for n in range(nmax+1):
for l in range(lmin, lmax+1, lstride):
for m in range(l+1):
if skip_m and m > 0: continue

logger.debug("Computing coefficients (n,l,m)=({},{},{})"
.format(n,l,m))

Snlm_integrand, ranges=limits,
args=(density_func, n, l, m, M, r_s, args),

if not S_only:
Tnlm_integrand, ranges=limits,
args=(density_func, n, l, m, M, r_s, args),

return (Snlm,Snlm_e), (Tnlm,Tnlm_e)

[docs]def compute_coeffs_discrete(xyz, mass, nmax, lmax, r_s,
skip_odd=False, skip_even=False, skip_m=False,
compute_var=False):
"""
Compute the expansion coefficients for representing the density distribution of input points
as a basis function expansion. The points, xyz, are assumed to be samples from the
density distribution.

Computing the coefficients involves computing triple integrals which are computationally
expensive. For an example of how to parallelize the computation of the coefficients, see
examples/parallel_compute_Anlm.py.

Parameters
----------
xyz : array_like
Samples from the density distribution. Should have shape (n_samples,3).
mass : array_like
Mass of each sample. Should have shape (n_samples,).
nmax : int
Maximum value of n for the radial expansion.
lmax : int
Maximum value of l for the spherical harmonics.
r_s : numeric
skip_odd : bool (optional)
Skip the odd terms in the angular portion of the expansion. For example, only
take :math:l=0,2,4,...
skip_even : bool (optional)
Skip the even terms in the angular portion of the expansion. For example, only
take :math:l=1,3,5,...
skip_m : bool (optional)
Ignore terms with :math:m > 0.
compute_var : bool (optional)
Also compute the variances of the coefficients. This does not compute the full covariance
matrix of the coefficients, just the individual variances.
TODO: separate function to compute full covariance matrix?

Returns
-------
Snlm : float
The value of the cosine expansion coefficient.
Tnlm : float
The value of the sine expansion coefficient.

"""
lmin = 0
lstride = 1

if skip_odd or skip_even:
lstride = 2

if skip_even:
lmin = 1

Snlm = np.zeros((nmax+1, lmax+1, lmax+1))
Tnlm = np.zeros((nmax+1, lmax+1, lmax+1))

if compute_var:
Snlm_var = np.zeros((nmax+1, lmax+1, lmax+1))
Tnlm_var = np.zeros((nmax+1, lmax+1, lmax+1))

# positions and masses of point masses
xyz = np.ascontiguousarray(np.atleast_2d(xyz))
mass = np.ascontiguousarray(np.atleast_1d(mass))

r = np.sqrt(np.sum(xyz**2, axis=-1))
s = r / r_s
phi = np.arctan2(xyz[:,1], xyz[:,0])
X = xyz[:,2] / r

for n in range(nmax+1):
for l in range(lmin, lmax+1, lstride):
for m in range(l+1):
if skip_m and m > 0: continue

logger.debug("Computing coefficients (n,l,m)=({},{},{})".format(n,l,m))

Snlm[n,l,m], Tnlm[n,l,m] = STnlm_discrete(s, phi, X, mass, n, l, m)
if compute_var:
Snlm_var[n,l,m], Tnlm_var[n,l,m] = STnlm_var_discrete(s, phi, X, mass, n, l, m)

if compute_var:
return (Snlm,Snlm_var), (Tnlm,Tnlm_var)
else:
return Snlm, Tnlm