Add pytests, numba optimizations, and switching between ne2001p/ne2025#3
Open
telegraphic wants to merge 28 commits intostella-ocker:mainfrom
Open
Add pytests, numba optimizations, and switching between ne2001p/ne2025#3telegraphic wants to merge 28 commits intostella-ocker:mainfrom
telegraphic wants to merge 28 commits intostella-ocker:mainfrom
Conversation
- Create tests/test_ne2025_main.py: Tests for NE2025 help/explain functionality - Create tests/test_ne2001_main.py: Tests for NE2001 help/explain functionality - Create tests/test_ne_input_main.py: Tests for read_nemod_parameters() function - Create tests/test_iss_mw_utils_main.py: Tests for scattering calculations - Existing tests/test_dmdsm.py: Tests for dmdsm_dm2d() with plotting enabled - Configure pytest with output directory fixture for plot files - Add .coveragerc for code coverage configuration - Add pytest-cov to environment.yaml dependencies - All 16 tests pass with code coverage tracking
sech²(z) = 4·exp(-2|z|) / (1 + exp(-2|z|))² This avoids np.cosh overflow warnings for large |z|
Applied Numba @njit to hottest functions in LoS integration: - nevoidN: Inner loop over void array (~589 calls per dmdsm_dm2d) - ne_outer/ne_inner: Disk component calculations (~40 calls per pass) JIT compilation automatically handles: - Loop unrolling and LLVM optimization - Elimination of Python function call overhead - Native machine code generation for math operations Performance gains (cumulative from session start): - Initial baseline: ~0.052s per dmdsm_dm2d call - After cumulative_trapezoid: ~0.035s (1.49× speedup) - After preallocation+cache: ~0.028s (1.86× speedup) - After Numba JIT (current): ~0.011s (4.65× total speedup) Speedup breakdown: - nevoidN: ~2.0-2.5× (inner void loop now compiled) - ne_outer/inner: ~1.5-2.0× (math operations optimized) - Overall: ~1.5-2.0× (amortized after first JIT compilation) First-call overhead: ~0.40s (JIT compFirst-call overhead: ~0.40s (JIT compFirst-call overhead: ~0.40s (JIT compFirst-call overhead: ~0.40s (JIT compFirst-→First-call overhead: ~0.40s (JIT compFirst-call overhead: ~baFirst-call overhead: ~0.40s (JIT compFirst-call overhead: ~0.40ba version: 0.64.0, llvmlite 0.46.0 - All regression tests passing (8/8 sm- All regression tests passing (8/8 sss- All regression tests passing (8/8 snc- All regression tests passing (8/8 sm- All rct- All regression tests passing (8/8 sm- All regressed - All regression tests passing (8/8 sm- All regre2)- All regression tests passing (8/8 sm- All regression tl,- All regression tests passing (8/8 sm- All regressioon- All regression tests passing (8/8 sm- All regressioests passing
…p achieved Added comprehensive documentation of two-phase optimization: PHASE 1: Algorithmic improvements - cumulative_trapezoid: O(N²) → O(N) integration (1.49× speedup) - Array preallocation: Removed append-in-loop overhead (1.86× cumulative) - Parameter caching: Reduced dict lookups in hot path PHASE 2: Numba JIT compilation - nevoidN: Inner void loop compiled (589 calls) - ne_outer/ne_inner: Density functions compiled (40 calls) - Graceful fallback if numba unavailable - First-call JIT overhead: ~0.40s (amortized over typical runs) TOTAL RESULT: 4.7× speedup - Baseline: 0.052s per dmdsm_dm2d call - Current: 0.011s per call (78.7% faster) - Typical use (1+100 calls): 14.9ms average TESTING STATUS: ✅ test_ne2001_main.py: 3/3 passing ✅ smooth_components_regression.py✅ smooth_components_regression.py✅ smooth_components_regression.py✅ smooth_components_regression.py✅ smooth_components_regression.py✅ smooth_components_regression.py✅ smooth_components_regression.py✅ smooth_components_regressiSpline construction overhead in spiral arm lookups - Future: JIT ne_lism functions (needs careful v- Future: JIT ne_lism functions (needs careful v- Future: JIT ne_litput checks.
…ween ne2001 and ne2025
telegraphic
changed the title
telegraphic
changed the title
…array (elimates numpy.diff calls)
Author
|
Update: code now runs significantly faster, mainly due to use of dm_cumulate_vec = \
pc_in_kpc * array([trapz(ne[:j], sf_vec[:j]) for j in range(1, Ns_fine+1) ])is much slower than dm_cumulate_vec = pc_in_kpc * cumulative_trapezoid(ne, sf_vec, initial=0.0)which is also neater. Vectorizing When integrating out to 50 kpc, it's now running about 115x faster (3231 ms -> 28 ms). |
Author
|
With these speed-ups, I was able to generate some all-sky maps of DM out to different distances on my laptop in a few hours. Here's a pretty picture of the NE2025 estimate for DM out to the galactic center:
|
Author
|
And example comparison between models: 
|
Owner
|
Hey @telegraphic the updated pygedm interface looks awesome. Really good to know about the cumulative_trapezoid speedup, sounds worth including in mwprop. Given the magnitude of this PR (and some of the specificity to your update of pygedm) I won't merge the whole thing, but I'll make sure to note which future updates are drawn from this fork. |
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
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.
Hi @stella-ocker, I'd like to update pygedm to provide access to NE2025, including on the web app at https://apps.datacentral.org.au/pygedm/. FYI the work to update pygedm is happening here: FRBs/pygedm#30
As a first port of call, I added references to the NE2025 paper to the pygedm README.md. (I note users have been great at citing the underlying models when publishing work using pygedm).
I wanted to generate some all-sky DM maps for the web app, but the 45x slowdown made it unfeasible. So I did some optimizations, adding numba in hotspots, and got a 4-5x speedup.
Before I started the optimizations, I wanted to make sure I didn't break anything, so added some unit tests and regression tests. I copied the test code in the
__main__functions as the baseline.I also wanted to enable switching between NE2001p and NE2025 without reimporting mwprop. I've added a
set_model()function to do this, which required some rejigging (using importlib)I realise this is a huge PR! So I won't be offended if it doesn't suit your plans and you don't wish to merge it.