From 565de7e38bca10ec10ae1e233baba9ed4246e09a Mon Sep 17 00:00:00 2001 From: Cara Giovanetti Date: Fri, 23 Oct 2026 14:31:19 -0700 Subject: [PATCH 1/4] dense output --- linx/background.py | 56 +++++++++++++++++++--------------------------- 1 file changed, 23 insertions(+), 33 deletions(-) diff --git a/linx/background.py b/linx/background.py index 8caaafa..ff3397e 100644 --- a/linx/background.py +++ b/linx/background.py @@ -31,6 +31,11 @@ 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 4096. throw : bool, optional Whether to raise exceptions on solver failure. Default is `True`. Set to `False` for parameter scans where some combinations may fail. @@ -42,9 +47,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=4096): """ Initialize the BackgroundModel with thermodynamic options. @@ -62,6 +68,11 @@ 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 4096. throw : bool, optional If True, raise exceptions on solver failure. Default is True. Set to False for parameter scans where some combinations may fail. @@ -73,6 +84,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 @@ -98,10 +110,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. @@ -150,7 +158,7 @@ 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 ), @@ -158,41 +166,23 @@ def T_EM_check(t, y, args, **kwargs): 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 From 0a902460db540b5501420dd299c34e2d4a9dbb20 Mon Sep 17 00:00:00 2001 From: Cara Giovanetti Date: Fri, 23 Oct 2026 15:07:30 -0700 Subject: [PATCH 2/4] reduce n_out default --- linx/background.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/linx/background.py b/linx/background.py index ff3397e..7a43e79 100644 --- a/linx/background.py +++ b/linx/background.py @@ -35,7 +35,8 @@ class BackgroundModel(eqx.Module): 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 4096. + 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. @@ -50,7 +51,7 @@ class BackgroundModel(eqx.Module): 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, n_out=4096): + 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. @@ -72,7 +73,8 @@ def __init__(self, decoupled=False, use_FD=True, collision_me=True, LO=True, NLO 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 4096. + 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. From c180a82ed2ba0e8c3940f6f66741bc8746a03f2b Mon Sep 17 00:00:00 2001 From: Cara Giovanetti Date: Fri, 23 Oct 2026 15:08:21 -0700 Subject: [PATCH 3/4] log interpolation --- linx/abundances.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/linx/abundances.py b/linx/abundances.py index 0bd2436..0af6c45 100644 --- a/linx/abundances.py +++ b/linx/abundances.py @@ -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( @@ -506,9 +506,9 @@ 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] @@ -516,13 +516,12 @@ def Y_prime(self, t, Y, args): 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 @@ -531,7 +530,7 @@ def Y_prime(self, t, Y, args): 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] + 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( From 1300e516656b11f048a548c981ab898d85cdc89e Mon Sep 17 00:00:00 2001 From: Cara Giovanetti Date: Fri, 23 Oct 2026 15:33:16 -0700 Subject: [PATCH 4/4] exp T fix --- linx/abundances.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/linx/abundances.py b/linx/abundances.py index 0af6c45..dd93391 100644 --- a/linx/abundances.py +++ b/linx/abundances.py @@ -529,9 +529,9 @@ def Y_prime(self, t, Y, args): # number density times amu. rhoBBN = rho0BmaOvermB * const.MeV4_to_gcmm3/a**3 - T_t = jnp.interp( + 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,