Python script which solves for the gravitational potential of a spherically symmetric star with a Lane-Emden polytrope density profile

lane-emden-potential.py

import numpy as np
import matplotlib.pyplot as plt

def density(r):
    if density.outside:
        return 0.0
    else:
        z = density.rho_c * np.sin(r/R) / (r/R) if r > 1e-8 else 0.0
        if z < 0.0:
            density.outside = True
            return 0.0
        return z

def f(y, r):
    psi, phi = y[0], y[1]
    psi_p = 4*np.pi*G*density(r) - (2 * psi/r if r > 1e-8 else 0.0)
    phi_p = psi
    return np.array([psi_p, phi_p])

G = 0.1  # Gravitational constant
R = 1.0  # Characteristic radius

psi = [ ]
phi = [ ]
rad = [ ]
rho = [ ]
dr = 1e-3
r = 0.0
y = np.array([0.0, 0.0])

density.outside = False
density.rho_c = 1.0

while r < 10 * R:
    k1 = f(y, r)
    k2 = f(y + 0.5*k1*dr, r + 0.5*dr)
    k3 = f(y + 0.5*k2*dr, r + 0.5*dr)
    k4 = f(y + 1.0*k3*dr, r + 1.0*dr)
    y += (k1 + 2*k2 + 2*k3 + k4) * dr / 6.0
    r += dr
    phi.append(y[1])
    rad.append(r)
    rho.append(density(r))

plt.plot(rad, phi, label=r"$\phi(r)$", ls='-.', lw=2)
plt.plot(rad, rho, label=r"$\rho(r)$", ls='--', lw=2)
plt.legend(loc='best')
plt.show()