Use interpax.interp1d to support non-monotonic temperature evolution (reheating) #16
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Replace jnp.interp with interpax.interp1d in abundances.py, nuclear.py, and weak_rates.py to support reheating scenarios where photon temperature may not be monotonically decreasing. jnp.interp requires monotonically increasing x-coordinates, which breaks when T_g_vec is non-monotonic (e.g., during reheating where temperature can increase before decreasing again). The fix sorts arrays by the x-coordinate before interpolation. Changes: - Add interpax dependency to requirements.txt - abundances.py: Sort by log(T_g) for a_start, t_start, t_end interpolations and by log(rho_tot) for P_tot interpolations in get_t/get_a - weak_rates.py: Sort by Tg for neutrino-to-photon temperature ratio lookup - nuclear.py: Sort by T_interval for weak rate interpolations Note: thermo.py and reactions.py are unchanged as they use pre-tabulated data that is guaranteed to be monotonic. Closes #15 Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
Owner
|
This looks great to me, thank you @smsharma ! I'm even seeing a ~25% speedup on M1 Mac (haven't tested other hardware yet)--with that in mind we should probably replace the jnp.interp's in thermo too, I see no reason not to. Passes the abundance test so I'm happy to merge pending @hongwanliu 's review. |
cgiovanetti
approved these changes
Jan 23, 2026
Owner
cgiovanetti
left a comment
cgiovanetti
left a comment
There was a problem hiding this comment.
Significant performance enhancement from interp1d over jnp.interp!
I'd hoped interp1d could take non-monotonic functions given their docs, but it seems like they just left that detail out; the sorting is a nice fix. Since we can still pass in T_g in any order we like this does not affect code output.
Collaborator
Author
|
Gonna merge in! |
smsharma
added a commit
that referenced
this pull request
Jan 24, 2026
Follow-up to PR #16 which replaced jnp.interp with interpax in abundances.py. This change applies the same improvement to thermo.py for consistency and better performance. Changes: - Import interpax module - Flip QED correction tables at load time (instead of at each call) for monotonically increasing x coordinates required by interpax - Replace 6 jnp.interp calls with interpax.interp1d: - rho_EM_std: 2 calls for QED corrections - p_EM_std: 1 call for QED correction - rho_plus_p_EM_std: 1 call for QED correction - G_nue_with_me: 1 call for collision factor interpolation - G_numt_with_me: 1 call for collision factor interpolation Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2 tasks
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary
jnp.interpwithinterpax.interp1dinabundances.py,nuclear.py, andweak_rates.pyto support reheating scenarios where photon temperature may not be monotonically decreasingthermo.pyandreactions.pyare unchanged as they use pre-tabulated monotonic dataProblem
jnp.interprequires monotonically increasing x-coordinates, which breaks whenT_g_vecis non-monotonic (e.g., during reheating where temperature can increase before decreasing again). The previous code usedjnp.flip()assuming temperature always decreases, which doesn't work for reheating.Solution
Use
interpax.interp1dwith sorted arrays:interpax.interp1dwithmethod='linear'andextrap=TrueTest plan
pytest/test_abundances.py)cc @cgiovanetti (as requested)
Closes #15
🤖 Generated with Claude Code