Skip to content

compile DynamicColoredNoise IIR loop with numba - #5

Merged
cboulay merged 1 commit into
devfrom
simulator-optimization
Jun 30, 2026
Merged

compile DynamicColoredNoise IIR loop with numba#5
cboulay merged 1 commit into
devfrom
simulator-optimization

Conversation

@kylmcgr

@kylmcgr kylmcgr commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

Summary

The per-sample Kasdin IIR recursion in DynamicColoredNoise was a pure-Python
double loop (over input bins × output samples) and the dominant cost in the
velocity→ecephys encoder. This ports the hot loop verbatim to a numba njit
kernel. The recursion is inherently sequential (output[i] feeds the delay
line that produces output[i+1]), so it can't vectorize over time — compiling
the scalar loop is the right tool.

No behavior change: the white-noise draw stays in NumPy (RNG sequence
preserved) and the MLX/array-api input path is untouched (the loop already ran
on a NumPy view). Channels are independent, so iterating over them inside the
sample loop is equivalent to the old vectorized-over-channels form; non-finite
β channels skip the coefficient update (an EMA-toward-self no-op).

Results

Per-stage encoder benchmark (bench_cosine_encoder.py, 256ch @ 30 kHz):

chunk lfp.noise before after speedup
1 (live path) 0.686 ms/call 0.038 ms/call ~18×
50 33.85 ms/call 1.15 ms/call ~29×

End-to-end (simulator over LSL, 50 Hz, 256ch): PINK_NOISE dropped from
~0.70 ms/msg to 0.088 ms/msg — no longer the pipeline's longest stage.

Correctness

  • Output matches the previous implementation to max abs diff 1.1e-13
    (floating-point reassociation only: the old code's np.einsum pole reduction
    sums in a different order than the scalar loop). Same-seed determinism holds.
  • All 117 unit tests pass (tests/unit/). The compiled loop also cuts the
    suite's runtime from ~28.5 s to ~2.5 s.

Dependency

Declares numba explicitly (it was already present transitively via
ezmsg-event).

@kylmcgr
kylmcgr requested a review from cboulay June 29, 2026 17:24
@kylmcgr kylmcgr self-assigned this Jun 29, 2026
@cboulay
cboulay merged commit f207956 into dev Jun 30, 2026
14 checks passed
@cboulay
cboulay deleted the simulator-optimization branch June 30, 2026 04:47
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