Hybrid numba - #775
Open
maedoc wants to merge 117 commits into
Open
Conversation
…n 2 pickled gpickle files pandas.read_pickle defaults to ASCII encoding which fails on binary numpy data pickled with Python 2. The gpickle format is just a pickled networkx graph, so use stdlib pickle.load with encoding='latin1' instead.
* added numba dfun for model * signature for elementwise if/else, njit for helper funcs * tuple unpack svars * test: fix and test numba kionex * fix: notebook fixes * fix: pickle load error * fix(kionex): shape error in notebook * fix(jansen-rit notebook): numpy 2.4 scalar assignment compat In numpy 2.4+, assigning a non-scalar ndarray to a scalar index raises ValueError instead of a deprecation warning. The expression phi_n_scaling = (jrm.a * jrm.A * ...) yields a (1,) array, so sigma[3] = phi_n_scaling now fails. Fix by adding .item() to extract the scalar value.
* added numba dfun for model * signature for elementwise if/else, njit for helper funcs * feat(kionex): extend numba backend and implement kionex * fix(kionex): fix dfun call
…n 2 pickled gpickle files pandas.read_pickle defaults to ASCII encoding which fails on binary numpy data pickled with Python 2. The gpickle format is just a pickled networkx graph, so use stdlib pickle.load with encoding='latin1' instead.
- Rebase hybrid-numba onto origin/master (picks up KIonEx numba backend commits: Add numba dfun for model KIonEx + Enable numba backend) - Extend NbHybridBackend._check_compatibility to accept KIonEx alongside MontbrioPazoRoxin - Add sin/cos/exp/log shorthands to nb-hybrid-sim.py.mako template header (required by KIonEx dfun_helpers which reference exp/log by name) - Extend nb-hybrid-sim.py.mako to emit per-subnet dfun_constants, dfun_helpers, and dfun_intermediates when the model provides them (KIonEx-style models with ionic current helpers) - New test: tvb/tests/library/simulator/hybrid/test_mpr_kionex.py Pure-Python Simulator.run() smoke test: MPR + KIonEx two-subnet hybrid (output shape, NaN-free at safe ICs, bidirectional inter-projections) - New test class: TestNbHybridMprKIonEx in test_nb_hybrid.py NbHybridBackend compiled kernel: acceptance, shape correctness, and Python-vs-Numba numerical consistency for MPR + KIonEx
Replace log(K_o/K_i) and log(Na_o/Na_i) with log(K_o)-log(K_i) and log(Na_o)-log(Na_i) in all four representation sites: - dfun_helpers strings (used by nb-hybrid-sim codegen) - _numpy_dfun inner helper lambdas - standalone @njit I_K_form / I_Na_form (two duplicate copies) In float32 the ratio K_o/K_i can underflow to zero before the log is taken, yielding -inf. Computing log(a)-log(b) avoids the intermediate ratio and keeps both operands in the normal float32 range as long as each concentration is individually representable. Note: log(Cl_o0/Cl_i0) is left as-is because both values are compile-time constants that are always positive and finite. Also document in TestNbHybridMprKIonEx.test_output_shapes that NaN-freedom for KIonEx in float32 is a separate open item (K_o can go negative under extreme states), while the Python float64 path is already checked in test_mpr_kionex.py.
Add Sigmoidal and SigmoidalJansenRit coupling functions to the numba
hybrid backend, and support the JansenRit model as a second model target
alongside MontbrioPazoRoxin and KIonEx.
Backend changes (nb_hybrid.py):
- _cfun_type(): recognise Sigmoidal -> 'sigmoidal', SigmoidalJR -> 'sigmoidal_jr'
- _cfun_params(): return float32[5] array (replaces (cfun_a, cfun_b) 2-tuple)
Layout: [a, sigma, midpoint, cmin, cmax] for Sigmoidal;
[a, e0, r, v0, 0] for SigmoidalJansenRit
- _check_compatibility(): accept JansenRit alongside MPR and KIonEx
- _run_compiled(): pass cfun_params array (single arg per projection)
Template changes (nb-hybrid-sim.py.mako):
- cfun_a/cfun_b scalar args -> cfun_params float32[5] array throughout
- Pre-cfun hook inside CSR inner loop for sigmoidal_jr (pre-synaptic
sigmoid on source state values before weighted-sum accumulation)
- Post-cfun dispatch: sigmoidal branch added; linear/scaling use cfun_params[0/1]
- Both mono_src (n_modes==1) and general paths updated
Model changes (jansen_rit.py):
- Add coupling_terms, parameter_names, dfun_helpers (sigm_jr helper),
dfun_intermediates, state_variable_dfuns for numba codegen
- coupling enters y4 equation only as Coupling_Term
- 13 scalar parameters baked at codegen time
Tests (test_nb_hybrid.py):
- TestNbHybridSigmoidalCfun: 6 tests (Sigmoidal + SigmoidalJansenRit,
shape + finite + matches-python)
- TestNbHybridJansenRit: 4 tests (acceptance, shape, finite, NB vs Python)
- TestNbHybridMultiMode: 3 tests (MPR with n_modes=2, shape + matches-python)
- Total: 260 tests passing (was 28)
Replace exec()-into-dict with _build_as_module(): renders generated source to a real .py file under $TMPDIR/tvb_nb_hybrid_cache/, registers it in sys.modules, then imports it. This gives Numba a real co_filename so that cache=True on @nb.njit writes .nbi/.nbc native-code files next to the .py. On second process startup with the same network topology the SHA-256 key matches the existing .py, Numba loads the native cache directly (~50 ms vs ~5 s cold JIT). The in-process _COMPILED_FN_CACHE dict is retained as an additional first-level fast path. Changes: - _build_as_module() helper with atomic os.replace write - _build() delegates to _build_as_module instead of exec() - NbHybridBackend.clear_cache() classmethod (in-process + disk) - NbHybridBackend.get_cache_dir() staticmethod - Template: cache=True added to all @nb.njit decorators (inline + plain) - TestNbHybridDiskCache: 3 tests (dir created, in-process hit, clear)
Mark as done: Sigmoidal cfun, SigmoidalJansenRit cfun, JansenRit model, n_modes>1 test, disk-persistent JIT cache. Update test count to 44 (263 total across backend + hybrid suites). Update §6 supported configuration table and §7/§8 to reflect 2026-04-10 work.
… strip CSR zeros Model codegen attributes added (coupling_terms, parameter_names, dfun_intermediates, state_variable_dfuns): - ReducedWongWang (wong_wang.py): 1 svar S, 8 params, 2 intermediates (x_ww, H_ww); S clipped to [0,1] via existing boundaries - Epileptor (epileptor.py): 6 svars, 16 params, 4 ternary intermediates (f1, zterm, h, f2); modification=False only - WilsonCowan (wilson_cowan.py): 2 svars E/I, 22 params, 4 intermediates (x_e, x_i, s_e, s_i); shifted sigmoid (shift_sigmoid=True only) - Generic2dOscillator (oscillator.py): already added in prior session Backend (nb_hybrid.py): - _check_compatibility: accept RWW, Epileptor, WilsonCowan (7 models total) - Per-model constraint checks: rejects Epileptor(modification=True) and WilsonCowan(shift_sigmoid=False) with NotImplementedError - _build_projection_info: strip structural epsilon zeros via eliminate_zeros() on a copy of p.weights before extracting .data/.indices/.indptr; idelays aligned with stripped structure Tests (test_nb_hybrid.py): - TestNbHybridReducedWongWang: 4 tests (accepted, shape, finite, py-match) - TestNbHybridEpileptor: 4 tests - TestNbHybridWilsonCowan: 4 tests (uses source_cvar=[0] for E-only coupling) - test_rejects_unsupported_model: updated to use SupHopf (WilsonCowan now accepted) - 283 tests passing total (was 52)
…/8.9/8.10)
§8.8 Resumable runs / snapshot API:
- CompiledNetworkFn.run(return_snapshot=True) -> (outputs, snapshot)
snapshot = {'states': [...], 'buffers': {...}}; captured from in-place
Numba writes, zero kernel changes required
- CompiledNetworkFn.resume(snapshot, nstep) restores state+buffers and
continues; numerically identical to a single longer run
- _run_compiled gains _initial_buffers=None parameter
§8.9 Testing gaps:
- TestNbHybridModeMap (4 tests): non-identity mode_map on 2-mode MPR nets;
verifies accepted, shape, finite, and that mixing actually differs from identity
- TestNbHybridLargeNScaling (2 tests): N=100 no-error smoke test; N=50
speedup regression (cached Numba vs Python loop with generous bound)
§8.10 Code quality:
- __all__ added to nb_hybrid.py (NbHybridBackend, CompiledNetworkFn,
NetworkAnalysis, SubnetworkInfo, ProjectionInfo)
- backend/__init__.py: lazy __getattr__ for NbHybridBackend + CompiledNetworkFn
(direct import caused circular dep via integrators -> equations -> backend)
- Plan doc: §8.3 marked DONE, §6 table extended, Phase C2 added to §7
Tests: 74 passing (was 68)
…§8.4/8.10) §8.10 Code quality: - compile(debug_nojit=True) / TVB_HYBRID_NO_JIT=1 env var: replaces @nb.njit decorators with no-ops for fast debugging; hash naturally differs (rendered source changes) so no cache collisions - run_network() also accepts debug_nojit kwarg - Type annotations added: _run_compiled, _analyse, _check_compatibility, _build_projection_info - Module docstring references nb_hybrid_plan.md - TestNbHybridDebugNojit: 2 tests (runs without error, matches JIT output) §8.4 Lazy stimulus infrastructure (stub): - _STIM_LAZY_THRESHOLD_MB = 64 (overridable via TVB_HYBRID_LAZY_STIM_MB) - _stim_estimate_mb(sn_info, nstep): projected stim array size in MiB - _compute_stimulus_lazy stub: windowed (step_start, step_end) computation with TODO explaining template change needed (network_chunk must use t_local indexing instead of global t-1 before lazy path can be wired to run_network) - TestStimulusMemoryEstimate: 2 tests (small <64MB, large >64MB) Plan doc (nb_hybrid_plan.md): - Status line updated; §6 table extended (8 new rows); Phase D + E added - §8.2 checkboxes all ticked; §8.6 FHN note clarified; §8.7 marked DONE - §8.8 marked DONE; §8.9 table fully checked; §8.10 checkboxes all ticked Tests: 78 passing (was 74)
nb_hybrid_next.md: plan for bulk model codegen (17 scalar models via Ralph loop), monitor support (M1 Python dispatch, M2 in-kernel subsample, M3 projection monitors), combined-mode dfun generation for ReducedSetFitzHughNagumo/HindmarshRose (unrolled matrix ops at codegen time). ralph_add_model_codegen.sh: bash script that iterates over 17 model files, calling opencode run with zai/glm-5.1 to add coupling_terms, parameter_names, dfun_intermediates, state_variable_dfuns. Includes validation (attr existence, key matching, coupling_term usage, parameter attribute check) and 3-retry with revert on failure. nb_hybrid_plan.md: updated status line, reclassified ReducedSetFHN from 'permanently deferred' to 'deferred to next phase (combined-mode)', updated file summary table.
…emplate (Phase F) Phase F — Bulk Model Codegen (Ralph Loop): - 15 scalar models via AI-assisted Ralph loop: SupHopf, Kuramoto, Epileptor2D, Hopfield, LarterBreakspear, EpileptorRestingState, EpileptorCodim3, EpileptorCodim3SlowMod, ZetterbergJansen, ReducedWongWangExcInh, CoombesByrne, CoombesByrne2D, GastSchmidtKnosche_SD, GastSchmidtKnosche_SF, DumontGutkin - Each model gets coupling_terms, parameter_names, dfun_intermediates, state_variable_dfuns codegen attributes - 2 Zerlaut models via custom nb-zerlaut-dfun.py.mako template: ZerlautAdaptationFirstOrder (5 svars, erfc transfer function pipeline) ZerlautAdaptationSecondOrder (8 svars, numerical-derivative covariance) Combined-mode dfun (nb_hybrid_next.md §3 / G4): - ReducedSetFitzHughNagumo (4 svars × 3 modes): dfun_mode='combined' - ReducedSetHindmarshRose (6 svars × 3 modes): dfun_mode='combined' - Template gains is_combined branch in nb-hybrid-sim.py.mako - Inter-mode matrix products (Aik, Bik, Cik) unrolled at codegen time - derived_matrix_names/ops model attributes for combined-mode metadata Gaps G1–G3: - G1: Multiplicative noise raises NotImplementedError (was silent wrong result) - G2: chunk_size > min_horizon raises ValueError (was silent wrong result) - G3: Monitor M1 Python-side dispatch: monitors= kwarg on run_network() supports TemporalAverage, Raw, SubSample, GlobalAverage, AfferentCoupling Test coverage: 78 → 145 tests (+67), all passing across 28 test classes Models supported: 8 → 27 (19 new via codegen) _ralph_models parametrized smoke tests (accepted + shape + finite) Zerlaut dedicated tests (shape + finite + python-match) Monitor dispatch tests (TestNbHybridMonitors)
…le pytest.ini option
…rs, model dfun corrections, test improvements, notebook fixes Phase 1-4 fixes from hybrid-numba code review (OVERALL_REVIEW.md): - cvar_utils: bare int handling, shared helper extraction - projection_utils: connectivity extraction for intra projections - inter_projection: target cvar validation, target symmetry check - stimulus: resolve_target_cvar instead of resolve_cvar_names - infinite_theta: alpha default 10.0→1.0, docstring equations - jansen_rit: ZetterbergJansen keke→kiki in derivative[9] - larter_breakspear/zerlaut: local_coupling deprecation docs - nb_hybrid_sweep_cpu: NotImplementedError guard for model-param sweeps - nb_hybrid_cuda_sweep: cuda.local.array(76)→MAX_NODES=1024 - nb_hybrid_cuda_sweep_backend: copy BOLD states from GPU - nb-hybrid-sweep-cuda.py.mako: per-source-subnet horizon - test fixes: CSR diagonal, CUDA fixture, eval() sort, unseeded RNG - test_subnetwork: test hybrid Stim instead of StimuliRegion - test_base: fix JansenRit source_cvar indices - test_classic_tvb_coupling_equiv: update stale docs/names - test_toy_synchronization: non-trivial ICs, Kuramoto coupling, sim length - test_model_codegen_parity: new file, 12 model dfun parity tests - test_stimulus_utils: new file, 7 stimulus utility tests - notebooks: axis fix, separate integrators, title restore, warnings
…ove unused dfe_plus, add local_coupling deprecation note
… noise floor corruption, add local_coupling deprecation note
…modification case
…omission in EpileptorRestingState state_variable_dfuns
… raise clear ValueError
…ow with clear error message
…Mod state_variable_dfuns and parameter_names
…e_dfuns and remove pdb debug line
…at dfun_intermediates only encode shift_sigmoid=True branch
…_dfuns in DecoBalancedExcInh to include M_i scaling
…ed for guard check)
…, EpileptorRestingState modification=True, Hopfield dynamic=True, and WilsonCowan shift_sigmoid=False
…field, stefanescu_jirsa
- simulator.py: fix 3 _run_numba bugs (time dup, concat fail, merged monitor IndexError); replace assert with explicit if/raise for dt validation
- base_projection.py: fix lengths=None crash by defaulting to zero-lengths CSR; replace assert with if/raise for nnz check
- hopfield.py: fix state_variable_dfuns theta expression (x→theta, taux→tauT); add tauT to parameter_names
- stefanescu_jirsa.py: fix FHN/HR codegen alpha expressions (c_alpha_{m}*mu→c_xi_{m}) to match Python dfun
NetworkSet no longer accepts 'stimuli' kwarg (moved to Subnetwork.stimuli). Update 28 NetworkSet(stimuli=[...]) calls to assign sn.stimuli = [...] instead.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
…dation - Two scenarios: CRBL-only (open-loop, 27 nodes) and WW+CRBL (closed-loop, 126 nodes) - 10s simulations at dt=1ms (~10s wall-clock total, python backend) - Per-subnet recorders capture all 4 CRBL populations (GrC, GoC, MLI, PC) - Welch PSD analysis with frequency band overlays (delta/theta/alpha/beta/gamma) - Carrier frequency comparison table across populations and scenarios - Scientific background on construct validation (Esaghei 2022, Lorenzi 2023/2025) - Inline figures in notebook via %pylab inline - CerebellarMF model registered in models/__init__.py
…F-TVB Critical fixes validated against monolithic reference simulation: Model parameter defaults (cerebellar_mf.py): - weight_noise: 10.5 → 4e-3 (was 2625x too high) - external_input_ex_ex: 0.0 → 3.15e-4 (was missing entirely) - tau_OU: 5.0 → 3.5 (1.43x off) - state_variable_range: match Lorenzi 2023 ICs [0.5e3, 5e3, 15e3, 38e3] Anatomical routing fractions (new parameters): - frac_mossy=0.57, frac_parallel=0.43, mf_to_grc=0.97, mf_to_goc=0.03, pf_to_goc=0.14, pf_to_mli=0.55, pf_to_pc=0.31 - These decompose the monolithic's single cvar coupling into the hybrid's explicit mossy/parallel pathways in the dfun Demo script updates: - Column-normalize SC weights (matching monolithic init) - Add P4 (GrC→mossy) IntraProjection alongside P5 (GrC→parallel) - Use nperseg=1024 for Welch PSD (avoids resolution artifacts) - Open-loop firing rates now match monolithic within ~10%: GrC 0.17 vs 0.18, GoC 1.54 vs 1.73, MLI 4.72 vs 5.26, PC 11.85 vs 13.22
…reproduction notebook - Add use_legacy_goc_e_e flag to CerebellarMF for reproducing published cMF-TVB results (GoC E_i bug in parallel_crbl.py:583) - Fix production parameter defaults: alpha_mli=5.0, external_input=0.0, frac_mossy/parallel=1.0, kHz-scale initial conditions - Add external_input_in_ex parameter for GoC drive (matches monolithic) - Add add_noise_mli_pc flag for CRBL-only anatomical routing - Update Numba Mako template (nb-cerebellar-dfun) for legacy mode support - Fix noise variable name collision in nb-hybrid-sim.py.mako: rename stochastic noise array param from 'noise' to '_nsig_arr' to avoid shadowing CerebellarMF's 'noise' state variable - Update tests: legacy+fixed GoC parity, production defaults - Add reproduction notebook: UC-1 (CRBL-only) + UC-2 (WW+CRBL) with both Python and Numba backends, reference comparison (39.1 Hz exact)
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
4 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.
In anticipation of merging #771 this adds a numba backend to accelerate hybrid simulations.