Compute an SCF representation of an arbitrary density distribution#

Basis function expansions are a useful tool for computing gravitational potentials and forces from an arbitrary density function that may not have an analytic solution to Poisson’s equation. They are also useful for generating smoothed or compressed representations of gravitational potentials from discrete particle distributions. For astronomical density distributions, a useful expansion technique is the Self-Consistent Field (SCF) method, as initially developed by Hernquist & Ostriker (1992). In this method, using the notation of Lowing et al. 2011, the density and potential functions are expressed as:

\[\begin{split}\rho(r, \phi, \theta) = \sum_{l=0}^{l_{\rm max}} \sum_{m=0}^{l} \sum_{n=0}^{n_{\rm max}} Y_{lm}(\theta) \, \rho_{nl}(r) \, \left[S_{nlm}\,\cos(m\phi) + T_{nlm}\,\sin(m\phi) \right] \\ \Phi(r, \phi, \theta) = \sum_{l=0}^{l_{\rm max}} \sum_{m=0}^{l} \sum_{n=0}^{n_{\rm max}} Y_{lm}(\theta) \, \Phi_{nl}(r) \, \left[S_{nlm}\,\cos(m\phi) + T_{nlm}\,\sin(m\phi) \right]\end{split}\]

where \(Y_{lm}(\theta)\) are the usual spherical harmonics, \(\rho_{nlm}(r)\) and \(\Phi_{nlm}(r)\) are bi-orthogonal radial basis functions, and \(S_{nlm}\) and \(T_{nlm}\) are expansion coefficients, which need to be computed from a given density function. In this notebook, we’ll estimate low-order expansion coefficients for an analytic density distribution (written as a Python function).

[3]:
# Some imports we'll need later:

# Third-party
import matplotlib.pyplot as plt
import numpy as np

# Gala
import gala.dynamics as gd
import gala.potential as gp
from gala.potential.scf import compute_coeffs

SCF representation of an analytic density distribution#

Custom spherical density function#

For this example, we’ll assume that we want a potential representation of the spherical density function:

\[\rho(r) = \frac{1}{r^{1.8} \, (1 + r)^{2.7}}\]

Let’s start by writing a density function that takes a single set of Cartesian coordinates (x, y, z) and returns the (scalar) value of the density at that location:

[4]:
def density_func(x, y, z):
    r = np.sqrt(x**2 + y**2 + z**2)
    return 1 / (r**1.8 * (1 + r) ** 2.7)

Let’s visualize this density function. For comparison, let’s also over-plot the Hernquist density distribution. The SCF expansion uses the Hernquist density for radial basis functions, so the similarity of the density we want to represent and the Hernquist function gives us a sense of how many radial terms we will need in the expansion:

[5]:
hern = gp.HernquistPotential(m=1, c=1)
[6]:
x = np.logspace(-1, 1, 128)
plt.plot(x, density_func(x, 0, 0), marker="", label="custom density")

# need a 3D grid for the potentials in Gala
xyz = np.zeros((3, len(x)))
xyz[0] = x
plt.plot(x, hern.density(xyz), marker="", label="Hernquist")

plt.xscale("log")
plt.yscale("log")

plt.xlabel("$r$")
plt.ylabel(r"$\rho(r)$")

plt.legend(loc="best")
[6]:
<matplotlib.legend.Legend at 0x7fd2989998a0>
../_images/tutorials_Arbitrary-density-SCF_8_1.png

These functions are not too different, implying that we probably don’t need too many radial expansion terms in order to well represent the density/potential from this custom function. As an arbitrary number, let’s choose to compute radial terms up to and including \(n = 10\). In this case, because the density we want to represent is spherical, we don’t need any \(l, m\) terms, so we set lmax=0. We can also neglect the sin() terms of the expansion (\(T_{nlm}\)):

[7]:
(S, Serr), _ = compute_coeffs(
    density_func, nmax=10, lmax=0, M=1.0, r_s=1.0, S_only=True
)

The above variable S will contain the expansion coefficients, and the variable Serr will contain an estimate of the error in this coefficient value. Let’s now construct an SCFPotential object with the coefficients we just computed:

[8]:
S
[8]:
array([[[ 1.10328037e+01]],

       [[-2.15355027e+00]],

       [[ 5.89670966e-01]],

       [[-2.51335957e-01]],

       [[ 1.23222803e-01]],

       [[-6.99340663e-02]],

       [[ 4.25020670e-02]],

       [[-2.78007144e-02]],

       [[ 1.89401928e-02]],

       [[-1.34871472e-02]],

       [[ 9.86744880e-03]]])
[9]:
pot = gp.SCFPotential(m=1.0, r_s=1, Snlm=S, Tnlm=np.zeros_like(S))

Now let’s visualize the SCF estimated density with the true density:

[10]:
x = np.logspace(-1, 1, 128)
plt.plot(x, density_func(x, 0, 0), marker="", label="custom density")

# need a 3D grid for the potentials in Gala
xyz = np.zeros((3, len(x)))
xyz[0] = x
plt.plot(x, pot.density(xyz), marker="", label="SCF density")

plt.xscale("log")
plt.yscale("log")

plt.xlabel("$r$")
plt.ylabel(r"$\rho(r)$")

plt.legend(loc="best")
[10]:
<matplotlib.legend.Legend at 0x7fd260130340>
../_images/tutorials_Arbitrary-density-SCF_15_1.png

This does a pretty good job of capturing the radial fall-off of our custom density function, but you may want to iterate a bit to satisfy your own constraints. For example, you may want the density to be represented with a less than 1% deviation over some range of radii, or whatever.

As a second example, let’s now try a custom axisymmetric density distribution:

Custom axisymmetric density function#

For this example, we’ll assume that we want a potential representation of the flattened Hernquist density function:

\[\begin{split}\rho(R, z) = \frac{1}{r \, (1 + r)^{3}}\\ r^2 = R^2 + \frac{z^2}{q^2}\end{split}\]

where \(q\) is the flattening, which we’ll set to \(q=0.6\).

Let’s again start by writing a density function that takes a single set of Cartesian coordinates (x, y, z) and returns the (scalar) value of the density at that location:

[11]:
def density_func_flat(x, y, z, q):
    r = np.sqrt(x**2 + y**2 + (z / q) ** 2)
    return 1 / (r * (1 + r) ** 3) / (2 * np.pi)

Let’s compute the density along a diagonal line for a few different flattenings and again compare to the non-flattened Hernquist profile:

[12]:
x = np.logspace(-1, 1, 128)
xyz = np.zeros((3, len(x)))
xyz[0] = x
xyz[2] = x

for q in np.arange(0.6, 1 + 1e-3, 0.2):
    plt.plot(
        x,
        density_func_flat(xyz[0], 0.0, xyz[2], q),
        marker="",
        label=f"custom density: q={q}",
    )

plt.plot(x, hern.density(xyz), marker="", ls="--", label="Hernquist")

plt.xscale("log")
plt.yscale("log")

plt.xlabel("$r$")
plt.ylabel(r"$\rho(r)$")

plt.legend(loc="best")
[12]:
<matplotlib.legend.Legend at 0x7fd26051a1a0>
../_images/tutorials_Arbitrary-density-SCF_20_1.png

Because this is an axisymmetric density distribution, we need to also compute \(l\) terms in the expansion, so we set lmax=6, but we can skip the \(m\) terms using skip_m=True. Because this computes more coefficients, we might want to see the progress in real time - if you install the Python package tqdm and pass progress=True, it will also display a progress bar:

[13]:
q = 0.6
(S_flat, Serr_flat), _ = compute_coeffs(
    density_func_flat,
    nmax=4,
    lmax=6,
    args=(q,),
    M=1.0,
    r_s=1.0,
    S_only=True,
    skip_m=True,
    progress=True,
)
100%|██████████| 35/35 [00:02<00:00, 16.94it/s]
[14]:
pot_flat = gp.SCFPotential(m=1.0, r_s=1, Snlm=S_flat, Tnlm=np.zeros_like(S_flat))
[15]:
x = np.logspace(-1, 1, 128)
xyz = np.zeros((3, len(x)))
xyz[0] = x
xyz[2] = x

plt.plot(
    x,
    density_func_flat(xyz[0], xyz[1], xyz[2], q),
    marker="",
    label=f"true density q={q}",
)

plt.plot(x, pot_flat.density(xyz), marker="", ls="--", label="SCF density")

plt.xscale("log")
plt.yscale("log")

plt.xlabel("$r$")
plt.ylabel(r"$\rho(r)$")

plt.legend(loc="best")
[15]:
<matplotlib.legend.Legend at 0x7fd25ffaf130>
../_images/tutorials_Arbitrary-density-SCF_24_1.png

The SCF potential object acts like any other gala.potential object, meaning we can, e.g., plot density or potential contours:

[16]:
grid = np.linspace(-8, 8, 128)

fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharex=True, sharey=True)
_ = pot_flat.plot_contours((grid, grid, 0), ax=axes[0])
axes[0].set_xlabel("$x$")
axes[0].set_ylabel("$y$")

_ = pot_flat.plot_contours((grid, 0, grid), ax=axes[1])
axes[1].set_xlabel("$x$")
axes[1].set_ylabel("$z$")

for ax in axes:
    ax.set_aspect("equal")
../_images/tutorials_Arbitrary-density-SCF_26_0.png

And numerically integrate orbits by passing in initial conditions and integration parameters:

[17]:
w0 = gd.PhaseSpacePosition(pos=[3.5, 0, 1], vel=[0, 0.4, 0.05])
[18]:
orbit_flat = pot_flat.integrate_orbit(w0, dt=1.0, n_steps=5000)
_ = orbit_flat.plot()
../_images/tutorials_Arbitrary-density-SCF_29_0.png