Hybrid numba - #777
Closed
Ziaeemehr wants to merge 36 commits into
Closed
Conversation
* 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)
…G1/G2/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 via monitors= kwarg
supports TemporalAverage, Raw, SubSample, GlobalAverage, AfferentCoupling
…igmoidal cfuns All 9 hybrid coupling functions now supported: - Linear, Scaling, Sigmoidal, SigmoidalJansenRit (existing) - Kuramoto: a * sin(x) post-hook - Difference: a * x (maps to scaling internally) - HyperbolicTangent: a * (1 + tanh((x-midpoint)/sigma)) pre-hook - PreSigmoidal: H * (Q + tanh(G*(P*x - theta))) pre-hook Template fix: pre-cfun transforms moved outside CSR accumulation loop to match Python pipeline (weighted sum -> pre -> scale -> post). This also fixes sigmoidal_jr correctness for heterogeneous source states. _cfun_params array expanded from 5 to 8 elements for new parameter layouts. 6 new tests in TestNbHybridCfunExtended (all python-vs-numba comparisons). 151 tests total, all passing.
…match tests for all models Template bug fix: tavg accumulation now uses actual variables_of_interest indices instead of state[0..n_voi-1]. This was wrong for any model where voi aren't the first N svars (e.g., ZetterbergJansen voi=v6,v7,v2-v5, Epileptor voi=x2-x1,z). Derived voi expressions (e.g., 'x2 - x1') are now parsed and generated as arithmetic in the template. _ic_from_range() fixed to produce correct multi-mode ICs for ReducedSet models (n_modes=3). Added parametrized test_ralph_model_matches_python for all 17 Ralph-loop models — compares full Numba backend output against Python reference loop. Epileptor match test updated for derived voi. 168 tests total, all passing.
…onitors (M2-M5)
All monitors implemented as Python-side post-processing in _apply_monitors():
M2 - SubSample: period enforcement via time-point masking, selecting chunks
whose mid-time falls within dt/2 of a period boundary.
M3 - SpatialAverage: node grouping via spatial_mean matrix (n_areas, n_nodes)
applied with einsum to reduce node dimension.
M4 - Projection (EEG/MEG/iEEG): gain-matrix projection from node space to
sensor space. Handles mode summation and adds singleton mode dim.
M5 - Bold: HRF convolution via rolling stock buffer with interim accumulation.
Stateful across multiple run_network() calls via _nb_ prefixed attrs.
Supports FirstOrderVolterra k1/V0 scaling.
Tests: 7 new monitor tests covering shape, values, period spacing, and
statefulness. Updated test_unsupported_monitor_raises to use ProgressLogger.
175 tests total, all passing.
…n tests Each test runs the same simulation through both the Python hybrid step-loop and the Numba backend, applies the same monitor logic, and asserts numerical equivalence at rtol=1e-2: - TemporalAverage (chunk_size=1 and chunk_size=5) - SubSample (step-based period masking, matching Python istep semantics) - GlobalAverage (mean across nodes) - SpatialAverage (spatial_mean matrix grouping) - Projection/EEG (gain-matrix sensor projection) - Bold (HRF convolution with FirstOrderVolterra scaling) Bug fix: SubSample mask changed from float time-point rounding to step-based modular arithmetic to avoid float32 precision artifacts. 182 tests total, all passing.
…umba hybrid stepping
…mparison between Python and Numba backends
…ng new models and dynamic simulation lengths
…dices in main function
…_case and main functions
… output formatting
Member
|
hi @Ziaeemehr (and @i-Zaak ) we could plan a sync this week or next, since I've now got hybrid simulation mostly covered by a new backend in #775 |
Author
|
works for me. I am available. |
Contributor
|
Sure! |
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.
Added two example for benchmarking and visualization of corresponding python and numba backends of multiple models.
we can also just add to tests but I think it's good demonstration of model usage.
Feel free to modify or remove.