"""
Functions for solving the TOV equations. Main source of inspiration: NMMA code https://github.com/nuclear-multimessenger-astronomy/nmma/blob/f03ea14fcd6cafc9b189c7a1dd088ddd5b7fe898/nmma/eos/tov.py
This script assumes geometric units.
"""
from . import utils
import jax.numpy as jnp
from diffrax import diffeqsolve, ODETerm, Dopri5, SaveAt, PIDController
[docs]
def tov_ode(h, y, eos):
# fetch the eos arrays
ps = eos["p"]
hs = eos["h"]
es = eos["e"]
dloge_dlogps = eos["dloge_dlogp"]
# actual equations
r, m, H, b = y
e = utils.interp_in_logspace(h, hs, es)
p = utils.interp_in_logspace(h, hs, ps)
dedp = e / p * jnp.interp(h, hs, dloge_dlogps)
A = 1.0 / (1.0 - 2.0 * m / r)
C1 = 2.0 / r + A * (2.0 * m / (r * r) + 4.0 * jnp.pi * r * (p - e))
C0 = A * (
-6 / (r * r)
+ 4.0 * jnp.pi * (e + p) * dedp
+ 4.0 * jnp.pi * (5.0 * e + 9.0 * p)
) - jnp.power(2.0 * (m + 4.0 * jnp.pi * r * r * r * p) / (r * (r - 2.0 * m)), 2.0)
drdh = -r * (r - 2.0 * m) / (m + 4.0 * jnp.pi * r * r * r * p)
dmdh = 4.0 * jnp.pi * r * r * e * drdh
dHdh = b * drdh
dbdh = -(C0 * H + C1 * b) * drdh
dydt = drdh, dmdh, dHdh, dbdh
return dydt
[docs]
def calc_k2(R, M, H, b):
y = R * b / H
C = M / R
num = (
(8.0 / 5.0)
* jnp.power(1 - 2 * C, 2.0)
* jnp.power(C, 5.0)
* (2 * C * (y - 1) - y + 2)
)
den = (
2
* C
* (
4 * (y + 1) * jnp.power(C, 4)
+ (6 * y - 4) * jnp.power(C, 3)
+ (26 - 22 * y) * C * C
+ 3 * (5 * y - 8) * C
- 3 * y
+ 6
)
)
den -= (
3
* jnp.power(1 - 2 * C, 2)
* (2 * C * (y - 1) - y + 2)
* jnp.log(1.0 / (1 - 2 * C))
)
return num / den
[docs]
def tov_solver(eos, pc):
# fetch the eos arrays
ps = eos["p"]
hs = eos["h"]
es = eos["e"]
dloge_dlogps = eos["dloge_dlogp"]
# central values
hc = utils.interp_in_logspace(pc, ps, hs)
ec = utils.interp_in_logspace(hc, hs, es)
dedp_c = ec / pc * jnp.interp(hc, hs, dloge_dlogps)
dhdp_c = 1.0 / (ec + pc)
dedh_c = dedp_c / dhdp_c
# initial values
dh = -1e-3 * hc
h0 = hc + dh
r0 = jnp.sqrt(3.0 * (-dh) / 2.0 / jnp.pi / (ec + 3.0 * pc))
r0 *= 1.0 - 0.25 * (ec - 3.0 * pc - 0.6 * dedh_c) * (-dh) / (ec + 3.0 * pc)
m0 = 4.0 * jnp.pi * ec * jnp.power(r0, 3.0) / 3.0
m0 *= 1.0 - 0.6 * dedh_c * (-dh) / ec
H0 = r0 * r0
b0 = 2.0 * r0
y0 = (r0, m0, H0, b0)
sol = diffeqsolve(
ODETerm(tov_ode),
Dopri5(scan_kind="bounded"),
t0=h0,
t1=0,
dt0=dh,
y0=y0,
args=eos,
saveat=SaveAt(t1=True),
stepsize_controller=PIDController(rtol=1e-5, atol=1e-6),
)
R = sol.ys[0][-1]
M = sol.ys[1][-1]
H = sol.ys[2][-1]
b = sol.ys[3][-1]
k2 = calc_k2(R, M, H, b)
return M, R, k2