import numpy as np
from pbhgr.background import evolve, conservation_residual
def test_scaling_no_decay():
res = evolve(rho_X0=1.0, rho_R0=1.0, Gamma=0.0, t_end=50.0, t0=0.1)
a, rX, rR = res["a"], res["rho_X"], res["rho_R"]
assert np.allclose(rX * a**3, rX[0], rtol=1e-7)
assert np.allclose(rR * a**4, rR[0], rtol=1e-7)
def test_radiation_dominated_power_law():
# start radiation dominated at Hubble time t0 = 1/(2H0)
rR0 = 1.0
H0 = np.sqrt(8 * np.pi / 3 * rR0)
t0 = 1 / (2 * H0)
res = evolve(rho_X0=0.0, rho_R0=rR0, Gamma=0.0, t_end=100 * t0, t0=t0)
assert np.allclose(res["a"], (res["t"] / t0) ** 0.5, rtol=1e-6)
def test_decay_transfers_energy_and_conserves_total():
G = 0.3
res = evolve(rho_X0=1.0, rho_R0=1e-3, Gamma=G, t_end=30.0, t0=0.1)
a, rX, rR, t = res["a"], res["rho_X"], res["rho_R"], res["t"]
# comoving X density decays as exp(-Gamma t)
assert np.allclose(rX * a**3, rX[0] * np.exp(-G * (t - t[0])), rtol=1e-6)
# integral form of total covariant conservation: (rho a^3)(t) - (rho a^3)(0) = -int H rho_R a^3 dt
res = evolve(rho_X0=1.0, rho_R0=1e-3, Gamma=G, t_end=30.0, t0=0.1, n_out=20000)
a, rX, rR, t, H = res["a"], res["rho_X"], res["rho_R"], res["t"], res["H"]
tot = (rX + rR) * a**3
work = np.concatenate([[0.0], np.cumsum(0.5 * (H * rR * a**3)[1:] + 0.5 * (H * rR * a**3)[:-1]) * np.diff(t)])
assert np.max(np.abs(tot - tot[0] + work)) < 1e-5 * tot[0]
# at late times the energy is in radiation
assert rR[-1] > 10 * rX[-1]