Skip to content

Add Numba JIT acceleration for KS computation (~2-3× end-to-end speedup)#1

Open
TorbenOestergaard wants to merge 1 commit intomarkus-schaffer:mainfrom
TorbenOestergaard:numba-speedup
Open

Add Numba JIT acceleration for KS computation (~2-3× end-to-end speedup)#1
TorbenOestergaard wants to merge 1 commit intomarkus-schaffer:mainfrom
TorbenOestergaard:numba-speedup

Conversation

@TorbenOestergaard
Copy link

Description:

This PR adds optional Numba JIT compilation to the inner loops of SAtom, significantly reducing computation time for large datasets.

Changes

_core.py

  • Added three @‌njit(cache=True) functions:
    • _ks_row_mode01 — fused KS distance for SScompare 0 and 1
    • _ks_row_mode2 — fused KS distance for SScompare 2
    • _build_mask — behavioural mask construction with zero temporary array allocations
  • Replaced YtempSorted (list of augmented arrays with 1-based MATLAB indices) with YsortedIdx — a clean (nOut, N) array of 0-based sort indices, used by both code paths
  • Numba import is optional: if not installed, falls back to pure NumPy with a warning
  • Added use_numba: bool = True parameter to SAtom() for easy A/B benchmarking

pyproject.toml

  • Removed pytest, pandas, openpyxl from runtime dependencies (not required by satom)
  • Added [fast] optional extra: pip install satom[fast] installs Numba
  • Moved test/dev tools to [dev] extras

Performance

Tested on the Ishigami-Homma function (3 inputs, 1 output, N=500) and the Viborg dataset (10 inputs, 3 outputs, N=5000):

Benchmark NumPy Numba Speedup
Isolated KS row (Ishigami). ~500 µs. ~31 µs. 16×
End-to-end SAtom (Viborg, J=1000). ~0.9 s. ~0.3–0.5 s. 2–3×

The first call incurs a one-time JIT compilation cost (~1–2 s); subsequent calls use the cached compilation.

Backward compatibility

  • All 19 existing tests pass unchanged
  • use_numba=False reproduces the original NumPy behavior exactly
  • Without Numba installed, the package works as before (with a warning)

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.

1 participant