diff --git a/linx/abundances.py b/linx/abundances.py index 0bd2436..dd93391 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 @@ -530,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_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, diff --git a/linx/background.py b/linx/background.py index 8caaafa..7a43e79 100644 --- a/linx/background.py +++ b/linx/background.py @@ -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. @@ -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. @@ -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. @@ -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 @@ -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. @@ -150,7 +160,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 +168,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