Skip to content

Smooth gradients #45

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
cgiovanetti merged 4 commits into from
Oct 23, 2026
Merged

Smooth gradients #45

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
565de7e
dense output
cgiovanetti Oct 23, 2026
0a90246
reduce n_out default
cgiovanetti Oct 23, 2026
c180a82
log interpolation
cgiovanetti Oct 23, 2026
1300e51
exp T fix
cgiovanetti Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 13 additions & 14 deletions linx/abundances.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -316,8 +316,8 @@ def __call__(
ODETerm(self.Y_prime), solver,
t0=t_start, t1=t_end, dt0=None, y0=Y_i,
args=(
a_vec, t_vec, T_g_vec, T_interval_nTOp, nTOp_frwrd,
nTOp_bkwrd, eta_fac, tau_n_fac, nuclear_rates_q
jnp.log(a_vec), jnp.log(t_vec), jnp.log(T_g_vec), T_interval_nTOp,
nTOp_frwrd, nTOp_bkwrd, eta_fac, tau_n_fac, nuclear_rates_q
),
saveat=saveat,
stepsize_controller=PIDController(
Expand Down Expand Up @@ -506,33 +506,32 @@ def Y_prime(self, t, Y, args):

"""

a_vec_in = args[0]
t_vec_in = args[1]
T_g_vec_in = args[2]
log_a_vec_in = args[0]
log_t_vec_in = args[1]
log_T_g_vec_in = args[2]
T_interval_in = args[3]
nTOp_frwrd_vec_in = args[4]
nTOp_bkwrd_vec_in = args[5]
eta_fac = args[6]
tau_n_fac = args[7]
nuclear_rates_q = args[8]

a_in = a_vec_in[0]
a_fin = a_vec_in[-1]
log_t = jnp.log(t)

a = jnp.interp(
t, t_vec_in, a_vec_in,
left=a_in,right=a_fin
)
a = jnp.exp(jnp.interp(
log_t, log_t_vec_in, log_a_vec_in,
left=log_a_vec_in[0], right=log_a_vec_in[-1]
))

# Baryon density rescaled by eta_fac.
n0B = eta_fac*const.n0CMB*const.eta0
rho0BmaOvermB = ma * n0B
# number density times amu.
rhoBBN = rho0BmaOvermB * const.MeV4_to_gcmm3/a**3

T_t = jnp.interp(
t, t_vec_in, T_g_vec_in, left=T_g_vec_in[0],right=T_g_vec_in[-1]
)
T_t = jnp.exp(jnp.interp(
log_t, log_t_vec_in, log_T_g_vec_in, left=log_T_g_vec_in[0],right=log_T_g_vec_in[-1]
))

dY = self.nuclear_net(
Y, T_t, rhoBBN, T_interval_in, nTOp_frwrd_vec_in,
Expand Down
58 changes: 25 additions & 33 deletions linx/background.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@ class BackgroundModel(eqx.Module):
Whether to use leading order QED correction. Default is `True`.
NLO : bool, optional
Whether to use next-to-leading order QED correction. Default is True.
max_steps : int, optional
Max number of steps for background solve. Default is 512.
n_out: int, optional
Number of points log-spaced on interpolation grid of thermodynamics solution.
Default is 512. It is recommended you raise this to 1024 or more if you need
LINX gradients.
throw : bool, optional
Whether to raise exceptions on solver failure. Default is `True`.
Set to `False` for parameter scans where some combinations may fail.
Expand All @@ -42,9 +48,10 @@ class BackgroundModel(eqx.Module):
LO : bool
NLO : bool
max_steps : int
n_out : int
throw : bool

def __init__(self, decoupled=False, use_FD=True, collision_me=True, LO=True, NLO = True, throw=True, max_steps=512):
def __init__(self, decoupled=False, use_FD=True, collision_me=True, LO=True, NLO = True, throw=True, max_steps=512, n_out=512):

"""
Initialize the BackgroundModel with thermodynamic options.
Expand All @@ -62,6 +69,12 @@ def __init__(self, decoupled=False, use_FD=True, collision_me=True, LO=True, NLO
If True, include leading order QED corrections. Default is True.
NLO : bool, optional
If True, include next-to-leading order QED corrections. Default is True.
max_steps : int, optional
Max number of steps for background solve. Default is 512.
n_out: int, optional
Number of log-spaced points on interpolation grid of thermodynamics solution.
Default is 512. It is recommended you raise this to 1024 or more if you need
LINX gradients.
throw : bool, optional
If True, raise exceptions on solver failure. Default is True.
Set to False for parameter scans where some combinations may fail.
Expand All @@ -73,6 +86,7 @@ def __init__(self, decoupled=False, use_FD=True, collision_me=True, LO=True, NLO
self.LO = LO
self.NLO = NLO
self.max_steps = max_steps
self.n_out = n_out
self.throw = throw

@eqx.filter_jit
Expand All @@ -98,10 +112,6 @@ def __call__(
Relative tolerance of the abundance solver. Default is `1e-8`.
atol : float, optional
Absolute tolerance of the abundance solver. Default is `1e-10`.
max_steps : int, optional
Maximum number of steps taken by the solver. Default is `4096`.
Increasing this slows down the code, while decreasing this could
mean that the solver cannot complete the solution.
solver : Diffrax ODE solver
The Diffrax ODE solver to use. A stiff solver is recommended.
Default is the Tsitouras' 5/4 solver.
Expand Down Expand Up @@ -150,49 +160,31 @@ def T_EM_check(t, y, args, **kwargs):
sol = diffeqsolve(
ODETerm(self.dY), solver, args=(lna_init, rho_extra_init, me),
t0 = t0, t1=jnp.inf, dt0=None, y0=Y0,
saveat=SaveAt(steps=True), event=Event(T_EM_check),
saveat=SaveAt(t0=True, t1=True, dense=True), event=Event(T_EM_check),
stepsize_controller=PIDController(
rtol=rtol, atol=atol
),
max_steps=self.max_steps,
throw=self.throw
)

a_vec = jnp.exp(sol.ys[0])
rho_g_vec = rho_massless_BE_v(sol.ys[1], 0., 2)
rho_nu_vec = rho_massless_FD_v(sol.ys[2], 0., 2)

# These vectors always have max_steps entries so that jit and grad
# work but the solver stops before hitting max_steps. Find the last
# legitimate step made by the solver, when T_g drops below T_end.

last_step_ind = jnp.max(
jnp.argwhere(
sol.ys[1] < T_end,
size=self.max_steps
)[:,0]
)

t_vec = jnp.exp(jnp.linspace(
jnp.log(sol.ts[0]), jnp.log(sol.ts[-1]), self.n_out
))
t_vec = t_vec.at[0].set(sol.ts[0]).at[-1].set(sol.ts[-1])
lna_vec, T_EM_vec, T_nu_vec = jax.vmap(sol.evaluate)(t_vec)

# Set every step after this in all vectors to be identical,
# so that there is no effect on interpolation.
t_vec = jnp.where(sol.ts == jnp.inf, sol.ts[last_step_ind], sol.ts)
a_vec = jnp.where(a_vec == jnp.inf, a_vec[last_step_ind], a_vec)
rho_g_vec = jnp.where(
rho_g_vec == jnp.inf, rho_g_vec[last_step_ind], rho_g_vec
)
rho_nu_vec = jnp.where(
rho_nu_vec == jnp.inf, rho_nu_vec[last_step_ind], rho_nu_vec
)
a_vec = jnp.exp(lna_vec)
rho_g_vec = rho_massless_BE_v(T_EM_vec, 0., 2)
rho_nu_vec = rho_massless_FD_v(T_nu_vec, 0., 2)

# Trivial relation between rho_extra and a.
rho_extra_vec = rho_extra_init * (jnp.exp(lna_init)**4 / a_vec**4)

# Rescale a so that the present day CMB temperature is correct.
final_a = const.T0CMB / (
sol.ys[1][last_step_ind]
)

final_a = const.T0CMB / T_EM_vec[-1]
a_vec *= final_a / a_vec[-1]

# old version -- used first step of solve and not initial a
Expand Down
Loading