Skip to content

Refactor: Lazily load or eliminate heavy dependencies (scipy, numba, h5py, astropy, cosmopy) to fix import overhead#140

Merged
kayhangultekin merged 4 commits into
devfrom
dev-kg-functioncaching
Jun 2, 2026
Merged

Refactor: Lazily load or eliminate heavy dependencies (scipy, numba, h5py, astropy, cosmopy) to fix import overhead#140
kayhangultekin merged 4 commits into
devfrom
dev-kg-functioncaching

Conversation

@kayhangultekin
Copy link
Copy Markdown
Collaborator

Description

This PR resolves the critical package initialization bottleneck by completely decoupling heavy scientific packages (scipy, numba, h5py, astropy, and cosmopy) from the global top-level import sequence.

Key Changes

  • Lazy-Loaded Cosmopy Integration: Deferred the initialization of cosmopy by encapsulating it in a lazy caching function. This completely eliminates the significant overhead of setting up the default cosmology framework on package boot.
  • Eliminated Astropy from Core Constants: Rewrote holodeck/constants.py to use hardcoded physical constant values rather than pulling from astropy.constants. Because constants.py is loaded directly in __init__.py, this completely immunizes the base package initialization from the astropy constants registry overhead.
  • Hardcoded 1-Sigma Constant: Replaced a dynamic scipy.stats.norm.cdf(1) call with its exact float64 constant (0.8413447460685429), eliminating a global scipy load in utils.py.
  • Localized SciPy & h5py: Shifted scipy.interpolate (for PchipInterpolator), scipy.special (for jn), and h5py imports directly into the specific function scopes where they are executed.
  • Lazy Numba JIT Decorator: Implemented a lightweight proxy decorator (_LazyNJIT) around _gw_ecc_func. This completely bypasses loading the heavy LLVM compiler machinery at boot time, delaying the import until the function is actually called during execution.
  • LaTeX Syntax Fix: Prepend r to the docstring on line 1073 to resolve a Python 3.12 SyntaxWarning regarding an invalid escape sequence (\l).

Performance Impact

Using time python -c "import holodeck", package startup overhead dropped by over 95% on subsequent runs (2600% if Trump is doing the math), dropping total wall-clock time from ~2.6 seconds down to around 0.1 seconds.

Side-by-Side Comparison:

⟩ git switch dev
Switched to branch 'dev'
Your branch is up to date with 'origin/dev'.
⟩ for i in $(seq 10 ) ; do time python -c "import holodeck" ; done
/Users/kayhan/Documents/projects/nanograv/holodeck/holodeck/utils.py:1058: SyntaxWarning: invalid escape sequence '\l'
  You may need to convert, e.g. $dN/dx = \frac{dN}{d \ln x} \frac{1}{x}$.
python -c "import holodeck"  1.37s user 0.31s system 52% cpu 3.186 total
python -c "import holodeck"  1.32s user 0.23s system 59% cpu 2.594 total
python -c "import holodeck"  1.30s user 0.24s system 59% cpu 2.577 total
python -c "import holodeck"  1.31s user 0.24s system 60% cpu 2.582 total
python -c "import holodeck"  1.31s user 0.24s system 59% cpu 2.584 total
python -c "import holodeck"  1.30s user 0.24s system 59% cpu 2.580 total
python -c "import holodeck"  1.31s user 0.24s system 59% cpu 2.576 total
python -c "import holodeck"  1.33s user 0.24s system 60% cpu 2.608 total
python -c "import holodeck"  1.33s user 0.24s system 60% cpu 2.618 total
python -c "import holodeck"  1.31s user 0.24s system 60% cpu 2.577 total

⟩ git switch dev-kg-functioncaching
Switched to branch 'dev-kg-functioncaching'
Your branch is up to date with 'origin/dev-kg-functioncaching'.
⟩ for i in $(seq 10 ) ; do time python -c "import holodeck" ; done
python -c "import holodeck"  0.07s user 0.03s system 57% cpu 0.181 total
python -c "import holodeck"  0.07s user 0.02s system 95% cpu 0.097 total
python -c "import holodeck"  0.07s user 0.02s system 92% cpu 0.096 total
python -c "import holodeck"  0.07s user 0.02s system 94% cpu 0.093 total
python -c "import holodeck"  0.07s user 0.02s system 89% cpu 0.100 total
python -c "import holodeck"  0.07s user 0.02s system 90% cpu 0.101 total
python -c "import holodeck"  0.07s user 0.02s system 91% cpu 0.108 total
python -c "import holodeck"  0.08s user 0.03s system 86% cpu 0.117 total
python -c "import holodeck"  0.07s user 0.02s system 91% cpu 0.101 total
python -c "import holodeck"  0.07s user 0.02s system 87% cpu 0.106 total

Kayhan Gultekin added 3 commits May 29, 2026 16:20
…ing. Made cosmopy lazily loaded to improve import time loading. Changed constatns to be hard coded to improve import loading time.
…ss it is being run in MPI based on environment variable check or if mpi4py has been imported elsewhere. Updated holodeck/utils.py so that it only imports scipy in the functions declared within utils that actually use scipy. Same for h5py. Then made a lazy loader for numba so that it is only imported when needed. This avoids a penalty for importing holodeck, which automatically imports utils. I need to run a few tests to make sure this has not broken anything.
@kayhangultekin kayhangultekin mentioned this pull request May 30, 2026
4 tasks
Copy link
Copy Markdown
Collaborator

@astrolamb astrolamb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey @kayhangultekin this looks great to me. I'm also seeing a similar speed-up on import. There's just some clean-up required for this PR.

I think we can also drop the requirement for python 3.9, which is the only version currently failing tests

Comment thread holodeck/constants.py Outdated
Comment thread holodeck/constants.py Outdated
Comment thread holodeck/logger.py Outdated
@astrolamb
Copy link
Copy Markdown
Collaborator

@kayhangultekin it looks like we have an assertion issue with the ubuntu tests... macos is ok though

@kayhangultekin kayhangultekin merged commit 3cd9dec into dev Jun 2, 2026
5 of 9 checks passed
@kayhangultekin kayhangultekin deleted the dev-kg-functioncaching branch June 2, 2026 14:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants