Skip to content

Speed up sorted-scale runs: 306 units 3258 s → 463 s (1.15.0) - #170

Merged
Timothysit merged 1 commit into
mainfrom
perf/nmf-ranks-numba-nulls
Oct 25, 2026
Merged

Timothysit merged 1 commit into
mainfrom
perf/nmf-ranks-numba-nulls

Conversation

@Timothysit

Copy link
Copy Markdown
Collaborator

Prepares the Python pipeline for spike-sorted data with hundreds of units. On the 306-unit sorted recording (steps 1–4, full figures, one recording):

Before After
Step 3 (STTC + 200 shuffles × 3 lags) 1,494 s 121 s
Step 4 1,749 s 332 s
Total 3,258 s 463 s

Adjacency matrices and the network CSVs are identical to the old run.

Changes

  • STTC (sttc.py): numba prange across rows, paired from both ends to balance load: 2.2 s → 0.22 s per matrix. If numba's threading layer is missing (e.g. a frozen app), it warns and runs serially with the same results. The build spec now collects numba.np.ufunc, and app_smoke checks it.
  • NMF (nmf.py): 933 s → 26 s.
    • Fits run on threads. Seeds are drawn in the original order and BLAS is pinned to 1 thread, so results don't depend on thread count.
    • The rank sweep stops at the 95% crossing, or when the curve levels off under the new nmf_plateau_gain (default 0.005; 0 = MATLAB's full sweep). Skipped ranks are NaN.
    • The setting has a widget under Advanced in the Network metrics box on the Connectivity tab and is read by both pipelines.
  • Normalised participation coefficient (null_models.py, network_metrics.py): 379 s → 33 s. All random draws for the 100 null models depend only on the original edge counts, so they are planned up front in order and the rewiring then runs on threads with identical results.
  • randmio/latmio v2 and randmio_und_signed: compiled with numba and consume the same RNG stream as before.

Caveats

  • NMF arrays differ from old runs by about 1e-14, because of the single-thread BLAS pin. Old runs already varied by this much between machines.
  • On sorted data the variance curve rises roughly linearly and never plateaus, so in practice the sweep stops at 95% (rank 228/306 here).
  • Figures are now most of the remaining time: about 250 s of step 4, of which plot_network takes about 190 s.
  • The frozen-app threading layer still needs to be confirmed on a real build.
  • This PR bumps to 1.15.0 / CAT-NAP 1.11.0, the same bump as the spatiotemporal-autocorrelation branch. Whichever merges second must re-bump and update the section count in test_gui_advanced.

Tests

  • test_pipeline_fast_paths: 91/91. It compares each rewritten piece against the replaced code (including RNG state) at 1 and many threads.
  • MATLAB parity tests for steps 3 and 4 pass (step 4: 60/60).
  • test_step4_optional_metrics 21/21; CAT-NAP 112/112; catnap_subnetwork 45/45; sorted_import, null_models, nmf and small_worldness pass.
  • Full suite: the only failures are test_pipeline_streaming_pool, test_viewer_server and test_stats_report, which also fail on main.

🤖 Generated with Claude Code

https://claude.ai/code/session_01LBjNAqNrYVFQSwnaKgi25j

…op, threaded exact null models (1.15.0)

On the 306-unit spike-sorted recording, steps 1-4 drop from 3258 s to 463 s.

- STTC: numba prange across rows (paired from both ends), ~10x per matrix;
  serial fallback if the threading layer is missing (frozen app).
- NMF: fits run on threads with seeds drawn in the original order and BLAS
  pinned to 1 thread, so results don't depend on thread count; the rank
  sweep stops at the 95% crossing or when the curve levels off
  (new nmf_plateau_gain, default 0.005; 0 = MATLAB's full sweep), with a
  GUI widget on the Connectivity tab.
- participation_coef_norm: null-model random draws are planned up front
  in order, then rewiring runs on threads with identical results.
- randmio/latmio v2 and randmio_und_signed compiled with numba, same RNG
  stream as before.
- Equivalence tests against the replaced code in test_pipeline_fast_paths.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LBjNAqNrYVFQSwnaKgi25j
@Timothysit
Timothysit merged commit 1f40eb4 into main Oct 25, 2026
2 of 3 checks passed
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