From 47f64778526ff53fba5868f37955d5d4e57e0e2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:09 +0200 Subject: [PATCH 01/48] feat(models): update exports to include dsl submodule --- src/neuronumba/simulator/models/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/neuronumba/simulator/models/__init__.py b/src/neuronumba/simulator/models/__init__.py index 547cbdc..b6a12ee 100644 --- a/src/neuronumba/simulator/models/__init__.py +++ b/src/neuronumba/simulator/models/__init__.py @@ -3,4 +3,5 @@ from neuronumba.simulator.models.deco2014 import Deco2014 from neuronumba.simulator.models.hopf import Hopf from neuronumba.simulator.models.montbrio import Montbrio -from neuronumba.simulator.models.zerlaut import ZerlautAdaptationFirstOrder, ZerlautAdaptationSecondOrder \ No newline at end of file +from neuronumba.simulator.models.zerlaut import ZerlautAdaptationFirstOrder, ZerlautAdaptationSecondOrder +from neuronumba.simulator.models import dsl # noqa: F401 \ No newline at end of file From 69914622085941d8dd0e1faf631d070ab682609e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:14 +0200 Subject: [PATCH 02/48] docs(dsl_models): add readme for dsl model specs and usage --- examples/dsl_models/README.md | 54 +++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 examples/dsl_models/README.md diff --git a/examples/dsl_models/README.md b/examples/dsl_models/README.md new file mode 100644 index 0000000..bdce7c7 --- /dev/null +++ b/examples/dsl_models/README.md @@ -0,0 +1,54 @@ +# Reference DSL specs + +This directory contains the canonical neuronumba models expressed in the DSL. +Each spec produces a `Model` subclass that is **bit-equivalent** to its +hand-written counterpart in `src/neuronumba/simulator/models/`. + +| Spec | Hand-written | LOC ratio | +|---|---|---| +| [hopf_dsl.py](hopf_dsl.py) | `hopf.py` | ~30 vs 152 | +| [deco2014_dsl.py](deco2014_dsl.py) | `deco2014.py` | ~50 vs 416 | +| [naskar2021_dsl.py](naskar2021_dsl.py) | `naskar2021.py` | ~60 vs 104 | +| [montbrio_dsl.py](montbrio_dsl.py) | `montbrio.py` | ~70 vs 160 | + +## Dual purpose + +These specs are both **examples** and **test references**: + +- **Examples** — read them top-to-bottom to learn the DSL. Copy any of them + into your own code and modify; you'll have a working model in minutes. +- **Test references** — `tests/test_dsl_equivalence.py` imports each spec, + builds the DSL model, and compares it against the hand-written model on + random inputs. Disagreement fails the test. + +## Drift policy + +If you change a hand-written model (bug fix, numerical refinement), update the +DSL spec here in the same commit. The equivalence test will fail loudly if you +forget — read the failure message: it tells you both files involved. + +If you change a DSL spec (e.g. for a new feature you want to add), the +equivalence test still has to pass against the hand-written reference. If they +genuinely diverge, that's a model change and belongs in a separate PR. + +## Models not included + +- **Zerlaut** — needs `@njit` helper subroutines (`get_fluct_regime_vars`, + `threshold_func`, `erfc_approx`). The DSL doesn't support that yet — see the + Phase 6 deferred items in `IMPLEMENTATION_PLAN.md`. +- **OrnsteinUhlenbeck** — needs non-scalar matrix parameters (`A` derived from + weights/g/tau). Deferred to v0.2. + +For these, subclass `Model` directly. + +## Usage + +```python +from examples.dsl_models import HopfDSL # if examples/ is on sys.path +# or +from neuronumba.simulator.models.dsl import build_model +from examples.dsl_models.hopf_dsl import hopf_spec +HopfDSL = build_model(hopf_spec) + +m = HopfDSL(g=0.5).configure(weights=W) +``` From c6deab65c91c579ee98b5efab0bde096cf33ceae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:18 +0200 Subject: [PATCH 03/48] feat(dsl_models): add package init with dsl model exports --- examples/dsl_models/__init__.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 examples/dsl_models/__init__.py diff --git a/examples/dsl_models/__init__.py b/examples/dsl_models/__init__.py new file mode 100644 index 0000000..abceb92 --- /dev/null +++ b/examples/dsl_models/__init__.py @@ -0,0 +1,30 @@ +"""Reference DSL specs for neuronumba's canonical models. + +These specs reproduce the existing hand-written models exactly. They serve two +purposes simultaneously: + +1. **User-facing examples.** Read them to learn how to express a model in the + DSL. Each file is self-contained and small (~30-60 LOC vs 100-400 LOC for + the hand-written equivalents). + +2. **Regression test references.** `tests/test_dsl_equivalence.py` imports + these specs and compares the resulting DSL-built models against the + hand-written ones bit-by-bit. If you change a hand-written model (e.g. a + numerical bug fix), the corresponding spec here must follow, or the + equivalence test will fail. + +Models out of scope for this set: + - Zerlaut: needs @njit helper subroutines (deferred to v0.2) + - OrnsteinUhlenbeck: needs non-scalar matrix parameters (deferred to v0.2) +""" +from .deco2014_dsl import deco_spec, Deco2014DSL +from .hopf_dsl import hopf_spec, HopfDSL +from .montbrio_dsl import montbrio_spec, MontbrioDSL +from .naskar2021_dsl import naskar_spec, Naskar2021DSL + +__all__ = [ + "Deco2014DSL", "deco_spec", + "HopfDSL", "hopf_spec", + "MontbrioDSL", "montbrio_spec", + "Naskar2021DSL", "naskar_spec", +] From c539bf2ca6ac6c3e9fe461dc7439d0dcb05b0903 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:23 +0200 Subject: [PATCH 04/48] feat(dsl_models): add deco2014 dsl spec with dfun and coupling --- examples/dsl_models/deco2014_dsl.py | 67 +++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) create mode 100644 examples/dsl_models/deco2014_dsl.py diff --git a/examples/dsl_models/deco2014_dsl.py b/examples/dsl_models/deco2014_dsl.py new file mode 100644 index 0000000..2149507 --- /dev/null +++ b/examples/dsl_models/deco2014_dsl.py @@ -0,0 +1,67 @@ +"""Reduced Wong-Wang / Deco 2014 dynamic mean-field, expressed in the DSL. + +DSL equivalent of `neuronumba.simulator.models.deco2014.Deco2014` (416 LOC). +Note this spec covers only the dfun + linear coupling — the hand-written class +also has `auto_fic` (FIC steady-state computation) which is not part of the +core differential equations. To use FIC with this DSL model you would override +`_init_dependant` in a subclass. See `tests/test_dsl_subclass.py` for the +subclassing pattern. + +`tests/test_dsl_equivalence.py::test_equivalence_deco2014` asserts dfun +equivalence against the hand-written class to 1e-12 (with `auto_fic=False` so +both share the same `J=1.0`). +""" +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +# The hand-written Deco2014 has an EPS-fix to avoid div-by-0 in the firing-rate +# sigmoid. We faithfully reproduce that with `np.where`. +deco_spec = ModelSpec( + name="Deco2014DSL", + state_vars=[ + StateVar("S_e", initial=0.001, bounds=(0.0, 1.0)), + StateVar("S_i", initial=0.001, bounds=(0.0, 1.0)), + ], + coupling_vars=[CouplingVar("S_e", kind="linear")], + observables=["Ie", "re"], + parameters=[ + Parameter("tau_e", default=100.0), + Parameter("tau_i", default=10.0), + Parameter("gamma_e", default=0.000641), + Parameter("gamma_i", default=0.001), + Parameter("I0", default=0.382), + Parameter("Jext_e", default=1.0), + Parameter("Jext_i", default=0.7), + Parameter("w", default=1.4), + Parameter("J_NMDA", default=0.15), + Parameter("J", default=1.0), + Parameter("ae", default=310.0), + Parameter("be", default=125.0), + Parameter("de", default=0.16), + Parameter("ai", default=615.0), + Parameter("bi", default=177.0), + Parameter("di", default=0.087), + Parameter("I_external", default=0.0), + ], + equations=""" + Ie = Jext_e * I0 + w * J_NMDA * S_e + J_NMDA * coupling.S_e - J * S_i + I_external + Ii = Jext_i * I0 + J_NMDA * S_e - S_i + + y_e = ae * Ie - be + denom_e = 1.0 - np.exp(-de * y_e) + denom_e = np.where(np.abs(denom_e) < 1e-12, 1e-12, denom_e) + re = y_e / denom_e + + y_i = ai * Ii - bi + denom_i = 1.0 - np.exp(-di * y_i) + denom_i = np.where(np.abs(denom_i) < 1e-12, 1e-12, denom_i) + ri = y_i / denom_i + + d_S_e = -S_e / tau_e + gamma_e * (1.0 - S_e) * re + d_S_i = -S_i / tau_i + gamma_i * ri + """, +) + +Deco2014DSL = build_model(deco_spec) From 666e369c8ce7b4f56db5b7eec2152ddde943c2fd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:27 +0200 Subject: [PATCH 05/48] feat(dsl_models): add hopf dsl spec with diffusive coupling --- examples/dsl_models/hopf_dsl.py | 38 +++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 examples/dsl_models/hopf_dsl.py diff --git a/examples/dsl_models/hopf_dsl.py b/examples/dsl_models/hopf_dsl.py new file mode 100644 index 0000000..4b37ef0 --- /dev/null +++ b/examples/dsl_models/hopf_dsl.py @@ -0,0 +1,38 @@ +"""Hopf supercritical bifurcation model, expressed in the neuronumba DSL. + +This is the DSL equivalent of `neuronumba.simulator.models.hopf.Hopf` (152 LOC +hand-written). The math is identical — `tests/test_dsl_equivalence.py:: +test_equivalence_hopf` compares the dfun and coupling-kernel outputs against +the hand-written class on random inputs and asserts agreement to 1e-12. + +If this file changes, that test must still pass; if `hopf.py` changes (e.g. a +numerical bug fix), this file must follow. +""" +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +hopf_spec = ModelSpec( + name="HopfDSL", + state_vars=[ + StateVar("x", initial=0.1), + StateVar("y", initial=0.1), + ], + coupling_vars=[ + CouplingVar("x", kind="diffusive"), + CouplingVar("y", kind="diffusive"), + ], + observables=[], + parameters=[ + Parameter("a", default=-0.5), + Parameter("omega", default=0.3), + Parameter("I_external", default=0.0), + ], + equations=""" + d_x = (a - x*x - y*y) * x - omega * y + coupling.x + I_external + d_y = (a - x*x - y*y) * y + omega * x + coupling.y + """, +) + +HopfDSL = build_model(hopf_spec) From e7bf02c94c26ae503945ef5dcca3fab32d64b912 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:31 +0200 Subject: [PATCH 06/48] feat(dsl_models): add montbrio dsl spec with six state variables --- examples/dsl_models/montbrio_dsl.py | 88 +++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 examples/dsl_models/montbrio_dsl.py diff --git a/examples/dsl_models/montbrio_dsl.py b/examples/dsl_models/montbrio_dsl.py new file mode 100644 index 0000000..eac3edc --- /dev/null +++ b/examples/dsl_models/montbrio_dsl.py @@ -0,0 +1,88 @@ +"""Montbrio firing-rate model, expressed in the DSL. + +DSL equivalent of `neuronumba.simulator.models.montbrio.Montbrio` (160 LOC). +Six state variables (r_e, r_i, u_e, u_i, S_ee, S_ie) and four *dependent* +parameters — values computed once at `configure()` time from the independent +parameters: + + J_N_ee = J_ee + g_ee * log(a_e) + J_N_ie = J_ie + g_ie * log(a_e) + J_G_ei = J_ei + g_ei * log(a_i) + J_G_ii = J_ii + g_ii * log(a_i) + +This is the canonical example of the DSL's dependent-parameter feature. Modify +any of `J_ee`, `g_ee`, `a_e`, `J_ie`, `g_ie`, `J_ei`, `g_ei`, `a_i`, `J_ii`, or +`g_ii` and call `configure()` again — the dependent values are re-derived +automatically. + +`tests/test_dsl_equivalence.py::test_equivalence_montbrio` asserts dfun +equivalence to 1e-10 against the hand-written class. +""" +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +montbrio_spec = ModelSpec( + name="MontbrioDSL", + state_vars=[ + StateVar("r_e", initial=0.0), + StateVar("r_i", initial=0.0), + StateVar("u_e", initial=0.0), + StateVar("u_i", initial=0.0), + StateVar("S_ee", initial=0.0), + StateVar("S_ie", initial=0.0), + ], + coupling_vars=[CouplingVar("S_ee", kind="linear")], + observables=[], + parameters=[ + # Time constants + Parameter("tau_e", default=10.0), + Parameter("tau_i", default=10.0), + Parameter("tau_N", default=10.0), + # Firing-rate parameters + Parameter("delta_e", default=1.0), + Parameter("delta_i", default=1.0), + Parameter("eta_e", default=1.0), + Parameter("eta_i", default=1.0), + # Synaptic parameters + Parameter("a_e", default=0.25), + Parameter("a_i", default=1.0), + Parameter("g_e", default=2.5), + Parameter("g_i", default=0.0), + Parameter("g_ee", default=2.5), + Parameter("g_ei", default=0.0), + Parameter("g_ie", default=2.5), + Parameter("g_ii", default=0.0), + # External inputs + Parameter("I_e_ext", default=0.0), + Parameter("I_i_ext", default=0.0), + # Coupling strengths + Parameter("J_A", default=1.0), + Parameter("J_ee", default=10.0), + Parameter("J_ei", default=10.0), + Parameter("J_ie", default=10.0), + Parameter("J_ii", default=10.0), + Parameter("J", default=10.0), + # Dependent parameters — recomputed every configure() + Parameter("J_N_ee", formula="J_ee + g_ee * np.log(a_e)"), + Parameter("J_N_ie", formula="J_ie + g_ie * np.log(a_e)"), + Parameter("J_G_ei", formula="J_ei + g_ei * np.log(a_i)"), + Parameter("J_G_ii", formula="J_ii + g_ii * np.log(a_i)"), + ], + equations=""" + I_e = I_e_ext + tau_e * S_ee - J * J_G_ei * tau_i * r_i + J_A * tau_e * coupling.S_ee + I_i = I_i_ext + tau_e * S_ie - J_G_ii * tau_i * r_i + + d_r_e = (delta_e / (np.pi * tau_e) + 2.0 * r_e * u_e - g_e * r_e) / tau_e + d_r_i = (delta_i / (np.pi * tau_i) + 2.0 * r_i * u_i - g_i * r_i) / tau_i + + d_u_e = (eta_e + u_e * u_e - r_e * np.pi * tau_e * (r_e * np.pi * tau_e) + I_e) / tau_e + d_u_i = (eta_i + u_i * u_i - r_i * np.pi * tau_i * (r_i * np.pi * tau_i) + I_i) / tau_i + + d_S_ee = (-S_ee + J_N_ee * r_e) / tau_N + d_S_ie = (-S_ie + J_N_ie * r_e) / tau_N + """, +) + +MontbrioDSL = build_model(montbrio_spec) From 6798263fea17ba83200b5e446a9245152d5b68e1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:34 +0200 Subject: [PATCH 07/48] feat(dsl_models): add naskar2021 dsl spec with inhibitory plasticity --- examples/dsl_models/naskar2021_dsl.py | 69 +++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) create mode 100644 examples/dsl_models/naskar2021_dsl.py diff --git a/examples/dsl_models/naskar2021_dsl.py b/examples/dsl_models/naskar2021_dsl.py new file mode 100644 index 0000000..5cfc56b --- /dev/null +++ b/examples/dsl_models/naskar2021_dsl.py @@ -0,0 +1,69 @@ +"""Naskar 2021 multiscale dynamic mean-field with inhibitory plasticity. + +DSL equivalent of `neuronumba.simulator.models.naskar2021.Naskar2021` (104 LOC). +Three state vars (S_e, S_i, J — where J is the plasticity term that's a state +variable, not a parameter) and two observables (Ie, re). + +Note: the hand-written sigmoid here has *no* eps-fix (unlike Deco2014). This +DSL spec faithfully reproduces that. `tests/test_dsl_equivalence.py:: +test_equivalence_naskar2021` asserts dfun equivalence to 1e-10 — slightly +looser than Hopf/Deco because tiny denominator differences from float +reordering propagate without the eps-fix. The math is identical. +""" +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +naskar_spec = ModelSpec( + name="Naskar2021DSL", + state_vars=[ + StateVar("S_e", initial=0.001, bounds=(0.0, 1.0)), + StateVar("S_i", initial=0.001, bounds=(0.0, 1.0)), + StateVar("J", initial=1.0), + ], + coupling_vars=[CouplingVar("S_e", kind="linear")], + observables=["Ie", "re"], + parameters=[ + Parameter("t_glu", default=7.46), # glutamate concentration + Parameter("t_gaba", default=1.82), # GABA concentration + Parameter("We", default=1.0), # external→excitatory scaling + Parameter("Wi", default=0.7), # external→inhibitory scaling + Parameter("I0", default=0.382), # overall external input + Parameter("w", default=1.4), # recurrent self-excitation weight + Parameter("J_NMDA", default=0.15), # NMDA current + Parameter("M_e", default=1.0), + Parameter("ae", default=310.0), + Parameter("be", default=125.0), + Parameter("de", default=0.16), + Parameter("ai", default=615.0), + Parameter("bi", default=177.0), + Parameter("di", default=0.087), + Parameter("M_i", default=1.0), + Parameter("alfa_e", default=0.072), # NMDA forward rate + Parameter("alfa_i", default=0.53), # GABA forward rate + Parameter("B_e", default=0.0066), # NMDA backward rate (ms^-1) + Parameter("B_i", default=0.18), # GABA backward rate (ms^-1) + Parameter("gamma", default=1.0), # plasticity learning rate + Parameter("rho", default=3.0), # target firing rate (Hz) + Parameter("I_external", default=0.0), + ], + # Divisions by 1000 convert to per-millisecond rates, matching the + # hand-written model exactly. + equations=""" + Ie = We * I0 + w * J_NMDA * S_e + J_NMDA * coupling.S_e - J * S_i + I_external + Ii = Wi * I0 + J_NMDA * S_e - S_i + + y_e = M_e * (ae * Ie - be) + re = y_e / (1.0 - np.exp(-de * y_e)) + + y_i = M_i * (ai * Ii - bi) + ri = y_i / (1.0 - np.exp(-di * y_i)) + + d_S_e = -S_e * B_e + alfa_e * t_glu * (1.0 - S_e) * re / 1000.0 + d_S_i = -S_i * B_i + alfa_i * t_gaba * (1.0 - S_i) * ri / 1000.0 + d_J = gamma * (ri / 1000.0) * (re - rho) / 1000.0 + """, +) + +Naskar2021DSL = build_model(naskar_spec) From 2a7f3a6dd407c3a61eee765176f6a5a72a3c5112 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:38 +0200 Subject: [PATCH 08/48] docs(dsl): add readme for dsl declarative model system --- src/neuronumba/simulator/models/dsl/README.md | 293 ++++++++++++++++++ 1 file changed, 293 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/README.md diff --git a/src/neuronumba/simulator/models/dsl/README.md b/src/neuronumba/simulator/models/dsl/README.md new file mode 100644 index 0000000..f0ceb60 --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/README.md @@ -0,0 +1,293 @@ +# DSL: declarative whole-brain models for neuronumba + +A `ModelSpec` describes a neuronumba model as data: state variables, coupling +variables, parameters, and an equations string. `build_model(spec)` compiles +that into a `Model` subclass that plugs into `Simulator` exactly like a +hand-written model — same numba compilation, same cache behavior, same +trajectories. + +## Why a DSL + +Every model in `neuronumba/simulator/models/` is a `Model` subclass that +hand-implements four things: state-variable names, coupling kernel, parameter +list, and a numba-compiled `dfun`. Most of that is boilerplate. Inside +`get_numba_dfun` you write 16 lines of `m[np.intp(P.x)]` parameter unpacking +and `state[i, :]` indexing before getting to the math. A "small variation" — +swap a sigmoid for a softplus in Deco2014 — currently means a new file. + +The DSL pitch: write only the math, the parameter list, and the topology +declarations. The boilerplate is generated. The result is the same numba +closure structure you'd write by hand — proven bit-equivalent on Hopf, Deco2014, +Naskar2021, and Montbrio (see `tests/test_dsl_equivalence.py`). + +## Hello world + +```python +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + +hopf_spec = ModelSpec( + name="HopfDSL", + state_vars=[ + StateVar("x", initial=0.1), + StateVar("y", initial=0.1), + ], + coupling_vars=[ + CouplingVar("x", kind="diffusive"), + CouplingVar("y", kind="diffusive"), + ], + observables=[], + parameters=[ + Parameter("a", default=-0.5), + Parameter("omega", default=0.3), + Parameter("I_external", default=0.0), + ], + equations=""" + d_x = (a - x*x - y*y) * x - omega * y + coupling.x + I_external + d_y = (a - x*x - y*y) * y + omega * x + coupling.y + """, +) + +HopfDSL = build_model(hopf_spec) + +m = HopfDSL(g=0.5).configure(weights=W) +``` + +The hand-written `Hopf` is 152 LOC. The DSL spec is ~30 LOC of declaration. +Both produce identical trajectories. + +## ModelSpec fields + +| Field | Type | Purpose | +|---|---|---| +| `name` | `str` | Class name of the generated model | +| `state_vars` | `list[StateVar]` | Row order matters: defines `state[0, :]`, `state[1, :]`, ... | +| `coupling_vars` | `list[CouplingVar]` | Subset of state vars that are coupled across regions | +| `observables` | `list[str]` | Names of intermediates the simulator should expose for monitoring | +| `parameters` | `list[Parameter]` | Independent + dependent (formula-driven) parameters | +| `equations` | `str` | Python-source body. See "The equation language" below | +| `initial_state_overrides` | `dict[str, float]` | Optional per-state-var initial-value overrides | + +`StateVar(name, initial=0.0, bounds=None)`. `bounds` is an optional +`(lo, hi)` tuple — passed straight to neuronumba's bounds-clamping pipeline. + +## Parameters + +### Independent + +A parameter is independent if it carries `default=` (or `required=True`) and +no `formula`. Independent parameters are user-settable and packed into +`self.m`: + +```python +Parameter("tau_e", default=100.0) # default value +Parameter("g", required=True) # caller must supply +``` + +### Dependent + +A parameter is dependent if it carries a `formula`. The formula is a Python +expression that may reference other parameters and `np`. Dependents are +computed at `configure()` time and recomputed every time the user calls +`configure()` again — that's the existing neuronumba protocol for "a +parameter was modified": + +```python +Parameter("J_ee", default=10.0), +Parameter("g_ee", default=2.5), +Parameter("a_e", default=0.25), +Parameter("J_N_ee", formula="J_ee + g_ee * np.log(a_e)"), +``` + +Order doesn't matter — dependents are topologically sorted at build time. +Cycles and self-references raise `ValueError`. Multi-level dependencies +work: a dependent can reference another dependent. + +### `regional` vs global + +By default parameters are `regional=True`: they're per-ROI scalars packed +into `self.m`. Pass `regional=False` to mark a parameter as global (not +per-region); these go into `self.m_aux`. Most parameters are regional. + +### Why no `default + formula`? + +Setting both is a signal of confused intent: a dependent value cannot also +have a default. The validator rejects this combination at `Parameter` +construction time. Same for `required=True` + `formula`. + +## Coupling kinds + +The DSL ships two pre-built kernels: + +### `linear` + +``` +coupling = g * (W^T @ S) +``` + +Standard structural-connectivity coupling. Most neuronumba models use this. + +### `diffusive` + +``` +coupling = g * (W^T @ S - sum(W^T, axis=1) * S) +``` + +Hopf-style: subtracts the local self-loop contribution from each region. + +For anything else — delays, custom expressions, mixed kernels — subclass +`Model` directly and override `get_numba_coupling`. The DSL deliberately +keeps the coupling surface narrow; complex coupling deserves the imperative +escape hatch. + +## The equation language + +The `equations` string is Python source with extra constraints. It's parsed +with `ast.parse`, validated, and rewritten — not interpreted as a new +language. So if you can write it in numba, you can write it here. + +### What's allowed + +- **State variables** declared in `state_vars`. Inside equations, `x` refers + to `state[i, :]` where `i` is the row index. The DSL handles the unpacking. +- **Parameters** by name. The DSL emits `tau_e = m[np.intp(P.tau_e)]` + bindings at the top of the generated function. +- **Intermediates** introduced by assignment: `Ie = ...` is a fine local + variable. It must be assigned before it's used. +- **`coupling.`** for declared coupling vars. Rewritten to + `coupling[k, :]` where `k` is the index of that coupling var. +- **`np.`** for these whitelisted numpy functions: + `exp, log, log1p, sqrt, sin, cos, tan, tanh, sinh, cosh, abs, where, + minimum, maximum, clip, stack, empty, zeros, ones, pi`. +- Standard arithmetic operators. + +### What's required + +- Every state variable `x` must have a `d_x = ...` assignment. +- Every declared `observables` entry must be assigned. + +### Names with special meaning + +- `d_` is the time derivative of state variable ``. +- `coupling.` is the per-region coupled state for that coupling var. +- `np` is `numpy`. + +### Errors caught at build time + +The AST rewriter validates names before numba ever runs. Typos surface with +line numbers in the spec, not as cryptic numba traceback frames: + +``` +ValueError: Equation errors in model 'HopfDSL': + - Line 2: unknown identifier 'omeg'. Not a state var, parameter, + intermediate, or allowed numpy func. +``` + +## Inspecting generated code + +`dump_generated(spec)` returns the synthesized dfun source as a string. +Useful when numba complains about compilation, or just to verify the rewriter +did what you expected: + +```python +from neuronumba.simulator.models.dsl import dump_generated +print(dump_generated(hopf_spec)) +``` + +The materialized `.py` file lives under `$TMPDIR/nndsl_generated/` and is the +real file numba caches against — same spec hash → same file → same cached +compilation. Two helpers for navigating the cache: + +```python +from neuronumba.simulator.models.dsl import get_source_file, get_cache_dir + +print(get_cache_dir()) # current cache root +print(get_source_file(HopfDSL)) # path to the .py for a built model +``` + +### Configurable cache directory + +By default the cache lives under `$TMPDIR/nndsl_generated/`. On shared +compute clusters where `$TMPDIR` is per-job and gets cleaned, this defeats +numba's compilation cache between runs. Override with the +`NEURONUMBA_DSL_CACHE_DIR` environment variable to point at persistent +storage: + +```bash +export NEURONUMBA_DSL_CACHE_DIR=/scratch/$USER/nndsl_cache +``` + +The env var is read on every cache-dir access, so changing it mid-session +takes effect for subsequent `build_model` calls. + +### Numba compile errors + +If numba can't compile a generated dfun, the `RuntimeError` you get includes +the spec name, the path to the synthesized `.py` file, and the original +spec equations — open the file, find the line numba was unhappy with, fix +the spec. + +### Cleaning up stale generated files + +Every spec change writes a new `.py` file. `cleanup_cache(max_age_days=7)` +removes files older than the cutoff. Files currently held in the in-process +import cache are never removed — deleting them out from under Python could +break already-built model classes: + +```python +from neuronumba.simulator.models.dsl import cleanup_cache +removed = cleanup_cache(max_age_days=1) # nuke yesterday's iterations +``` + +## When to subclass `Model` directly + +The DSL covers the 80% case. Reach for a hand-written `Model` (or +`LinearCouplingModel`) subclass when: + +1. **You need `@njit` helper functions inside the dfun.** Zerlaut's + `get_fluct_regime_vars` and `erfc_approx` are 30-line jitted helpers that + the DSL doesn't yet support. +2. **You need imperative setup in `_init_dependant`.** Naskar/Deco2014 + compute FIC iteratively via `FICHerzog2022().compute_J(weights, g)` — that + needs Python control flow, not a single expression. You can also + *subclass a DSL-built class* and override `_init_dependant` to call + `super()` and then do your imperative setup. See + `tests/test_dsl_subclass.py`. +3. **You need an analytic Jacobian** (`get_jacobian`). Not auto-derived. +4. **You need non-scalar matrix parameters** (e.g. OrnsteinUhlenbeck's `A` + from `weights/g/tau`). The DSL only handles regional scalars right now. +5. **You need delays beyond linear/diffusive coupling.** + +These are documented as deferred items in `IMPLEMENTATION_PLAN.md` Phase 6. + +## Reference specs + +Four canonical neuronumba models live as DSL specs in +`examples/dsl_models/`: + +- [hopf_dsl.py](../../../../../examples/dsl_models/hopf_dsl.py) +- [deco2014_dsl.py](../../../../../examples/dsl_models/deco2014_dsl.py) +- [naskar2021_dsl.py](../../../../../examples/dsl_models/naskar2021_dsl.py) +- [montbrio_dsl.py](../../../../../examples/dsl_models/montbrio_dsl.py) + +These also serve as regression-test references — see +`tests/test_dsl_equivalence.py`. They're the best starting point for porting +a new model. + +## Public API + +```python +from neuronumba.simulator.models.dsl import ( + # Spec types + ModelSpec, StateVar, CouplingVar, Parameter, + # Compilation + build_model, dump_generated, + # Cache & introspection + get_source_file, get_cache_dir, cleanup_cache, +) +``` + +Everything else in the subpackage is internal — names prefixed with `_` and +helpers under `rewriter.py`, `dependents.py`, `materialize.py`, +`builder.py`. Don't depend on them; they may move between releases. From 420dfe438fb236ee56c185635450df630e3bf4ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:45 +0200 Subject: [PATCH 09/48] feat(dsl): add public api init for declarative model dsl --- .../simulator/models/dsl/__init__.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/__init__.py diff --git a/src/neuronumba/simulator/models/dsl/__init__.py b/src/neuronumba/simulator/models/dsl/__init__.py new file mode 100644 index 0000000..a2e9bef --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/__init__.py @@ -0,0 +1,28 @@ +""" +Declarative model DSL for neuronumba. + +A `ModelSpec` describes a whole-brain model in terms of state variables, +coupling variables, parameters, and equations. `build_model(spec)` compiles it +into a `Model` subclass that plugs into `Simulator` exactly like a hand-written +model. + +Public API: + ModelSpec, StateVar, CouplingVar, Parameter + build_model, dump_generated, get_source_file + get_cache_dir, cleanup_cache +""" +from .builder import build_model, dump_generated, get_source_file +from .materialize import cleanup_cache, get_cache_dir +from .spec import CouplingVar, ModelSpec, Parameter, StateVar + +__all__ = [ + "CouplingVar", + "ModelSpec", + "Parameter", + "StateVar", + "build_model", + "cleanup_cache", + "dump_generated", + "get_cache_dir", + "get_source_file", +] From 413e917836f4ba31ac99a7dd2a77e9a1a2fc52a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:49 +0200 Subject: [PATCH 10/48] feat(dsl): add builder module to generate model from spec --- .../simulator/models/dsl/builder.py | 342 ++++++++++++++++++ 1 file changed, 342 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/builder.py diff --git a/src/neuronumba/simulator/models/dsl/builder.py b/src/neuronumba/simulator/models/dsl/builder.py new file mode 100644 index 0000000..fb58c9f --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/builder.py @@ -0,0 +1,342 @@ +""" +Source generators and `build_model` — the entry point that turns a `ModelSpec` +into a numba-friendly `Model` subclass. +""" +from __future__ import annotations + +import ast +import textwrap +from typing import Any, Dict, List, Type + +import numba as nb +import numpy as np + +from neuronumba.basic.attr import Attr +from neuronumba.numba_tools.config import NUMBA_CACHE, NUMBA_FASTMATH, NUMBA_NOGIL +from neuronumba.simulator.models.model import Model, LinearCouplingModel + +from .dependents import _topo_sort_dependents +from .materialize import _materialize_module +from .rewriter import _DfunRewriter +from .spec import ModelSpec + + +def _build_dfun_source(spec: ModelSpec) -> str: + """Return executable Python source for a function `dfun(state, coupling)` + that uses captured `m` and `P` variables (added by exec env).""" + + state_index = {sv.name: i for i, sv in enumerate(spec.state_vars)} + coupling_index = {cv.name: i for i, cv in enumerate(spec.coupling_vars)} + param_names = {p.name for p in spec.parameters} + obs_names = set(spec.observables) + + src = textwrap.dedent(spec.equations).strip() + try: + tree = ast.parse(src, mode="exec") + except SyntaxError as e: + raise ValueError( + f"Could not parse equations for model '{spec.name}': {e}" + ) from e + + rewriter = _DfunRewriter( + state_index, coupling_index, param_names, obs_names + ) + new_body = [rewriter.visit(stmt) for stmt in tree.body] + + if rewriter.errors: + raise ValueError( + "Equation errors in model '{}':\n - {}".format( + spec.name, "\n - ".join(rewriter.errors) + ) + ) + + # Verify that every state var has a derivative `d_` introduced. + expected_d = {f"d_{sv.name}" for sv in spec.state_vars} + user_derivatives = { + node.targets[0].id + for node in new_body + if isinstance(node, ast.Assign) + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id in expected_d + } + missing = expected_d - user_derivatives + if missing: + raise ValueError( + f"Model '{spec.name}': missing derivative(s) {missing}. " + f"You must assign d_ for every state variable." + ) + + # Verify each declared observable has been assigned. + user_assigns = { + node.targets[0].id + for node in new_body + if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name) + } + missing_obs = obs_names - user_assigns + if missing_obs: + raise ValueError( + f"Model '{spec.name}': declared observable(s) {missing_obs} are " + f"never assigned in equations." + ) + + # Build the return tuple: + # (np.stack((d_, d_, ...)), np.stack((, , ...))) + d_stack_args = ", ".join(f"d_{sv.name}" for sv in spec.state_vars) + if spec.observables: + obs_stack_args = ", ".join(spec.observables) + ret_obs = f"np.stack(({obs_stack_args},))" + else: + ret_obs = "np.empty((1, 1))" + + body_src = "\n".join(ast.unparse(stmt) for stmt in new_body) + + # Header bindings: state vars from `state` array, params from `m`. + state_binds = "\n".join( + f" {sv.name} = state[{i}, :]" + for i, sv in enumerate(spec.state_vars) + if sv.name in rewriter.used_state + ) + param_binds = "\n".join( + f" {p} = m[np.intp(P.{p})]" + for p in sorted(rewriter.used_params) + ) + + indented_body = textwrap.indent(body_src, " ") + + # Factory function: `m` and `P` are passed in and become real closure + # variables for the inner dfun — exactly the pattern the hand-written + # neuronumba models use. + parts = [f"def make_{spec.name}_dfun(m, P):"] + parts.append(f" def {spec.name}_dfun(state, coupling):") + if state_binds: + parts.append(state_binds) + if param_binds: + parts.append(param_binds) + parts.append(indented_body) + parts.append(f" return np.stack(({d_stack_args},)), {ret_obs}") + parts.append(f" return {spec.name}_dfun") + fn_src = "\n".join(parts) + "\n" + return fn_src + + +def _build_coupling_kernel(spec: ModelSpec): + """Return a factory function that, given a configured model instance, + builds and returns a numba-jitted coupling closure.""" + + kinds = {cv.kind for cv in spec.coupling_vars} + + if kinds == {"linear"}: + def make_coupling(self): + wtg = self.g * self.weights_t.copy() + + @nb.njit(nb.f8[:, :](nb.f8[:, :]), cache=NUMBA_CACHE) + def _linear_coupling(state): + return state @ wtg + return _linear_coupling + return make_coupling + + if kinds == {"diffusive"}: + def make_coupling(self): + wt = self.weights_t.copy() + ink = wt.sum(axis=1) + g = self.g + + @nb.njit(nb.f8[:, :](nb.f8[:, :]), cache=NUMBA_CACHE) + def _diffusive_coupling(state): + r = state @ wt + return g * (r - ink * state) + return _diffusive_coupling + return make_coupling + + raise NotImplementedError( + f"Unsupported coupling kinds {sorted(kinds)}. The DSL ships 'linear' " + f"and 'diffusive'; subclass `Model` directly and override " + f"`get_numba_coupling` for anything else." + ) + + +def build_model(spec: ModelSpec) -> Type[Model]: + """Compile a ModelSpec into a `Model` subclass equivalent to a hand-written + neuronumba model. The returned class has the same external contract.""" + + # Validate that coupling vars are state vars. + state_names = [sv.name for sv in spec.state_vars] + for cv in spec.coupling_vars: + if cv.name not in state_names: + raise ValueError( + f"Coupling var '{cv.name}' is not declared as a state var." + ) + + # Build the dfun source. We embed it into a complete module that imports + # numpy (so numba can later compile from a real file path). + dfun_body_src = _build_dfun_source(spec) + full_src = ( + "# Auto-generated by nndsl. Do not edit.\n" + "import numpy as np\n" + "\n" + f"{dfun_body_src}\n" + ) + gen_module = _materialize_module(spec.name, full_src) + factory = getattr(gen_module, f"make_{spec.name}_dfun") + + # Build the class body programmatically. + cls_dict: Dict[str, Any] = {} + + # Class-level model-info attributes. + cls_dict["_state_var_names"] = [sv.name for sv in spec.state_vars] + cls_dict["_coupling_var_names"] = [cv.name for cv in spec.coupling_vars] + cls_dict["_observable_var_names"] = list(spec.observables) + bounds = {sv.name: sv.bounds for sv in spec.state_vars if sv.bounds is not None} + if bounds: + cls_dict["_state_var_bounds"] = bounds + + # Parameter Attrs. Independent params get `default`; dependents get + # `dependant=True` so neuronumba's existing machinery treats them like + # regional params that the model itself fills in during `_init_dependant`. + independent_params = [p for p in spec.parameters if not p.is_dependent] + dependent_params = [p for p in spec.parameters if p.is_dependent] + for p in independent_params: + attr_kwargs: Dict[str, Any] = { + "default": p.default, + "doc": p.doc, + } + if p.required: + attr_kwargs["required"] = True + attr_kwargs["attributes"] = ( + Model.Tag.REGIONAL if p.regional else Model.Tag.GLOBAL + ) + cls_dict[p.name] = Attr(**attr_kwargs) + for p in dependent_params: + # `dependant=True` tells HasAttr/Model to skip default-setting and + # required-checking; we'll fill these in during _init_dependant. + cls_dict[p.name] = Attr( + dependant=True, + doc=p.doc, + attributes=(Model.Tag.REGIONAL if p.regional else Model.Tag.GLOBAL), + ) + + # Topo-sort dependents once at build time and compile each formula. + sorted_deps = _topo_sort_dependents(spec) + compiled_formulas = [ + (p.name, compile(p.formula, f"", "eval")) + for p in sorted_deps + ] + + # _init_dependant: evaluate formulas in order, on the instance. Runs at + # configure() time, i.e. on construction and any subsequent configure() call, + # which is the existing protocol for "parameter was modified" in neuronumba. + # + # Mutable container that will hold the generated class once `type()` returns + # it below. Captured by the closure so `super()` always resolves against the + # DSL-generated class, not against `type(self)`. Without this indirection, + # subclassing the generated class would re-enter this method via super() + # forever (because `type(self)` would be the user's subclass). + _class_cell: List[Type[Model]] = [None] + + def _init_dependant(self): + # Run the base-class hook first (sets n_rois, weights_t for + # LinearCouplingModel, etc.). Without this, `self.weights` exists but + # things like `self.weights_t` don't yet. + super(_class_cell[0], self)._init_dependant() + # Evaluation env: numpy + every parameter currently on self (independent + # ones from the user, plus dependents already computed earlier in the + # topo order). + env: Dict[str, Any] = {"np": np} + # Snapshot all params that exist so far. + for p in spec.parameters: + val = getattr(self, p.name, None) + if val is not None: + env[p.name] = val + for name, code_obj in compiled_formulas: + try: + value = eval(code_obj, env) + except Exception as e: + raise RuntimeError( + f"Error evaluating dependent parameter '{name}' in model " + f"'{spec.name}': {e}" + ) from e + setattr(self, name, value) + env[name] = value + cls_dict["_init_dependant"] = _init_dependant + + # `g` is pulled in by LinearCouplingModel; we don't redeclare it. + + initial_overrides = dict(spec.initial_state_overrides) + state_initials = [sv.initial for sv in spec.state_vars] + + def initial_state(self, n_rois: int) -> np.ndarray: + out = np.empty((len(state_initials), n_rois)) + for i, init_val in enumerate(state_initials): + out[i] = init_val + for name, val in initial_overrides.items(): + if name in cls_dict["_state_var_names"]: + out[cls_dict["_state_var_names"].index(name)] = val + return out + cls_dict["initial_state"] = initial_state + + # Capture spec metadata in the closure for error reporting. + _spec_name = spec.name + _spec_equations = spec.equations + _source_file = gen_module.__file__ + + def get_numba_dfun(self): + m_local = self.m.copy() + P_local = self.P + py_dfun = factory(m_local, P_local) + try: + jitted = nb.njit( + nb.types.UniTuple(nb.f8[:, :], 2)(nb.f8[:, :], nb.f8[:, :]), + cache=NUMBA_CACHE, fastmath=NUMBA_FASTMATH, nogil=NUMBA_NOGIL, + )(py_dfun) + except nb.errors.NumbaError as e: + raise RuntimeError( + f"Could not compile dfun for DSL model '{_spec_name}'.\n" + f"Generated source: {_source_file}\n" + f"Spec equations:\n{_spec_equations}\n" + f"Original numba error: {e}" + ) from e + return jitted + cls_dict["get_numba_dfun"] = get_numba_dfun + + coupling_factory = _build_coupling_kernel(spec) + cls_dict["get_numba_coupling"] = coupling_factory + + # Choose base class. LinearCouplingModel only adds `g` + `weights_t`. We + # always inherit from it: even diffusive Hopf needs `g` and `weights_t`, + # and for 'linear' it gives us the default coupling (which we override + # only when we want — but we override always for clarity here). + base = LinearCouplingModel + + cls_dict["__doc__"] = ( + f"DSL-generated whole-brain model: {spec.name}\n\n" + f"State: {[sv.name for sv in spec.state_vars]}\n" + f"Coupling: {[(cv.name, cv.kind) for cv in spec.coupling_vars]}\n" + f"Observables: {spec.observables}\n" + f"Parameters: {[p.name for p in spec.parameters]}\n" + ) + cls_dict["__module__"] = "neuronumba.simulator.models.dsl.generated" + + cls = type(spec.name, (base,), cls_dict) + cls.__nndsl_source_file__ = gen_module.__file__ + _class_cell[0] = cls + return cls + + +def dump_generated(spec: ModelSpec) -> str: + """Return the generated dfun source for inspection.""" + return _build_dfun_source(spec) + + +def get_source_file(cls: Type[Model]) -> str: + """Return the on-disk path to the generated dfun source for a built model. + + Useful when numba reports a compile error pointing at the synthesized + file: open the path, read the offending line, fix the spec. + """ + try: + return cls.__nndsl_source_file__ + except AttributeError: + raise TypeError( + f"{cls!r} does not look like a DSL-built model " + f"(no `__nndsl_source_file__` attribute)." + ) from None From e2d31cc94dc1d5883d2fcab15efa25422e958ff4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:53 +0200 Subject: [PATCH 11/48] feat(dsl): add dependent parameter dag analysis and topo sort --- .../simulator/models/dsl/dependents.py | 90 +++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/dependents.py diff --git a/src/neuronumba/simulator/models/dsl/dependents.py b/src/neuronumba/simulator/models/dsl/dependents.py new file mode 100644 index 0000000..01ce3e3 --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/dependents.py @@ -0,0 +1,90 @@ +""" +Dependent-parameter analysis. + +Parses each formula's AST, builds a DAG of dependencies, topo-sorts, and +returns dependent parameters in evaluation order. Cycles and self-references +raise ValueError. +""" +from __future__ import annotations + +import ast +from typing import Dict, List + +from .rewriter import ALLOWED_NP_FUNCS +from .spec import ModelSpec, Parameter + + +def _formula_deps(name: str, src: str, allowed: set) -> List[str]: + """Return the parameter names referenced in `src`. Validate identifiers.""" + try: + tree = ast.parse(src, mode="eval") + except SyntaxError as e: + raise ValueError(f"Parameter '{name}' formula is not valid Python: {e}") + + deps: List[str] = [] + seen = set() + + class V(ast.NodeVisitor): + def visit_Name(self, node): + n = node.id + if n == "np": + return + if n in ALLOWED_NP_FUNCS: + return + if n not in allowed: + raise ValueError( + f"Parameter '{name}': formula references unknown identifier " + f"'{n}' (line {node.lineno}). Must be another declared " + f"parameter or numpy." + ) + if n not in seen: + deps.append(n) + seen.add(n) + + def visit_Attribute(self, node): + # Allow `np.something`, recurse into base. + self.generic_visit(node) + + V().visit(tree) + if name in deps: + raise ValueError(f"Parameter '{name}' formula references itself.") + return deps + + +def _topo_sort_dependents(spec: ModelSpec) -> List[Parameter]: + """Return dependent parameters in evaluation order. Raise on cycles.""" + by_name = {p.name: p for p in spec.parameters} + all_names = set(by_name.keys()) + + deps_of: Dict[str, List[str]] = {} + for p in spec.parameters: + if p.is_dependent: + deps_of[p.name] = _formula_deps(p.name, p.formula, all_names) + + # Kahn's algorithm restricted to the dependent subgraph. + dependent_names = set(deps_of.keys()) + indeg: Dict[str, int] = {n: 0 for n in dependent_names} + rev: Dict[str, List[str]] = {n: [] for n in dependent_names} + for n, ds in deps_of.items(): + for d in ds: + if d in dependent_names: + indeg[n] += 1 + rev[d].append(n) + + ready = [n for n, k in indeg.items() if k == 0] + order: List[str] = [] + while ready: + n = ready.pop() + order.append(n) + for m in rev[n]: + indeg[m] -= 1 + if indeg[m] == 0: + ready.append(m) + + if len(order) != len(dependent_names): + unresolved = dependent_names - set(order) + raise ValueError( + f"Cycle detected among dependent parameters: {sorted(unresolved)}" + ) + + return [by_name[n] for n in order] From e757c63d8a3745b9907694110cb96f7475eb5f50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:38:57 +0200 Subject: [PATCH 12/48] feat(dsl): add on-disk module materialization for numba cache --- .../simulator/models/dsl/materialize.py | 100 ++++++++++++++++++ 1 file changed, 100 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/materialize.py diff --git a/src/neuronumba/simulator/models/dsl/materialize.py b/src/neuronumba/simulator/models/dsl/materialize.py new file mode 100644 index 0000000..1850de4 --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/materialize.py @@ -0,0 +1,100 @@ +""" +Materialize generated source to a real on-disk module. + +Numba's `cache=True` requires a real file path (the cache key is the source +file path). Synthesized filenames via `linecache` do not work. Writing to a +real `.py` file under `tempfile.gettempdir()` (or a user-specified directory) +and importing it as a module gives us: + - working numba cache, + - real tracebacks pointing into the file, + - debuggable code you can inspect with `inspect.getsource`. + +Same approach SymPy's `lambdify` and TVB's RateML take. +""" +from __future__ import annotations + +import hashlib +import importlib.util +import os +import sys +import tempfile +import time +from typing import Any, Dict + + +_GENERATED_MODULES: Dict[str, Any] = {} # cache by source-hash + + +def get_cache_dir() -> str: + """Return the directory where generated module files are written. + + Defaults to ``$TMPDIR/nndsl_generated/``. Override with the + ``NEURONUMBA_DSL_CACHE_DIR`` environment variable — useful on shared + compute clusters where ``$TMPDIR`` is per-job and gets cleaned, defeating + numba's compile cache between runs. + + Resolved fresh on each call so tests (and users who change the env var + mid-session) see the current value. + """ + return os.environ.get( + "NEURONUMBA_DSL_CACHE_DIR", + os.path.join(tempfile.gettempdir(), "nndsl_generated"), + ) + + +def cleanup_cache(max_age_days: float = 7.0) -> int: + """Delete generated ``.py`` files older than ``max_age_days``. + + Useful when iterating in Jupyter: every spec change creates a new module + file, and stale ones accumulate. Returns the number of files removed. + + Files currently held in the in-process import cache are NOT removed even + if they're old, because deleting them out from under Python could break + further use of already-built model classes. + """ + cache_dir = get_cache_dir() + if not os.path.isdir(cache_dir): + return 0 + + cutoff = time.time() - max_age_days * 86400.0 + in_use = {getattr(mod, "__file__", None) for mod in _GENERATED_MODULES.values()} + removed = 0 + for entry in os.listdir(cache_dir): + if not entry.endswith(".py"): + continue + path = os.path.join(cache_dir, entry) + if path in in_use: + continue + try: + if os.path.getmtime(path) < cutoff: + os.remove(path) + removed += 1 + except OSError: + # File may have been removed by another process between listdir + # and stat — fine to ignore. + pass + return removed + + +def _materialize_module(spec_name: str, src: str) -> Any: + """Write `src` (a complete Python file) to disk and import it as a module. + Returns the imported module. Same source -> cached module reuse.""" + key = hashlib.md5(src.encode("utf-8")).hexdigest()[:12] + if key in _GENERATED_MODULES: + return _GENERATED_MODULES[key] + + cache_dir = get_cache_dir() + os.makedirs(cache_dir, exist_ok=True) + + mod_name = f"nndsl_gen_{spec_name}_{key}" + fname = os.path.join(cache_dir, f"{mod_name}.py") + if not os.path.exists(fname): + with open(fname, "w", encoding="utf-8") as f: + f.write(src) + + spec = importlib.util.spec_from_file_location(mod_name, fname) + mod = importlib.util.module_from_spec(spec) + sys.modules[mod_name] = mod + spec.loader.exec_module(mod) + _GENERATED_MODULES[key] = mod + return mod From a7bbdbcb77dc3b8aecdea8d9de5395b59827bce8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:00 +0200 Subject: [PATCH 13/48] feat(dsl): add ast rewriter to transform equations to numba form --- .../simulator/models/dsl/rewriter.py | 124 ++++++++++++++++++ 1 file changed, 124 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/rewriter.py diff --git a/src/neuronumba/simulator/models/dsl/rewriter.py b/src/neuronumba/simulator/models/dsl/rewriter.py new file mode 100644 index 0000000..cf33853 --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/rewriter.py @@ -0,0 +1,124 @@ +""" +AST rewriter for DSL equation bodies. + +Transforms identifiers in user-supplied equations to the +`m[np.intp(P.x)]` / `state[i, :]` / `coupling[k, :]` form used by the +hand-written neuronumba dfuns, and validates that every identifier is +declared (state var, parameter, intermediate, or whitelisted numpy func). +""" +from __future__ import annotations + +import ast +from typing import Dict, List + + +# Numpy functions allowed inside dfun bodies. Using a whitelist is the cheapest +# way to catch typos and to refuse names that numba won't compile anyway. +ALLOWED_NP_FUNCS = { + "exp", "log", "log1p", "sqrt", "sin", "cos", "tan", + "tanh", "sinh", "cosh", "abs", "where", + "minimum", "maximum", "clip", + "stack", "empty", "zeros", "ones", + "pi", +} + + +class _DfunRewriter(ast.NodeTransformer): + """Walks the user's equation AST and rewrites Names / Attributes.""" + + def __init__( + self, + state_index: Dict[str, int], + coupling_index: Dict[str, int], + param_names: set, + observable_names: set, + ): + self.state_index = state_index + self.coupling_index = coupling_index + self.param_names = param_names + self.observable_names = observable_names + self.intermediates: set = set() # names introduced by user assignments + self.used_state: set = set() + self.used_params: set = set() + self.errors: List[str] = [] + + # --- assignment LHS handling ------------------------------------------ + def visit_Assign(self, node: ast.Assign): + # Targets: only simple names allowed (e.g. `Ie = ...` or `dSe = ...`). + if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name): + self.errors.append( + f"Line {node.lineno}: only single-name assignments are supported" + ) + return node + + target = node.targets[0].id + # If the target is `d_`, mark it as a derivative (no rewriting needed — + # we collect derivatives at emit time and assemble the return value ourselves). + # Anything else is an intermediate (or an observable). + if not ( + target.startswith("d_") + and target[2:] in self.state_index + ): + self.intermediates.add(target) + # Recurse into the RHS only. + node.value = self.visit(node.value) + return node + + # --- name resolution -------------------------------------------------- + def visit_Name(self, node: ast.Name): + n = node.id + + if n in self.state_index: + # Don't rewrite. State vars get bound as locals (`Se = state[0,:]`). + self.used_state.add(n) + return node + + if n in self.param_names: + # Don't rewrite parameters. They get bound as locals at the top of + # the generated function (e.g. `tau_e = m[np.intp(P.tau_e)]`). + self.used_params.add(n) + return node + + # Allowed: intermediates introduced earlier, observables, the special + # tokens 'state', 'coupling', 'm', 'P', 'np', and known numpy functions. + if n in self.intermediates: + return node + if n in {"state", "coupling", "m", "P", "np"}: + return node + if n in ALLOWED_NP_FUNCS: + return node + + # Anything else is an undeclared identifier — error. + self.errors.append( + f"Line {node.lineno}: unknown identifier '{n}'. " + f"Not a state var, parameter, intermediate, or allowed numpy func." + ) + return node + + # --- coupling. rewriting ---------------------------------------- + def visit_Attribute(self, node: ast.Attribute): + # Match `coupling.S_e` etc. -> `coupling[k, :]` + if isinstance(node.value, ast.Name) and node.value.id == "coupling": + cname = node.attr + if cname not in self.coupling_index: + self.errors.append( + f"Line {node.lineno}: coupling.{cname} but '{cname}' is not " + f"declared as a coupling variable." + ) + return node + k = self.coupling_index[cname] + return ast.copy_location( + ast.Subscript( + value=ast.Name(id="coupling", ctx=ast.Load()), + slice=ast.Tuple( + elts=[ast.Constant(value=k), + ast.Slice(lower=None, upper=None, step=None)], + ctx=ast.Load(), + ), + ctx=node.ctx, + ), + node, + ) + # Otherwise allow (e.g. `np.exp`) — but recurse into the value side. + node.value = self.visit(node.value) + return node From e83ecd14485202e0aae477f0f9b4f20143c9c908 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:04 +0200 Subject: [PATCH 14/48] feat(dsl): add spec dataclasses for model declaration --- src/neuronumba/simulator/models/dsl/spec.py | 99 +++++++++++++++++++++ 1 file changed, 99 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/spec.py diff --git a/src/neuronumba/simulator/models/dsl/spec.py b/src/neuronumba/simulator/models/dsl/spec.py new file mode 100644 index 0000000..1e9a8a9 --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/spec.py @@ -0,0 +1,99 @@ +""" +Spec dataclasses for the neuronumba DSL. + +These are the user-facing declaration objects: a `ModelSpec` describes a whole +model in terms of its `StateVar`s, `CouplingVar`s, `Parameter`s, and an +equations string. They are pure data — no compilation logic lives here. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Tuple + + +@dataclass +class StateVar: + name: str + initial: float = 0.0 + bounds: Optional[Tuple[float, float]] = None + + +@dataclass +class CouplingVar: + """A state variable that participates in inter-region coupling. + + kind: + 'linear' — coupling = g * (W^T @ S) (standard SC coupling) + 'diffusive' — coupling = g * (W^T @ S - sum(W^T,1) * S) (Hopf-style) + + For coupling shapes outside this set, subclass `Model` directly and override + `get_numba_coupling`. + """ + name: str + kind: str = "linear" + + +@dataclass +class Parameter: + """A model parameter. + + A parameter is *independent* if it carries a `default` (or is `required`) and + no `formula`. Independent params are user-settable and packed into `self.m`. + + A parameter is *dependent* if it has a `formula`: a Python expression that + may reference other parameters (independent or dependent, as long as no + cycles). Dependents are NOT user-settable; they're computed at + `configure()` time and re-computed every time the user calls `configure()` + (i.e. every time params are modified — which is already the protocol the + existing neuronumba models follow). + + Examples: + Parameter("J_ee", default=10.0) + Parameter("g_ee", default=2.5) + Parameter("a_e", default=0.25) + Parameter("J_N_ee", formula="J_ee + g_ee * np.log(a_e)") # Montbrio-style + + Formulas may use any of: numpy (as `np`), the standard math operators, and + references to other parameter names. They evaluate in NumPy land (not numba) + so anything numpy supports works, including matrix ops. + """ + name: str + default: Optional[float] = None + regional: bool = True + required: bool = False + doc: str = "" + formula: Optional[str] = None + + def __post_init__(self): + if self.formula is not None: + if self.default is not None: + raise ValueError( + f"Parameter '{self.name}' has both `default` and `formula`; " + f"a dependent parameter cannot have a default value." + ) + if self.required: + raise ValueError( + f"Parameter '{self.name}' has `formula` and is `required`; " + f"dependent parameters are computed, not user-supplied." + ) + else: + if self.default is None and not self.required: + raise ValueError( + f"Parameter '{self.name}' needs either `default`, " + f"`required=True`, or a `formula`." + ) + + @property + def is_dependent(self) -> bool: + return self.formula is not None + + +@dataclass +class ModelSpec: + name: str + state_vars: List[StateVar] + coupling_vars: List[CouplingVar] + observables: List[str] + parameters: List[Parameter] + equations: str # multiline; assignments of intermediates and `d_` rates + initial_state_overrides: Dict[str, float] = field(default_factory=dict) From ca7dc90b298aba304a6720717439cece8702f3bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:07 +0200 Subject: [PATCH 15/48] test(tests): add shared fixtures for dsl test suite --- tests/conftest.py | 97 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 tests/conftest.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..489dca1 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,97 @@ +"""Shared fixtures for the DSL test suite. + +Spec definitions live in `examples/dsl_models/` (single source of truth for +both user-facing examples and equivalence-test references). This file just +re-exports them as fixtures. + +Built classes are session-scoped so numba's compile cache amortizes across +test modules. +""" +import os +import sys + +import numpy as np +import pytest + +# Make the `examples/dsl_models` package importable. `examples/` itself is not +# a Python package (it's a flat collection of demo scripts); we add it to +# sys.path so `import dsl_models` resolves the sub-package. +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir)) +sys.path.insert(0, os.path.join(_REPO_ROOT, "examples")) + +from dsl_models import ( # noqa: E402 + hopf_spec as _hopf_spec, + deco_spec as _deco_spec, + naskar_spec as _naskar_spec, + montbrio_spec as _montbrio_spec, + HopfDSL as _HopfDSL, + Deco2014DSL as _Deco2014DSL, + Naskar2021DSL as _Naskar2021DSL, + MontbrioDSL as _MontbrioDSL, +) + + +def _make_W(n=8, seed=0): + """Symmetric, zero-diagonal connectivity normalized to max 1.0.""" + rng = np.random.default_rng(seed) + W = rng.random((n, n)) + W = (W + W.T) / 2 + np.fill_diagonal(W, 0) + W /= W.max() + return W + + +@pytest.fixture(scope="session") +def n_rois(): + return 8 + + +@pytest.fixture(scope="session") +def weights(n_rois): + return _make_W(n=n_rois) + + +# ----- canonical model specs (re-exported from examples/dsl_models) ---------- + +@pytest.fixture(scope="session") +def hopf_spec(): + return _hopf_spec + + +@pytest.fixture(scope="session") +def deco_spec(): + return _deco_spec + + +@pytest.fixture(scope="session") +def naskar_spec(): + return _naskar_spec + + +@pytest.fixture(scope="session") +def montbrio_spec(): + return _montbrio_spec + + +# ----- materialized model classes -------------------------------------------- +# Re-export the built classes from `examples/dsl_models/` directly. Tests +# exercise the exact same class objects a user would import. + +@pytest.fixture(scope="session") +def hopf_dsl_cls(): + return _HopfDSL + + +@pytest.fixture(scope="session") +def deco_dsl_cls(): + return _Deco2014DSL + + +@pytest.fixture(scope="session") +def naskar_dsl_cls(): + return _Naskar2021DSL + + +@pytest.fixture(scope="session") +def montbrio_dsl_cls(): + return _MontbrioDSL From cda71e94764466d5e722d49d8364619508f500ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:12 +0200 Subject: [PATCH 16/48] test(tests): add spec and equation validation unit tests --- tests/test_dsl_basics.py | 72 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 72 insertions(+) create mode 100644 tests/test_dsl_basics.py diff --git a/tests/test_dsl_basics.py b/tests/test_dsl_basics.py new file mode 100644 index 0000000..87d0ea5 --- /dev/null +++ b/tests/test_dsl_basics.py @@ -0,0 +1,72 @@ +"""Spec-level validation: Parameter() rules and dfun-equation rules.""" +import pytest + +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +# ----- Parameter validation rules (Parameter.__post_init__) ------------------ + +def test_param_default_and_formula_rejected(): + with pytest.raises(ValueError, match="both `default` and `formula`"): + Parameter("x", default=1.0, formula="y + 1") + + +def test_param_required_and_formula_rejected(): + with pytest.raises(ValueError, match="`formula` and is `required`"): + Parameter("x", required=True, formula="y + 1") + + +def test_param_neither_default_nor_required_nor_formula(): + with pytest.raises(ValueError, match="needs either `default`"): + Parameter("x") + + +# ----- AST rewriter validation ----------------------------------------------- + +def _minimal_spec(equations, observables=()): + return ModelSpec( + name="MinModel", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=list(observables), + parameters=[Parameter("a", default=1.0)], + equations=equations, + ) + + +def test_dfun_unknown_identifier_raises(): + spec = _minimal_spec("d_x = -x * a + bogus_thing + coupling.x") + with pytest.raises(ValueError, match="unknown identifier 'bogus_thing'"): + build_model(spec) + + +def test_dfun_missing_derivative_raises(): + spec = _minimal_spec("foo = a * 2") + with pytest.raises(ValueError, match="missing derivative"): + build_model(spec) + + +def test_dfun_missing_observable_raises(): + spec = _minimal_spec( + "d_x = -x * a + coupling.x", + observables=["r"], + ) + with pytest.raises(ValueError, match="declared observable.*never assigned"): + build_model(spec) + + +# ----- Spec-level structural validation ------------------------------------- + +def test_coupling_var_must_be_state_var(): + spec = ModelSpec( + name="BadCoupling", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("y", kind="linear")], # 'y' not in state_vars + observables=[], + parameters=[Parameter("a", default=1.0)], + equations="d_x = -x * a + coupling.y", + ) + with pytest.raises(ValueError, match="not declared as a state var"): + build_model(spec) From 4b4ac1af151c5d8617a773b8ca103c7fdfa011d1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:16 +0200 Subject: [PATCH 17/48] test(tests): add dependent parameter dag and evaluation tests --- tests/test_dsl_dependents.py | 118 +++++++++++++++++++++++++++++++++++ 1 file changed, 118 insertions(+) create mode 100644 tests/test_dsl_dependents.py diff --git a/tests/test_dsl_dependents.py b/tests/test_dsl_dependents.py new file mode 100644 index 0000000..f3c3e2c --- /dev/null +++ b/tests/test_dsl_dependents.py @@ -0,0 +1,118 @@ +"""Dependent-parameter analysis and runtime evaluation.""" +import numpy as np +import pytest + +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +def _spec_with_params(params, name="DepModel", equations="d_x = -x + coupling.x"): + return ModelSpec( + name=name, + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=params, + equations=equations, + ) + + +def test_dependents_simple_chain(): + spec = _spec_with_params( + [ + Parameter("a", default=2.0), + Parameter("b", default=3.0), + Parameter("c", formula="a + b"), + ], + name="ChainDep", + equations="d_x = -x / c + coupling.x", + ) + Cls = build_model(spec) + n = 4 + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + assert np.allclose(m.c, 5.0) + assert np.allclose(m.m[int(m.P.c)], 5.0) + + +def test_dependents_multilevel_topo_sort(): + spec = _spec_with_params( + [ + Parameter("a", default=2.0), + Parameter("b", default=3.0), + Parameter("c", default=4.0), + Parameter("d", formula="a + b"), + Parameter("e", formula="c * d"), + Parameter("f", formula="np.exp(0) * e"), + ], + name="MultiDep", + equations="d_x = -x + f * coupling.x", + ) + Cls = build_model(spec) + m = Cls(g=0.0).configure(weights=np.zeros((3, 3))) + assert np.allclose(m.d, 5.0) + assert np.allclose(m.e, 20.0) + assert np.allclose(m.f, 20.0) + + +def test_dependents_montbrio_pattern(): + """The recurrent Montbrio pattern: J_N_ee = J_ee + g_ee * log(a_e).""" + spec = _spec_with_params( + [ + Parameter("J_ee", default=10.0), + Parameter("g_ee", default=2.5), + Parameter("a_e", default=0.25), + Parameter("J_N_ee", formula="J_ee + g_ee * np.log(a_e)"), + ], + name="MontbrioDep", + equations="d_x = -x * J_N_ee + coupling.x", + ) + Cls = build_model(spec) + m = Cls(g=0.0).configure(weights=np.zeros((3, 3))) + expected = 10.0 + 2.5 * np.log(0.25) + assert np.allclose(m.J_N_ee, expected) + + +def test_dependents_cycle_detected(): + spec = _spec_with_params( + [ + Parameter("a", formula="b + 1"), + Parameter("b", formula="a + 1"), + ], + name="CycleDep", + equations="d_x = -x * a + coupling.x", + ) + with pytest.raises(ValueError, match="Cycle detected"): + build_model(spec) + + +def test_dependents_self_reference_detected(): + spec = _spec_with_params( + [Parameter("a", formula="a + 1")], + name="SelfDep", + equations="d_x = -x * a + coupling.x", + ) + with pytest.raises(ValueError, match="references itself"): + build_model(spec) + + +def test_dependents_recompute_on_configure(): + """Modifying an independent param + re-configure() recomputes dependents.""" + spec = _spec_with_params( + [ + Parameter("a", default=2.0), + Parameter("b", formula="a * 10"), + ], + name="RecomputeDep", + equations="d_x = -x / b + coupling.x", + ) + Cls = build_model(spec) + n = 3 + W = np.zeros((n, n)) + m = Cls(g=0.0).configure(weights=W) + assert np.allclose(m.b, 20.0) + + m.a = 5.0 + m.configure(weights=W) + assert np.allclose(m.b, 50.0) + assert np.allclose(m.m[int(m.P.b)], 50.0) From ae65e4a50fc9433ea39e159cb95f0e4824f4bd39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:22 +0200 Subject: [PATCH 18/48] test(tests): add dfun and coupling equivalence tests for dsl models --- tests/test_dsl_equivalence.py | 68 +++++++++++++++++++++++++++++++++++ 1 file changed, 68 insertions(+) create mode 100644 tests/test_dsl_equivalence.py diff --git a/tests/test_dsl_equivalence.py b/tests/test_dsl_equivalence.py new file mode 100644 index 0000000..d4671b1 --- /dev/null +++ b/tests/test_dsl_equivalence.py @@ -0,0 +1,68 @@ +"""DSL-built models produce identical outputs to hand-written ones. + +For each canonical model we compare: + - the coupling kernel on a random coupling-state batch, + - the dfun on a random (state, coupling) pair. +""" +import numpy as np + +from neuronumba.simulator.models import Hopf, Deco2014, Naskar2021, Montbrio + + +def _compare_dfun(m_ref, m_dsl, n_rois, rng, atol=1e-12): + state = rng.standard_normal((m_ref.n_state_vars, n_rois)) * 0.1 + coupling = rng.standard_normal((len(m_ref.c_vars), n_rois)) * 0.05 + + f_ref = m_ref.get_numba_dfun() + f_dsl = m_dsl.get_numba_dfun() + ds_ref, obs_ref = f_ref(state.copy(), coupling.copy()) + ds_dsl, obs_dsl = f_dsl(state.copy(), coupling.copy()) + + assert np.allclose(ds_ref, ds_dsl, atol=atol, rtol=atol) + if obs_ref.size > 1 and obs_dsl.size > 1: + assert obs_ref.shape == obs_dsl.shape + assert np.allclose(obs_ref, obs_dsl, atol=atol, rtol=atol) + + +def _compare_coupling(m_ref, m_dsl, n_rois, rng, atol=1e-12): + f_ref = m_ref.get_numba_coupling() + f_dsl = m_dsl.get_numba_coupling() + n_c = len(m_ref.c_vars) + s = rng.standard_normal((n_c, n_rois)) * 0.1 + a = f_ref(s.copy()) + b = f_dsl(s.copy()) + assert np.allclose(a, b, atol=atol, rtol=atol) + + +def test_equivalence_hopf(weights, n_rois, hopf_dsl_cls): + rng = np.random.default_rng(42) + m_ref = Hopf(g=0.5).configure(weights=weights) + m_dsl = hopf_dsl_cls(g=0.5).configure(weights=weights) + _compare_coupling(m_ref, m_dsl, n_rois, rng) + _compare_dfun(m_ref, m_dsl, n_rois, rng) + + +def test_equivalence_deco2014(weights, n_rois, deco_dsl_cls): + rng = np.random.default_rng(42) + m_ref = Deco2014(g=0.5, auto_fic=False).configure(weights=weights) + m_dsl = deco_dsl_cls(g=0.5).configure(weights=weights) + _compare_coupling(m_ref, m_dsl, n_rois, rng) + _compare_dfun(m_ref, m_dsl, n_rois, rng) + + +def test_equivalence_naskar2021(weights, n_rois, naskar_dsl_cls): + rng = np.random.default_rng(42) + m_ref = Naskar2021(g=0.5).configure(weights=weights) + m_dsl = naskar_dsl_cls(g=0.5).configure(weights=weights) + _compare_coupling(m_ref, m_dsl, n_rois, rng) + # Naskar's sigmoid has no eps-fix; tiny denominator differences can + # propagate. 1e-10 still proves bit-equivalent math. + _compare_dfun(m_ref, m_dsl, n_rois, rng, atol=1e-10) + + +def test_equivalence_montbrio(weights, n_rois, montbrio_dsl_cls): + rng = np.random.default_rng(42) + m_ref = Montbrio(g=0.5, auto_fic=False).configure(weights=weights) + m_dsl = montbrio_dsl_cls(g=0.5).configure(weights=weights) + _compare_coupling(m_ref, m_dsl, n_rois, rng) + _compare_dfun(m_ref, m_dsl, n_rois, rng, atol=1e-10) From 0bf4043a0661d4b09115003d50a3d53eb93ced77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:26 +0200 Subject: [PATCH 19/48] test(tests): add ux polish and cache dir configuration tests --- tests/test_dsl_polish.py | 96 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 tests/test_dsl_polish.py diff --git a/tests/test_dsl_polish.py b/tests/test_dsl_polish.py new file mode 100644 index 0000000..bc69f42 --- /dev/null +++ b/tests/test_dsl_polish.py @@ -0,0 +1,96 @@ +"""Tests for the Phase 5 UX-polish helpers and the configurable cache dir. + +Covers: + - get_cache_dir() honors NEURONUMBA_DSL_CACHE_DIR (Phase 0.3, deferred here) + - get_source_file() returns a real, readable on-disk file for built models + - cleanup_cache() removes old files but leaves in-use and recent files alone + - The numba-error wrapping in get_numba_dfun is wired up (smoke check) +""" +import os +import time + +import pytest + +from neuronumba.simulator.models.dsl import ( + cleanup_cache, get_cache_dir, get_source_file, +) + + +def test_get_cache_dir_env_var_respected(monkeypatch, tmp_path): + monkeypatch.setenv("NEURONUMBA_DSL_CACHE_DIR", str(tmp_path)) + assert get_cache_dir() == str(tmp_path) + + +def test_get_cache_dir_default_when_unset(monkeypatch): + monkeypatch.delenv("NEURONUMBA_DSL_CACHE_DIR", raising=False) + cache = get_cache_dir() + # Default lives under the system tempdir. + import tempfile + assert cache.startswith(tempfile.gettempdir()) + assert cache.endswith("nndsl_generated") + + +def test_get_source_file_returns_existing_path(hopf_dsl_cls): + path = get_source_file(hopf_dsl_cls) + assert os.path.isfile(path) + with open(path, encoding="utf-8") as f: + src = f.read() + # Sanity: the dfun factory we expect for this spec must be present. + assert "make_HopfDSL_dfun" in src + + +def test_get_source_file_rejects_non_dsl_class(): + from neuronumba.simulator.models import Hopf + with pytest.raises(TypeError, match="DSL-built model"): + get_source_file(Hopf) + + +def test_cleanup_cache_removes_old_only(monkeypatch, tmp_path): + """cleanup_cache deletes stale files but spares recent ones.""" + monkeypatch.setenv("NEURONUMBA_DSL_CACHE_DIR", str(tmp_path)) + + old = tmp_path / "old_module.py" + fresh = tmp_path / "fresh_module.py" + not_a_py = tmp_path / "stray.txt" + old.write_text("# old") + fresh.write_text("# fresh") + not_a_py.write_text("ignored") + + # Backdate `old` by 30 days; leave `fresh` at now. + old_time = time.time() - 30 * 86400 + os.utime(str(old), (old_time, old_time)) + + removed = cleanup_cache(max_age_days=7.0) + assert removed == 1 + assert not old.exists() + assert fresh.exists() + assert not_a_py.exists() # only .py files are touched + + +def test_numba_error_wrapping_smoke(hopf_dsl_cls, weights): + """The error-wrapping path is not exercised in the happy case, but we + verify the wrapper is in place by monkeypatching numba to raise.""" + import numba as nb + + m = hopf_dsl_cls(g=0.5).configure(weights=weights) + + # Monkeypatch nb.njit to raise the kind of error users might hit. + # (We can't easily produce a real TypingError with our whitelisted + # equations; this confirms the wrapper exists and reports the right + # context.) + real_njit = nb.njit + def _broken_njit(*args, **kwargs): + def _decorator(fn): + raise nb.errors.TypingError("synthetic typing failure") + return _decorator + nb.njit = _broken_njit + try: + with pytest.raises(RuntimeError) as excinfo: + m.get_numba_dfun() + finally: + nb.njit = real_njit + + msg = str(excinfo.value) + assert "HopfDSL" in msg # spec name surfaced + assert "Generated source:" in msg # path surfaced + assert "synthetic typing failure" in msg # original error preserved From 986c43c9540ff9aa27e27fb17a4c558df52b6233 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:29 +0200 Subject: [PATCH 20/48] test(tests): add end-to-end simulator trajectory equivalence tests --- tests/test_dsl_simulator.py | 39 +++++++++++++++++++++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 tests/test_dsl_simulator.py diff --git a/tests/test_dsl_simulator.py b/tests/test_dsl_simulator.py new file mode 100644 index 0000000..dc5e917 --- /dev/null +++ b/tests/test_dsl_simulator.py @@ -0,0 +1,39 @@ +"""End-to-end equivalence: simulator trajectories match hand-written models. + +Both runs use `EulerDeterministic` so the only source of variation is the dfun +and coupling kernel — both already proved bit-equal in test_dsl_equivalence. +We seed numpy globally before each `simulate_nodelay` call because that helper +draws random tract lengths internally; HistoryNoDelays ignores lengths but we +still seed for reproducibility. +""" +import numpy as np + +from neuronumba.simulator.simulator import simulate_nodelay +from neuronumba.simulator.integrators.euler import EulerDeterministic +from neuronumba.simulator.models import Hopf, Deco2014 + + +def _run(model_cls, weights, obs_var, *, dt=0.1, t_max=200.0, **model_kwargs): + np.random.seed(0) + model = model_cls(**model_kwargs) + integrator = EulerDeterministic(dt=dt) + return simulate_nodelay( + model, integrator, weights, obs_var, + sampling_period=1.0, + t_max_neuronal=t_max, + t_warmup=0.0, + ) + + +def test_simulator_end_to_end_hopf_deterministic(weights, hopf_dsl_cls): + ref = _run(Hopf, weights, obs_var="x", g=0.5) + dsl = _run(hopf_dsl_cls, weights, obs_var="x", g=0.5) + assert ref.shape == dsl.shape + assert np.allclose(ref, dsl, atol=1e-12, rtol=1e-12) + + +def test_simulator_end_to_end_deco2014_deterministic(weights, deco_dsl_cls): + ref = _run(Deco2014, weights, obs_var="S_e", g=0.5, auto_fic=False) + dsl = _run(deco_dsl_cls, weights, obs_var="S_e", g=0.5) + assert ref.shape == dsl.shape + assert np.allclose(ref, dsl, atol=1e-12, rtol=1e-12) From e3ba89defff8a425f40e0d30742e4b4779fc8dc5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:39:33 +0200 Subject: [PATCH 21/48] test(tests): add subclass recursion regression tests for dsl models --- tests/test_dsl_subclass.py | 76 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 tests/test_dsl_subclass.py diff --git a/tests/test_dsl_subclass.py b/tests/test_dsl_subclass.py new file mode 100644 index 0000000..d44a256 --- /dev/null +++ b/tests/test_dsl_subclass.py @@ -0,0 +1,76 @@ +"""Regression test for the `super(type(self), self)` recursion trap in +`build_model`. + +Background: + The original prototype used `super(type(self), self)._init_dependant()` inside + the closure attached to the generated class. This works only as long as no one + subclasses the generated class — the moment a user does, `type(self)` becomes + the user's subclass and `super(type(self), self)` re-resolves to the SAME + generated method, causing infinite recursion. + + The fix captures the generated class in a closure cell so `super()` always + walks past it, regardless of `type(self)`. +""" +import numpy as np + +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +def test_subclass_overrides_init_dependant(): + """Direct subclass: override _init_dependant + super() must not recurse.""" + spec = ModelSpec( + name="SubclassableDSL", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[ + Parameter("a", default=2.0), + Parameter("b", default=3.0), + Parameter("c", formula="a + b"), + ], + equations="d_x = -x / c + coupling.x", + ) + Base = build_model(spec) + + class Extended(Base): + def _init_dependant(self): + super()._init_dependant() + self.user_extra = 42 + + n = 4 + m = Extended(g=0.0).configure(weights=np.zeros((n, n))) + assert np.allclose(m.c, 5.0) + assert m.user_extra == 42 + + +def test_grandchild_subclass_chain(): + """Two-level subclass: A(DSL) -> B(A). Each calls super(); none recurses.""" + spec = ModelSpec( + name="GrandparentDSL", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[ + Parameter("a", default=1.0), + Parameter("d", formula="a * 10"), + ], + equations="d_x = -x * d + coupling.x", + ) + Base = build_model(spec) + + class Mid(Base): + def _init_dependant(self): + super()._init_dependant() + self.mid_marker = "mid" + + class Leaf(Mid): + def _init_dependant(self): + super()._init_dependant() + self.leaf_marker = "leaf" + + m = Leaf(g=0.0).configure(weights=np.zeros((3, 3))) + assert np.allclose(m.d, 10.0) + assert m.mid_marker == "mid" + assert m.leaf_marker == "leaf" From 81eaa824fcb697954135dc5cac50e2c7ab392a50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 17:45:25 +0200 Subject: [PATCH 22/48] feat(dsl_models): add paired simulation comparison script --- examples/dsl_models/compare_simulations.py | 97 ++++++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 examples/dsl_models/compare_simulations.py diff --git a/examples/dsl_models/compare_simulations.py b/examples/dsl_models/compare_simulations.py new file mode 100644 index 0000000..845160f --- /dev/null +++ b/examples/dsl_models/compare_simulations.py @@ -0,0 +1,97 @@ +"""Run paired (hand-written, DSL) simulations and show that they agree. + +For each pair we: + 1. Run the hand-written model and the DSL spec on the same connectivity, + params, and deterministic integrator. + 2. Plot the two trajectories overlaid (visually identical). + 3. Plot the difference, which should sit at machine-precision noise. + +Run: PYTHONPATH=src python examples/dsl_models/compare_simulations.py +""" +from __future__ import annotations + +import os +import sys + +import numpy as np +import matplotlib.pyplot as plt + +from neuronumba.simulator.simulator import simulate_nodelay +from neuronumba.simulator.integrators.euler import EulerDeterministic +from neuronumba.simulator.models import Hopf, Deco2014 + +# Make the dsl_models package importable when run as a script. The script +# lives inside the package, so we add its parent directory to sys.path. +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from dsl_models import HopfDSL, Deco2014DSL # noqa: E402 + + +def make_W(n: int = 8, seed: int = 0) -> np.ndarray: + """Symmetric, zero-diagonal connectivity normalized to max 1.0.""" + rng = np.random.default_rng(seed) + W = rng.random((n, n)) + W = (W + W.T) / 2 + np.fill_diagonal(W, 0) + W /= W.max() + return W + + +def run(model_cls, weights, obs_var: str, *, t_max: float = 1000.0, + dt: float = 0.1, sampling_period: float = 1.0, **model_kwargs): + """Run a deterministic simulation and return the observed signal.""" + np.random.seed(0) # simulate_nodelay draws random tract lengths internally + model = model_cls(**model_kwargs) + integrator = EulerDeterministic(dt=dt) + return simulate_nodelay( + model, integrator, weights, obs_var, + sampling_period=sampling_period, + t_max_neuronal=t_max, + t_warmup=0.0, + ) + + +def plot_pair(ax_overlay, ax_diff, ref, dsl, *, title: str, n_show: int = 4): + """Overlay (left axis) + difference (right axis) for one model pair.""" + t = np.arange(ref.shape[0]) + for i in range(min(n_show, ref.shape[1])): + ax_overlay.plot(t, ref[:, i], color="black", lw=1.0, alpha=0.6, + label="hand-written" if i == 0 else None) + ax_overlay.plot(t, dsl[:, i], color="tab:red", lw=0.8, ls="--", + label="DSL" if i == 0 else None) + ax_overlay.set_title(f"{title} — trajectories (first {n_show} ROIs)") + ax_overlay.set_xlabel("sample index") + ax_overlay.legend(loc="upper right", fontsize=9) + + diff = ref - dsl + max_abs = np.abs(diff).max() + ax_diff.plot(t, diff, lw=0.6) + ax_diff.set_title(f"{title} — difference (max |Δ| = {max_abs:.2e})") + ax_diff.set_xlabel("sample index") + ax_diff.axhline(0.0, color="black", lw=0.5, alpha=0.3) + + +def main(): + n_rois = 8 + W = make_W(n=n_rois) + + # Hopf: diffusive coupling, observe state variable x. + hopf_ref = run(Hopf, W, obs_var="x", g=0.5) + hopf_dsl = run(HopfDSL, W, obs_var="x", g=0.5) + + # Deco2014: linear coupling, observe state variable S_e. + deco_ref = run(Deco2014, W, obs_var="S_e", g=0.5, auto_fic=False) + deco_dsl = run(Deco2014DSL, W, obs_var="S_e", g=0.5) + + fig, axes = plt.subplots(2, 2, figsize=(12, 7)) + plot_pair(axes[0, 0], axes[0, 1], hopf_ref, hopf_dsl, title="Hopf") + plot_pair(axes[1, 0], axes[1, 1], deco_ref, deco_dsl, title="Deco2014") + fig.suptitle("DSL vs hand-written: deterministic Euler, n=8 ROIs", y=1.0) + fig.tight_layout() + + print(f"Hopf max |Δ|: {np.abs(hopf_ref - hopf_dsl).max():.3e}") + print(f"Deco2014 max |Δ|: {np.abs(deco_ref - deco_dsl).max():.3e}") + plt.show() + + +if __name__ == "__main__": + main() From 84b5aa6df6790ee88474d52a6df7a15f419ab549 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:08:58 +0200 Subject: [PATCH 23/48] feat(simulator): rename HistoryDense to HistoryDelays and fix buffer init --- src/neuronumba/simulator/history.py | 74 ++++++++++++++++++----------- 1 file changed, 46 insertions(+), 28 deletions(-) diff --git a/src/neuronumba/simulator/history.py b/src/neuronumba/simulator/history.py index 7b7fb5e..089179a 100644 --- a/src/neuronumba/simulator/history.py +++ b/src/neuronumba/simulator/history.py @@ -20,10 +20,23 @@ def _init_dependant(self): self.n_rois = self.weights.shape[0] -class HistoryDense(History): +class HistoryDelays(History): + """Ring-buffer history with per-(target, source) conduction delays. - # Global linear coupling - g = Attr(default=None, required=True) + Contract — for compatibility with the simulator's m_couple step: + + `get_numba_sample(step)` returns a (n_cvars, n_rois) array where + entry [k, i] = g * sum_j W[i, j] * buffer[k, t_idx[i, j], j], + with t_idx[i, j] = (step - 1 - i_delays[i, j]) mod n_time. + + The per-pair time index is the only thing that distinguishes this from + the no-delay path; the (n_cvars, n_rois) output shape lets the existing + DSL `linear` coupling kernel run as an identity passthrough on top. + """ + + g = Attr(default=1.0, required=False) + # Delay matrix in seconds: delays[i, j] is the time it takes for a + # signal from region j to arrive at region i. delays = Attr(default=None, required=True) dt = Attr(required=True) @@ -32,53 +45,58 @@ class HistoryDense(History): n_time = Attr(dependant=True) def _init_dependant(self): + super()._init_dependant() self.i_delays = np.rint(self.delays / self.dt).astype(np.int32) - self.n_time = np.max(self.i_delays) + 1 - self.buffer = np.zeros((len(self.c_vars), self.n_time, self.n_rois)) + # +1 because step indices count from 1; we always need at least one + # past slot even when all delays round to zero. + self.n_time = int(np.max(self.i_delays)) + 1 + # Buffer shape: (n_cvars, n_time, n_rois). Initialized to zero — + # early steps before the buffer fills get implicit zero coupling + # from far-delay regions, which matches the convention. + self.buffer = np.zeros( + (self.n_cvars, self.n_time, self.n_rois), dtype=np.float64 + ) def get_numba_update(self): - # buffer = self.buffer n_cvars = self.n_cvars + n_time = self.n_time c_vars = self.c_vars - # addr = buffer.ctypes.data b_addr, b_shape, b_dtype = addr.get_addr(self.buffer) - @nb.njit((nb.void)(nb.int, nb.f8[:,:]), cache=NUMBA_CACHE) - def c_update(step: nb.int, state: NDA_f8_2d): - # data = nb.carray(addr.address_as_void_pointer(addr), buffer.shape, - # dtype=buffer.dtype) - data = addr.create_carray(b_addr, b_shape, b_dtype) - for i in range(n_cvars): - data[i, step % self.n_time, :] = state[c_vars[i], :] + @nb.njit(nb.void(nb.i8, nb.f8[:, :]), cache=NUMBA_CACHE) + def c_update(step, state): + buf = addr.create_carray(b_addr, b_shape, b_dtype) + slot = step % n_time + for k in range(n_cvars): + buf[k, slot, :] = state[c_vars[k], :] return c_update def get_numba_sample(self): - buffer = self.buffer + n_cvars = self.n_cvars + n_rois = self.n_rois n_time = self.n_time weights = self.weights i_delays = self.i_delays - c_vars = self.c_vars - n_cvars = self.n_cvars - n_rois = self.n_rois g = self.g + b_addr, b_shape, b_dtype = addr.get_addr(self.buffer) - @nb.njit((nb.f8[:,:])(nb.int), cache=NUMBA_CACHE) + @nb.njit(nb.f8[:, :](nb.i8), cache=NUMBA_CACHE) def h_sample(step): - time_idx = (step - 1 - i_delays + n_time) % n_time - result = np.empty((n_cvars, n_rois)) - for v in c_vars: - delayed_state = np.empty((n_rois, n_rois)) + buf = addr.create_carray(b_addr, b_shape, b_dtype) + result = np.empty((n_cvars, n_rois), dtype=np.float64) + for k in range(n_cvars): for i in range(n_rois): - delayed_state[i] = buffer[v, time_idx[i], i] - result[v] = np.sum(weights * delayed_state, axis=0) - return g * result + s = 0.0 + for j in range(n_rois): + t_idx = (step - 1 - i_delays[i, j] + n_time) % n_time + s += weights[i, j] * buf[k, t_idx, j] + result[k, i] = g * s + return result return h_sample - - class HistoryNoDelays(History): buffer = Attr(dependant=True) From 86c8101935703e8caa7e3b129c86ebd97942fc53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:01 +0200 Subject: [PATCH 24/48] docs(dsl): update readme with parameter and coupling clarifications --- src/neuronumba/simulator/models/dsl/README.md | 223 +++++++++++++++--- 1 file changed, 193 insertions(+), 30 deletions(-) diff --git a/src/neuronumba/simulator/models/dsl/README.md b/src/neuronumba/simulator/models/dsl/README.md index f0ceb60..6df53db 100644 --- a/src/neuronumba/simulator/models/dsl/README.md +++ b/src/neuronumba/simulator/models/dsl/README.md @@ -10,15 +10,16 @@ trajectories. Every model in `neuronumba/simulator/models/` is a `Model` subclass that hand-implements four things: state-variable names, coupling kernel, parameter -list, and a numba-compiled `dfun`. Most of that is boilerplate. Inside -`get_numba_dfun` you write 16 lines of `m[np.intp(P.x)]` parameter unpacking -and `state[i, :]` indexing before getting to the math. A "small variation" — -swap a sigmoid for a softplus in Deco2014 — currently means a new file. +list, and a numba-compiled `dfun`. Most of that is boilerplate — parameter +unpacking, state-variable indexing, return-tuple assembly. A "small +variation" — swap a sigmoid for a softplus in Deco2014 — currently means a +new file. The DSL pitch: write only the math, the parameter list, and the topology -declarations. The boilerplate is generated. The result is the same numba -closure structure you'd write by hand — proven bit-equivalent on Hopf, Deco2014, -Naskar2021, and Montbrio (see `tests/test_dsl_equivalence.py`). +declarations. The boilerplate is generated. The DSL-built class satisfies +the same `Model` interface as a hand-written one — proven on Hopf, Deco2014, +Naskar2021, and Montbrio with random-input dfun comparisons that agree to +machine precision (see `tests/test_dsl_equivalence.py`). ## Hello world @@ -74,11 +75,15 @@ Both produce identical trajectories. ## Parameters +A parameter has one of two flavors. Both kinds become regular Python +attributes on the configured model instance; the dfun captures them by +closure. + ### Independent A parameter is independent if it carries `default=` (or `required=True`) and -no `formula`. Independent parameters are user-settable and packed into -`self.m`: +no `formula`. Independent parameters are user-settable through the +constructor: ```python Parameter("tau_e", default=100.0) # default value @@ -88,27 +93,34 @@ Parameter("g", required=True) # caller must supply ### Dependent A parameter is dependent if it carries a `formula`. The formula is a Python -expression that may reference other parameters and `np`. Dependents are -computed at `configure()` time and recomputed every time the user calls -`configure()` again — that's the existing neuronumba protocol for "a -parameter was modified": +expression that may reference other parameters, `np`, and the model-context +names `weights`, `weights_t`, `g`, `n_rois`. Dependents are computed at +`configure()` time and recomputed every time the user calls `configure()` +again — that's the existing neuronumba protocol for "a parameter was +modified": ```python Parameter("J_ee", default=10.0), Parameter("g_ee", default=2.5), Parameter("a_e", default=0.25), Parameter("J_N_ee", formula="J_ee + g_ee * np.log(a_e)"), + +# Matrix-shaped dependents: just return a 2D ndarray from the formula. +Parameter("A", formula="(-g * weights + np.diag(weights.sum(axis=1))) / tau"), ``` -Order doesn't matter — dependents are topologically sorted at build time. -Cycles and self-references raise `ValueError`. Multi-level dependencies -work: a dependent can reference another dependent. +Order of declaration doesn't matter — dependents are topologically sorted +at build time. Cycles and self-references raise `ValueError`. Multi-level +dependencies work: a dependent can reference another dependent. -### `regional` vs global +### Shapes -By default parameters are `regional=True`: they're per-ROI scalars packed -into `self.m`. Pass `regional=False` to mark a parameter as global (not -per-region); these go into `self.m_aux`. Most parameters are regional. +There's no explicit shape declaration. Whatever the user supplies (or the +formula returns) is the parameter's value: scalar, per-region 1D array, or +2D matrix. The dfun captures it via closure and numba's type inference +handles it. Use `A @ x` for matrix-vector, `tau_e * S_e` for scalar or +per-region — the same expression works whether `tau_e` is a scalar or an +`(n_rois,)` array. ### Why no `default + formula`? @@ -116,6 +128,90 @@ Setting both is a signal of confused intent: a dependent value cannot also have a default. The validator rejects this combination at `Parameter` construction time. Same for `required=True` + `formula`. +## The `on_configure` escape hatch + +Some setup doesn't fit the declarative `formula=` form: FIC needs `fsolve`, +steady-state computations need iteration, parameters might need to be +derived from `self.weights` and `self.g` together. For those cases pass an +`on_configure` callback to `ModelSpec`: + +```python +from neuronumba.fitting.fic.fic import FICHerzog2022 + +def auto_fic(self): + self.J = FICHerzog2022().compute_J(self.weights, self.g) + +deco_spec = ModelSpec( + ..., + parameters=[ + ..., + Parameter("J", default=1.0), # placeholder; auto_fic overwrites + ], + on_configure=auto_fic, +) +``` + +The callback fires once per `configure()`, *after* dependent-parameter +evaluation. Whatever the callback assigns to `self` (e.g. `self.J = ...`) +persists on the instance and is visible to the dfun the next time +`get_numba_dfun()` is called — the factory reads `model.` fresh +at that point. + +When to use: +- **Imperative computation** that needs Python control flow (loops, + fsolve, iterative steady-state). +- **Cross-parameter setup** that depends on multiple inputs in non-formula + ways (e.g. solving a linear system). + +When **not** to use: +- A simple formula like `J_N_ee = J_ee + g_ee * np.log(a_e)` — use + `Parameter(..., formula=...)` instead. The DSL handles ordering and + recompute-on-modify automatically. + +## Helpers: user `@nb.njit` functions + +Some models need custom subroutines that aren't in the whitelisted numpy +function set — e.g. Zerlaut's `erfc_approx` or threshold-piecewise +functions. Pass them via `helpers=[...]`: + +```python +import numba as nb + +@nb.njit(nb.f8[:](nb.f8[:]), cache=False) +def erfc_approx(x): + # ... implementation ... + +@nb.njit(nb.f8[:](nb.f8[:], nb.f8[:]), cache=False) +def threshold_func(x, theta): + # ... implementation ... + +zerlaut_spec = ModelSpec( + ..., + helpers=[erfc_approx, threshold_func], + equations=""" + s = erfc_approx(some_combination) + rate = threshold_func(I, theta) + d_S = ... + """, +) +``` + +The DSL binds each helper as a closure variable inside the generated dfun, +so numba sees and compiles them as regular jitted callees. + +### Constraints + +- **Each helper must be `@nb.njit`-decorated.** Passing a plain Python + function will compile-fail when numba reaches the call site. +- **Each helper needs a unique `__name__`.** Lambdas and other callables + without a usable name are rejected up front. +- **Names must not clash** with state vars, coupling vars, parameter + names, declared observables, whitelisted numpy functions, or the + reserved tokens (`state`, `coupling`, `m`, `P`, `np`). +- **Unused helpers are silently dropped** — list one in `helpers=[...]` + and never call it from the equations and the generated dfun simply + doesn't bind it. + ## Coupling kinds The DSL ships two pre-built kernels: @@ -136,9 +232,46 @@ coupling = g * (W^T @ S - sum(W^T, axis=1) * S) Hopf-style: subtracts the local self-loop contribution from each region. -For anything else — delays, custom expressions, mixed kernels — subclass -`Model` directly and override `get_numba_coupling`. The DSL deliberately -keeps the coupling surface narrow; complex coupling deserves the imperative +### `delayed` + +``` +coupling[i] = g * sum_j W[i, j] * state[j, t - delay[i, j]] +``` + +Conduction-delay coupling: each region receives signals from others +shifted in time by `delay[i, j] = lengths[i, j] / speed`. Use this when +modelling realistic propagation along white-matter tracts where the +~5-30 ms delay matters for phase relations. + +To run a delayed simulation, swap the simulator's history class: + +```python +from neuronumba.simulator.history import HistoryDelays +from neuronumba.simulator.connectivity import Connectivity +from neuronumba.simulator.simulator import Simulator + +sim = Simulator( + connectivity=Connectivity(weights=W, lengths=L, speed=10.0), + model=DelayedDSL(g=0.5), + history=HistoryDelays(), # not HistoryNoDelays + integrator=integrator, + monitors=[monitor], +) +sim.run(0, t_max) +``` + +The `Simulator` automatically passes `lengths/speed` (and the model's `g`, +the integrator's `dt`) to `HistoryDelays.configure()`. + +**Limitations in v0.2:** +- All coupling vars in a spec must use the same kind. Mixing `delayed` + with `linear` or `diffusive` raises at `build_model` time. +- `get_jacobian()` raises `NotImplementedError` for delayed kernels — + the Jacobian of a delay-differential system isn't a single matrix. + +For anything else — custom expressions, mixed kernels — subclass `Model` +directly and override `get_numba_coupling`. The DSL deliberately keeps +the coupling surface narrow; complex coupling deserves the imperative escape hatch. ## The equation language @@ -151,8 +284,8 @@ language. So if you can write it in numba, you can write it here. - **State variables** declared in `state_vars`. Inside equations, `x` refers to `state[i, :]` where `i` is the row index. The DSL handles the unpacking. -- **Parameters** by name. The DSL emits `tau_e = m[np.intp(P.tau_e)]` - bindings at the top of the generated function. +- **Parameters** by name. The DSL emits `tau_e = model.tau_e` bindings at + the top of the generated factory; the inner dfun captures them by closure. - **Intermediates** introduced by assignment: `Ie = ...` is a fine local variable. It must be assigned before it's used. - **`coupling.`** for declared coupling vars. Rewritten to @@ -240,6 +373,36 @@ from neuronumba.simulator.models.dsl import cleanup_cache removed = cleanup_cache(max_age_days=1) # nuke yesterday's iterations ``` +## Jacobians + +DSL-built models ship a numerical `get_jacobian(state)` that returns the +network Jacobian at a given operating point: + +```python +m = HopfDSL(g=0.5, a=-0.5, omega=0.3).configure(weights=W) +J = m.get_jacobian(np.zeros((m.n_state_vars, m.n_rois))) +# shape: (n_state_vars * n_rois, n_state_vars * n_rois) +# index: J[u*N + i, v*N + j] = d(d_state_u at region i)/d(state_v at region j) +``` + +Implementation: centered finite differences on the dfun for local partials +(d_state w.r.t. state, d_state w.r.t. coupling) plus the closed-form +linearization of the coupling kernel (which the DSL knows from each +`CouplingVar.kind`). No SymPy, no new dependencies, ~1e-7 accuracy with the +default step. Validated against a brute-force naive Jacobian for Hopf, Deco, +and Naskar in `tests/test_dsl_jacobian.py`. + +Tune the FD step size with the optional `eps` keyword if your model has +parameters at unusual scales: + +```python +J = m.get_jacobian(state, eps=1e-8) # tighter; risk of roundoff +``` + +If you need an analytic Jacobian (machine-precision, faster) you can still +hand-write `get_jacobian` on a subclass — the DSL-built class is just a +regular `Model` subclass, so override and you're done. + ## When to subclass `Model` directly The DSL covers the 80% case. Reach for a hand-written `Model` (or @@ -254,12 +417,12 @@ The DSL covers the 80% case. Reach for a hand-written `Model` (or *subclass a DSL-built class* and override `_init_dependant` to call `super()` and then do your imperative setup. See `tests/test_dsl_subclass.py`. -3. **You need an analytic Jacobian** (`get_jacobian`). Not auto-derived. -4. **You need non-scalar matrix parameters** (e.g. OrnsteinUhlenbeck's `A` - from `weights/g/tau`). The DSL only handles regional scalars right now. -5. **You need delays beyond linear/diffusive coupling.** +3. **You need delays beyond linear/diffusive/delayed coupling** (e.g. + per-coupling-var custom delays, mixed-kind specs). -These are documented as deferred items in `IMPLEMENTATION_PLAN.md` Phase 6. +These are documented as deferred items in [ROADMAP.md](ROADMAP.md), which +covers what each missing feature would look like, why it was deferred, and +the design questions to settle before implementing. ## Reference specs From 2234dc6271e47b89af3798cebadb78c61b5634fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:05 +0200 Subject: [PATCH 25/48] feat(dsl): update dfun builder to use closure-based param capture --- .../simulator/models/dsl/builder.py | 393 ++++++++++++++---- 1 file changed, 302 insertions(+), 91 deletions(-) diff --git a/src/neuronumba/simulator/models/dsl/builder.py b/src/neuronumba/simulator/models/dsl/builder.py index fb58c9f..da476e2 100644 --- a/src/neuronumba/simulator/models/dsl/builder.py +++ b/src/neuronumba/simulator/models/dsl/builder.py @@ -1,12 +1,18 @@ """ Source generators and `build_model` — the entry point that turns a `ModelSpec` into a numba-friendly `Model` subclass. + +Architecture: every name referenced in the dfun (parameters, helpers, etc.) +resolves through ``model.`` at the factory level. The factory takes a +single argument — the model instance — and the inner dfun closes over the +captured locals. There is no `m`/`P`/`m_aux` packing for DSL-built classes; +parameters live as regular Python attributes on the instance. """ from __future__ import annotations import ast import textwrap -from typing import Any, Dict, List, Type +from typing import Any, Dict, List, Tuple, Type import numba as nb import numpy as np @@ -17,18 +23,24 @@ from .dependents import _topo_sort_dependents from .materialize import _materialize_module -from .rewriter import _DfunRewriter +from .rewriter import ALLOWED_NP_FUNCS, _DfunRewriter from .spec import ModelSpec -def _build_dfun_source(spec: ModelSpec) -> str: - """Return executable Python source for a function `dfun(state, coupling)` - that uses captured `m` and `P` variables (added by exec env).""" +def _build_dfun_source(spec: ModelSpec) -> Tuple[str, List[str]]: + """Return ``(source, used_names)`` for the dfun factory of `spec`. + + The source defines a ``make__dfun(model)`` factory whose inner + function is the dfun. ``used_names`` is the sorted list of param + helper + names referenced by the equations; each gets a ``name = model.name`` + binding at the factory level so the inner dfun captures it via closure. + """ state_index = {sv.name: i for i, sv in enumerate(spec.state_vars)} coupling_index = {cv.name: i for i, cv in enumerate(spec.coupling_vars)} param_names = {p.name for p in spec.parameters} obs_names = set(spec.observables) + helper_names = {h.__name__ for h in spec.helpers} src = textwrap.dedent(spec.equations).strip() try: @@ -39,7 +51,8 @@ def _build_dfun_source(spec: ModelSpec) -> str: ) from e rewriter = _DfunRewriter( - state_index, coupling_index, param_names, obs_names + state_index, coupling_index, param_names, obs_names, + helper_names=helper_names, ) new_body = [rewriter.visit(stmt) for stmt in tree.body] @@ -90,33 +103,35 @@ def _build_dfun_source(spec: ModelSpec) -> str: body_src = "\n".join(ast.unparse(stmt) for stmt in new_body) - # Header bindings: state vars from `state` array, params from `m`. + # State-var bindings live inside the inner dfun (read from `state`). state_binds = "\n".join( f" {sv.name} = state[{i}, :]" for i, sv in enumerate(spec.state_vars) if sv.name in rewriter.used_state ) - param_binds = "\n".join( - f" {p} = m[np.intp(P.{p})]" - for p in sorted(rewriter.used_params) + + # All parameters and helpers bind once at the factory level via + # `model.`. Numba captures them as compile-time constants when the + # inner dfun is JIT'd. Helper functions (also model attributes since + # build_model stashes them on the class) come through the same path. + used_names = sorted(rewriter.used_params | rewriter.used_helpers) + factory_binds = "\n".join( + f" {name} = model.{name}" for name in used_names ) indented_body = textwrap.indent(body_src, " ") - # Factory function: `m` and `P` are passed in and become real closure - # variables for the inner dfun — exactly the pattern the hand-written - # neuronumba models use. - parts = [f"def make_{spec.name}_dfun(m, P):"] + parts = [f"def make_{spec.name}_dfun(model):"] + if factory_binds: + parts.append(factory_binds) parts.append(f" def {spec.name}_dfun(state, coupling):") if state_binds: parts.append(state_binds) - if param_binds: - parts.append(param_binds) parts.append(indented_body) parts.append(f" return np.stack(({d_stack_args},)), {ret_obs}") parts.append(f" return {spec.name}_dfun") fn_src = "\n".join(parts) + "\n" - return fn_src + return fn_src, used_names def _build_coupling_kernel(spec: ModelSpec): @@ -148,16 +163,186 @@ def _diffusive_coupling(state): return _diffusive_coupling return make_coupling + if kinds == {"delayed"}: + # `HistoryDelays.get_numba_sample` already computed + # g * sum_j W[i,j] * state[k, j, t-delay[i,j]] and returned shape + # (n_cvars, n_rois). The coupling step is therefore identity. + def make_coupling(self): + @nb.njit(nb.f8[:, :](nb.f8[:, :]), cache=NUMBA_CACHE) + def _delayed_coupling(state): + return state + return _delayed_coupling + return make_coupling + + if "delayed" in kinds: + raise NotImplementedError( + f"Mixing 'delayed' coupling with other kinds ({sorted(kinds)}) is " + f"not supported in this version. Either make all coupling vars " + f"'delayed' or none." + ) + raise NotImplementedError( - f"Unsupported coupling kinds {sorted(kinds)}. The DSL ships 'linear' " - f"and 'diffusive'; subclass `Model` directly and override " + f"Unsupported coupling kinds {sorted(kinds)}. The DSL ships 'linear', " + f"'diffusive', and 'delayed'; subclass `Model` directly and override " f"`get_numba_coupling` for anything else." ) +# Names that the generated dfun source uses directly (function args, np +# alias) and that helpers must therefore not shadow. +_RESERVED_HELPER_NAMES = {"state", "coupling", "model", "np"} + + +def _validate_helpers(spec: ModelSpec) -> None: + """Reject helper lists that would clash with other names in the dfun. + + A helper conflicts with the dfun namespace if its `__name__` matches a + state var, coupling var, parameter, observable, whitelisted numpy + function, a reserved token, or any inherited `LinearCouplingModel` + method/attribute name (helpers are stashed as class attributes; a + collision would shadow the inherited member). Two helpers with the same + `__name__` are also a conflict. + """ + if not spec.helpers: + return + + seen: Dict[str, int] = {} + for i, h in enumerate(spec.helpers): + name = getattr(h, "__name__", None) + if not name: + raise ValueError( + f"helpers[{i}] has no __name__; pass a regular function, not " + f"a lambda or callable object." + ) + if name in seen: + raise ValueError( + f"helpers[{i}] and helpers[{seen[name]}] both have " + f"__name__='{name}'; helper names must be unique." + ) + seen[name] = i + + # Names already on the base class (Model methods, HasAttr machinery). + # Stashing a helper with the same __name__ would shadow them and break + # the model contract. + base_attrs = {n for n in dir(LinearCouplingModel) if not n.startswith("__")} + + forbidden = { + sv.name for sv in spec.state_vars + } | { + cv.name for cv in spec.coupling_vars + } | { + p.name for p in spec.parameters + } | set(spec.observables) | _RESERVED_HELPER_NAMES | ALLOWED_NP_FUNCS | base_attrs + + clashes = set(seen) & forbidden + if clashes: + raise ValueError( + f"Helper name(s) {sorted(clashes)} clash with state vars, " + f"coupling vars, parameters, observables, reserved tokens, " + f"whitelisted numpy functions, or base-class methods. " + f"Rename the helpers." + ) + + +def _coupling_kernel_jacobian(g: float, W: np.ndarray, kind: str) -> np.ndarray: + """Closed-form Jacobian of one coupling kernel. + + Returns ``C`` of shape ``(N, N)`` where + ``C[i, j] = d coupling[i] / d state_at_coupling_var[j]``. + """ + if kind == "linear": + # coupling[i] = sum_j state[j] * g * W[i, j] + return g * W + if kind == "diffusive": + # coupling[i] = g * (sum_j state[j] * W[i, j] - ink[i] * state[i]) + # ink[i] = sum_l W[l, i] (col-sum of W; row-sum if W is symmetric) + ink = W.T.sum(axis=1) + return g * (W - np.diag(ink)) + if kind == "delayed": + # The Jacobian of a delay-differential system isn't a single matrix + # — it's a delay-coupled spectrum. Linear stability for delayed + # systems needs different tooling (DDE-Biftool, Lambert-W methods). + raise NotImplementedError( + "get_jacobian is not supported for 'delayed' coupling kernels. " + "The Jacobian of a delay-differential system is not a single " + "matrix. Use a no-delay variant (replace 'delayed' with 'linear' " + "or 'diffusive') for linear stability analysis, or analyse the " + "system with DDE-aware tooling." + ) + raise NotImplementedError( + f"Jacobian not implemented for coupling kind '{kind}'. " + f"Subclass and override `get_jacobian` for custom kernels." + ) + + +def _compute_jacobian(model, state, eps: float = 1e-6) -> np.ndarray: + """Two-piece numerical Jacobian: local FD + closed-form coupling assembly. + + Used as the body of `get_jacobian` on every DSL-built class. See the + docstring on the class method for shape and indexing conventions. + """ + state = np.asarray(state, dtype=np.float64) + nsv = model.n_state_vars + N = model.n_rois + if state.shape != (nsv, N): + raise ValueError( + f"state must have shape ({nsv}, {N}); got {state.shape}" + ) + + c_vars = list(model.c_vars) + ncv = len(c_vars) + + # Get the current dfun + coupling closures (capture current params). + dfun = model.get_numba_dfun() + coupling_factory = model.get_numba_coupling() + coupling = coupling_factory(state[c_vars, :].copy()) # shape (ncv, N) + + # 1) Local partials: dD_u/d(state_v) at each region. + # The dfun is per-region pointwise (no cross-region dependence), + # so perturbing the entire row of state_v gives the per-region + # derivative with no cross-talk. + local = np.empty((nsv, nsv, N)) + for v in range(nsv): + s_plus = state.copy(); s_plus[v] += eps + s_minus = state.copy(); s_minus[v] -= eps + out_plus, _ = dfun(s_plus, coupling.copy()) + out_minus, _ = dfun(s_minus, coupling.copy()) + local[:, v, :] = (out_plus - out_minus) / (2.0 * eps) + + # 2) Coupling partials: dD_u/d(coupling_k) at each region. + coup = np.empty((nsv, ncv, N)) + for k in range(ncv): + c_plus = coupling.copy(); c_plus[k] += eps + c_minus = coupling.copy(); c_minus[k] -= eps + out_plus, _ = dfun(state.copy(), c_plus) + out_minus, _ = dfun(state.copy(), c_minus) + coup[:, k, :] = (out_plus - out_minus) / (2.0 * eps) + + # 3) Closed-form coupling-kernel Jacobians, one (N, N) matrix per cv. + coupling_jacobians = [ + _coupling_kernel_jacobian(model.g, model.weights, kind) + for kind in model._coupling_var_kinds + ] + + # 4) Assemble the network Jacobian: + # J[u*N + i, v*N + j] = local[u, v, i] * delta_ij + # + sum_k (v == c_vars[k]) * coup[u, k, i] * C_k[i, j] + J = np.zeros((nsv * N, nsv * N)) + for u in range(nsv): + for v in range(nsv): + J[u * N:(u + 1) * N, v * N:(v + 1) * N] = np.diag(local[u, v, :]) + for k in range(ncv): + v = c_vars[k] + C_k = coupling_jacobians[k] + for u in range(nsv): + J[u * N:(u + 1) * N, v * N:(v + 1) * N] += coup[u, k, :, None] * C_k + + return J + + def build_model(spec: ModelSpec) -> Type[Model]: - """Compile a ModelSpec into a `Model` subclass equivalent to a hand-written - neuronumba model. The returned class has the same external contract.""" + """Compile a ModelSpec into a `Model` subclass with the same external + contract as a hand-written neuronumba model.""" # Validate that coupling vars are state vars. state_names = [sv.name for sv in spec.state_vars] @@ -167,9 +352,12 @@ def build_model(spec: ModelSpec) -> Type[Model]: f"Coupling var '{cv.name}' is not declared as a state var." ) - # Build the dfun source. We embed it into a complete module that imports - # numpy (so numba can later compile from a real file path). - dfun_body_src = _build_dfun_source(spec) + # Validate helpers up-front. + _validate_helpers(spec) + + # Build the dfun source. Embed in a module that imports numpy so numba + # can compile from a real file path. + dfun_body_src, _used_names = _build_dfun_source(spec) full_src = ( "# Auto-generated by nndsl. Do not edit.\n" "import numpy as np\n" @@ -185,81 +373,47 @@ def build_model(spec: ModelSpec) -> Type[Model]: # Class-level model-info attributes. cls_dict["_state_var_names"] = [sv.name for sv in spec.state_vars] cls_dict["_coupling_var_names"] = [cv.name for cv in spec.coupling_vars] + cls_dict["_coupling_var_kinds"] = [cv.kind for cv in spec.coupling_vars] cls_dict["_observable_var_names"] = list(spec.observables) bounds = {sv.name: sv.bounds for sv in spec.state_vars if sv.bounds is not None} if bounds: cls_dict["_state_var_bounds"] = bounds - # Parameter Attrs. Independent params get `default`; dependents get - # `dependant=True` so neuronumba's existing machinery treats them like - # regional params that the model itself fills in during `_init_dependant`. + # Parameter Attrs. No `attributes=Tag.REGIONAL` — DSL params do not flow + # through the inherited `_init_dependant_automatic` packing path + # (see the no-op override below). The Attr is purely for default-setting + # and `__init__` kwarg validation. independent_params = [p for p in spec.parameters if not p.is_dependent] dependent_params = [p for p in spec.parameters if p.is_dependent] for p in independent_params: - attr_kwargs: Dict[str, Any] = { - "default": p.default, - "doc": p.doc, - } + attr_kwargs: Dict[str, Any] = {"default": p.default, "doc": p.doc} if p.required: attr_kwargs["required"] = True - attr_kwargs["attributes"] = ( - Model.Tag.REGIONAL if p.regional else Model.Tag.GLOBAL - ) cls_dict[p.name] = Attr(**attr_kwargs) for p in dependent_params: - # `dependant=True` tells HasAttr/Model to skip default-setting and - # required-checking; we'll fill these in during _init_dependant. - cls_dict[p.name] = Attr( - dependant=True, - doc=p.doc, - attributes=(Model.Tag.REGIONAL if p.regional else Model.Tag.GLOBAL), - ) - - # Topo-sort dependents once at build time and compile each formula. + # `dependant=True` skips default-setting and tells HasAttr to refuse + # the param as an __init__ kwarg. We populate it in _init_dependant. + cls_dict[p.name] = Attr(dependant=True, doc=p.doc) + + # Stash helpers as class attributes so the factory can resolve them + # via `model.`. They're shared across all instances. + for helper in spec.helpers: + cls_dict[helper.__name__] = helper + + # Override the inherited m/m_aux packing to a no-op. DSL parameters live + # as plain instance attributes; the dfun captures them via closure. + def _init_dependant_automatic(self): + pass + cls_dict["_init_dependant_automatic"] = _init_dependant_automatic + + # Topo-sort dependents once at build time and pre-compile each formula. sorted_deps = _topo_sort_dependents(spec) compiled_formulas = [ (p.name, compile(p.formula, f"", "eval")) for p in sorted_deps ] - # _init_dependant: evaluate formulas in order, on the instance. Runs at - # configure() time, i.e. on construction and any subsequent configure() call, - # which is the existing protocol for "parameter was modified" in neuronumba. - # - # Mutable container that will hold the generated class once `type()` returns - # it below. Captured by the closure so `super()` always resolves against the - # DSL-generated class, not against `type(self)`. Without this indirection, - # subclassing the generated class would re-enter this method via super() - # forever (because `type(self)` would be the user's subclass). - _class_cell: List[Type[Model]] = [None] - - def _init_dependant(self): - # Run the base-class hook first (sets n_rois, weights_t for - # LinearCouplingModel, etc.). Without this, `self.weights` exists but - # things like `self.weights_t` don't yet. - super(_class_cell[0], self)._init_dependant() - # Evaluation env: numpy + every parameter currently on self (independent - # ones from the user, plus dependents already computed earlier in the - # topo order). - env: Dict[str, Any] = {"np": np} - # Snapshot all params that exist so far. - for p in spec.parameters: - val = getattr(self, p.name, None) - if val is not None: - env[p.name] = val - for name, code_obj in compiled_formulas: - try: - value = eval(code_obj, env) - except Exception as e: - raise RuntimeError( - f"Error evaluating dependent parameter '{name}' in model " - f"'{spec.name}': {e}" - ) from e - setattr(self, name, value) - env[name] = value - cls_dict["_init_dependant"] = _init_dependant - - # `g` is pulled in by LinearCouplingModel; we don't redeclare it. + on_configure_cb = spec.on_configure initial_overrides = dict(spec.initial_state_overrides) state_initials = [sv.initial for sv in spec.state_vars] @@ -274,15 +428,13 @@ def initial_state(self, n_rois: int) -> np.ndarray: return out cls_dict["initial_state"] = initial_state - # Capture spec metadata in the closure for error reporting. + # Capture spec metadata for error reporting. _spec_name = spec.name _spec_equations = spec.equations _source_file = gen_module.__file__ def get_numba_dfun(self): - m_local = self.m.copy() - P_local = self.P - py_dfun = factory(m_local, P_local) + py_dfun = factory(self) try: jitted = nb.njit( nb.types.UniTuple(nb.f8[:, :], 2)(nb.f8[:, :], nb.f8[:, :]), @@ -301,10 +453,28 @@ def get_numba_dfun(self): coupling_factory = _build_coupling_kernel(spec) cls_dict["get_numba_coupling"] = coupling_factory - # Choose base class. LinearCouplingModel only adds `g` + `weights_t`. We - # always inherit from it: even diffusive Hopf needs `g` and `weights_t`, - # and for 'linear' it gives us the default coupling (which we override - # only when we want — but we override always for clarity here). + def get_jacobian(self, state, eps: float = 1e-6) -> np.ndarray: + """Network Jacobian d(d_state)/d(state) at the operating point ``state``. + + Computed via centered finite differences on the dfun, combined with + the closed-form linearization of the coupling kernel (which the DSL + already knows from each `CouplingVar.kind`). + + Args: + state: shape ``(n_state_vars, n_rois)`` operating point. + eps: finite-difference step. Default ``1e-6`` gives roughly + ``1e-10`` accuracy with centered differences. + + Returns: + Jacobian of shape ``(n_state_vars * n_rois, n_state_vars * n_rois)``. + Indexing follows the convention ``J[u*N + i, v*N + j] = d(dfun_u + at region i)/d(state_v at region j)``. + """ + return _compute_jacobian(self, state, eps) + cls_dict["get_jacobian"] = get_jacobian + + # Inherit from LinearCouplingModel: we get `g`, `weights_t`, the base + # `_init_dependant` (sets weights_t/n_rois), and `get_numba_validate`. base = LinearCouplingModel cls_dict["__doc__"] = ( @@ -318,13 +488,54 @@ def get_numba_dfun(self): cls = type(spec.name, (base,), cls_dict) cls.__nndsl_source_file__ = gen_module.__file__ - _class_cell[0] = cls + + # Define `_init_dependant` AFTER class creation so the closure captures + # the real `cls` directly (no `_class_cell[0]` indirection). Defines + # super() resolution to walk past `cls` regardless of the runtime type + # of `self` — crucial when users subclass the DSL-generated class. + def _init_dependant(self): + super(cls, self)._init_dependant() + # Evaluation env: numpy + model-context attrs + every parameter + # currently on self (independents from the user, dependents already + # computed earlier in the topo order). + env: Dict[str, Any] = {"np": np} + for ctx_name in ("weights", "weights_t", "g", "n_rois"): + val = getattr(self, ctx_name, None) + if val is not None: + env[ctx_name] = val + for p in spec.parameters: + val = getattr(self, p.name, None) + if val is not None: + env[p.name] = val + for name, code_obj in compiled_formulas: + try: + value = eval(code_obj, env) + except Exception as e: + raise RuntimeError( + f"Error evaluating dependent parameter '{name}' in model " + f"'{spec.name}': {e}" + ) from e + # Contiguify any 2D-or-higher array result so the dfun's + # closure capture sees C-contiguous memory (numba prefers it). + if isinstance(value, np.ndarray) and value.ndim >= 2: + value = np.ascontiguousarray(value) + setattr(self, name, value) + env[name] = value + # User imperative hook runs last — sees fully-evaluated dependents, + # can freely mutate self. Mutations are visible to the dfun the next + # time `get_numba_dfun()` is called (factory reads `model.` + # fresh at that point). + if on_configure_cb is not None: + on_configure_cb(self) + cls._init_dependant = _init_dependant + return cls def dump_generated(spec: ModelSpec) -> str: """Return the generated dfun source for inspection.""" - return _build_dfun_source(spec) + src, _ = _build_dfun_source(spec) + return src def get_source_file(cls: Type[Model]) -> str: From a9ef80bd4f37536a679e291f4af8d58fb479a720 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:08 +0200 Subject: [PATCH 26/48] feat(dsl): update dependents to allow model context names in formulas --- src/neuronumba/simulator/models/dsl/dependents.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/neuronumba/simulator/models/dsl/dependents.py b/src/neuronumba/simulator/models/dsl/dependents.py index 01ce3e3..464c619 100644 --- a/src/neuronumba/simulator/models/dsl/dependents.py +++ b/src/neuronumba/simulator/models/dsl/dependents.py @@ -14,6 +14,15 @@ from .spec import ModelSpec, Parameter +# Model-context attributes available inside dependent-param formulas. These +# are real attributes on the configured model instance (set by +# `LinearCouplingModel._init_dependant` before our hook runs); we expose +# them by name so formulas can reference e.g. `g * weights` or +# `np.eye(n_rois)`. They're allowed in formulas but do NOT create +# parameter-graph dependencies (they aren't declared params). +MODEL_CONTEXT_NAMES = frozenset({"weights", "weights_t", "g", "n_rois"}) + + def _formula_deps(name: str, src: str, allowed: set) -> List[str]: """Return the parameter names referenced in `src`. Validate identifiers.""" try: @@ -31,11 +40,14 @@ def visit_Name(self, node): return if n in ALLOWED_NP_FUNCS: return + if n in MODEL_CONTEXT_NAMES: + # Allowed but doesn't create a param-graph dependency. + return if n not in allowed: raise ValueError( f"Parameter '{name}': formula references unknown identifier " f"'{n}' (line {node.lineno}). Must be another declared " - f"parameter or numpy." + f"parameter, numpy, or one of {sorted(MODEL_CONTEXT_NAMES)}." ) if n not in seen: deps.append(n) From 6e9b7c64d38f110e6a94cef64557f4b3c68483d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:12 +0200 Subject: [PATCH 27/48] feat(dsl): update rewriter to support helper functions in dfun bodies --- .../simulator/models/dsl/rewriter.py | 29 ++++++++++++------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/src/neuronumba/simulator/models/dsl/rewriter.py b/src/neuronumba/simulator/models/dsl/rewriter.py index cf33853..1b4ba96 100644 --- a/src/neuronumba/simulator/models/dsl/rewriter.py +++ b/src/neuronumba/simulator/models/dsl/rewriter.py @@ -1,15 +1,16 @@ """ AST rewriter for DSL equation bodies. -Transforms identifiers in user-supplied equations to the -`m[np.intp(P.x)]` / `state[i, :]` / `coupling[k, :]` form used by the -hand-written neuronumba dfuns, and validates that every identifier is -declared (state var, parameter, intermediate, or whitelisted numpy func). +Transforms identifiers in user-supplied equations and validates that every +identifier is declared (state var, parameter, intermediate, helper, or +whitelisted numpy func). State vars and coupling refs get rewritten to +indexed access; everything else is left as-is and bound at the factory +level via `model.` closure capture. """ from __future__ import annotations import ast -from typing import Dict, List +from typing import Dict, List, Optional # Numpy functions allowed inside dfun bodies. Using a whitelist is the cheapest @@ -32,14 +33,17 @@ def __init__( coupling_index: Dict[str, int], param_names: set, observable_names: set, + helper_names: Optional[set] = None, ): self.state_index = state_index self.coupling_index = coupling_index self.param_names = param_names self.observable_names = observable_names + self.helper_names = helper_names or set() self.intermediates: set = set() # names introduced by user assignments self.used_state: set = set() self.used_params: set = set() + self.used_helpers: set = set() self.errors: List[str] = [] # --- assignment LHS handling ------------------------------------------ @@ -74,24 +78,29 @@ def visit_Name(self, node: ast.Name): return node if n in self.param_names: - # Don't rewrite parameters. They get bound as locals at the top of - # the generated function (e.g. `tau_e = m[np.intp(P.tau_e)]`). + # Parameters get bound at factory level via `model.` closure + # capture; we don't rewrite the AST node. self.used_params.add(n) return node # Allowed: intermediates introduced earlier, observables, the special - # tokens 'state', 'coupling', 'm', 'P', 'np', and known numpy functions. + # tokens 'state', 'coupling', 'model', 'np', known numpy functions, + # and user-supplied helper functions. if n in self.intermediates: return node - if n in {"state", "coupling", "m", "P", "np"}: + if n in {"state", "coupling", "model", "np"}: return node if n in ALLOWED_NP_FUNCS: return node + if n in self.helper_names: + self.used_helpers.add(n) + return node # Anything else is an undeclared identifier — error. self.errors.append( f"Line {node.lineno}: unknown identifier '{n}'. " - f"Not a state var, parameter, intermediate, or allowed numpy func." + f"Not a state var, parameter, intermediate, helper, or allowed " + f"numpy func." ) return node From e6dbc71cabe9f5d8c5ffb4d62b88232d5a209c05 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:16 +0200 Subject: [PATCH 28/48] feat(dsl): update spec to support matrix params and delayed coupling --- src/neuronumba/simulator/models/dsl/spec.py | 53 ++++++++++++++++----- 1 file changed, 41 insertions(+), 12 deletions(-) diff --git a/src/neuronumba/simulator/models/dsl/spec.py b/src/neuronumba/simulator/models/dsl/spec.py index 1e9a8a9..35c7a9e 100644 --- a/src/neuronumba/simulator/models/dsl/spec.py +++ b/src/neuronumba/simulator/models/dsl/spec.py @@ -8,7 +8,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple @dataclass @@ -25,6 +25,7 @@ class CouplingVar: kind: 'linear' — coupling = g * (W^T @ S) (standard SC coupling) 'diffusive' — coupling = g * (W^T @ S - sum(W^T,1) * S) (Hopf-style) + 'delayed' — coupling = g * sum_j W[i,j] * state[j, t-delay[i,j]] For coupling shapes outside this set, subclass `Model` directly and override `get_numba_coupling`. @@ -38,28 +39,30 @@ class Parameter: """A model parameter. A parameter is *independent* if it carries a `default` (or is `required`) and - no `formula`. Independent params are user-settable and packed into `self.m`. + no `formula`. Independent params are user-settable. A parameter is *dependent* if it has a `formula`: a Python expression that may reference other parameters (independent or dependent, as long as no - cycles). Dependents are NOT user-settable; they're computed at - `configure()` time and re-computed every time the user calls `configure()` - (i.e. every time params are modified — which is already the protocol the - existing neuronumba models follow). + cycles), the model-context names ``weights``, ``weights_t``, ``g``, + ``n_rois``, and ``np``. Dependents are NOT user-settable; they're computed + at ``configure()`` time and re-computed every time the user calls + ``configure()`` (i.e. every time params are modified). + + Whatever Python value the formula returns is what the parameter is — a + scalar, a per-region 1D array, or a 2D matrix all work. The dfun captures + the value by closure and numba's type inference handles it. There is no + explicit shape declaration; if the formula returns ``np.eye(n_rois)``, + the parameter is a matrix. Examples: Parameter("J_ee", default=10.0) Parameter("g_ee", default=2.5) Parameter("a_e", default=0.25) - Parameter("J_N_ee", formula="J_ee + g_ee * np.log(a_e)") # Montbrio-style - - Formulas may use any of: numpy (as `np`), the standard math operators, and - references to other parameter names. They evaluate in NumPy land (not numba) - so anything numpy supports works, including matrix ops. + Parameter("J_N_ee", formula="J_ee + g_ee * np.log(a_e)") # scalar dependent + Parameter("A", formula="(-g * weights + np.diag(weights.sum(axis=1))) / tau") # matrix dependent """ name: str default: Optional[float] = None - regional: bool = True required: bool = False doc: str = "" formula: Optional[str] = None @@ -97,3 +100,29 @@ class ModelSpec: parameters: List[Parameter] equations: str # multiline; assignments of intermediates and `d_` rates initial_state_overrides: Dict[str, float] = field(default_factory=dict) + + # Imperative escape hatch. Called as `on_configure(model)` once per + # `configure()`, AFTER dependent-parameter evaluation. Use it for setup + # that doesn't fit the declarative formula form: fsolve-based FIC, + # iterative steady-state computation, anything that needs Python control + # flow. Mutations to the model instance (e.g. `model.J = computed`) are + # visible to the dfun the next time `get_numba_dfun()` is called. + # + # For parameters set by this callback, declare them with a placeholder + # `default=` value; the callback overwrites it. The user can still + # supply an override at construction time, in which case the callback + # would also overwrite that — design accordingly. + on_configure: Optional[Callable[[Any], None]] = None + + # User-supplied @nb.njit functions usable from `equations`. Each must + # have a unique `__name__` that doesn't collide with state/coupling + # variables, parameter names, declared observables, whitelisted numpy + # functions, the reserved tokens (`state`, `coupling`, `model`, `np`), + # or any inherited `Model`/`LinearCouplingModel` method name. Helpers + # are stashed as class attributes on the generated class and referenced + # via `model.` in the dfun factory; numba sees the closure-captured + # callable and dispatches to it during compilation. + # + # Use case: models with custom subroutines (e.g. Zerlaut's + # `erfc_approx`, `threshold_func`, `get_fluct_regime_vars`). + helpers: List[Callable] = field(default_factory=list) From c2f312fd7551d7592a4caca650fd416162f2a6f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:20 +0200 Subject: [PATCH 29/48] feat(simulator): update configure to wire HistoryDelays with tract delays --- src/neuronumba/simulator/simulator.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/src/neuronumba/simulator/simulator.py b/src/neuronumba/simulator/simulator.py index 9b803ed..487c164 100644 --- a/src/neuronumba/simulator/simulator.py +++ b/src/neuronumba/simulator/simulator.py @@ -6,7 +6,7 @@ from neuronumba.numba_tools import address_as_void_pointer, addr from neuronumba.numba_tools.types import NDA_f8_2d, NDA_f8_1d from neuronumba.simulator.connectivity import Connectivity -from neuronumba.simulator.history import HistoryNoDelays +from neuronumba.simulator.history import HistoryDelays, HistoryNoDelays from neuronumba.simulator.integrators import EulerStochastic from neuronumba.simulator.monitors import RawSubSample, TemporalAverage from neuronumba.numba_tools.config import NUMBA_CACHE, NUMBA_FASTMATH, NUMBA_NOGIL @@ -34,7 +34,21 @@ def run(self, t_start=0, t_end=10000, stimulus=None): self.integrator.configure() self.connectivity.configure() self.model.configure(weights=self.connectivity.weights) - self.history.configure(c_vars=self.model.c_vars, weights=self.connectivity.weights) + if isinstance(self.history, HistoryDelays): + # Delays in seconds; the history converts to integer step indices + # using the integrator's dt. + delays = self.connectivity.lengths / self.connectivity.speed + self.history.configure( + c_vars=self.model.c_vars, + weights=self.connectivity.weights, + g=getattr(self.model, "g", 1.0), + delays=delays, + dt=self.integrator.dt, + ) + else: + self.history.configure( + c_vars=self.model.c_vars, weights=self.connectivity.weights + ) dt = self.integrator.dt t_max = t_end - t_start From 5b286e03bf1a121f453efc653fadc4422a7f5797 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:23 +0200 Subject: [PATCH 30/48] test(tests): update dependents tests to verify dfun closure capture --- tests/test_dsl_dependents.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_dsl_dependents.py b/tests/test_dsl_dependents.py index f3c3e2c..f756001 100644 --- a/tests/test_dsl_dependents.py +++ b/tests/test_dsl_dependents.py @@ -32,7 +32,6 @@ def test_dependents_simple_chain(): n = 4 m = Cls(g=0.0).configure(weights=np.zeros((n, n))) assert np.allclose(m.c, 5.0) - assert np.allclose(m.m[int(m.P.c)], 5.0) def test_dependents_multilevel_topo_sort(): @@ -115,4 +114,8 @@ def test_dependents_recompute_on_configure(): m.a = 5.0 m.configure(weights=W) assert np.allclose(m.b, 50.0) - assert np.allclose(m.m[int(m.P.b)], 50.0) + # Verify the dfun sees the updated dependent: -x / b at b=50, x=1 → -0.02. + f = m.get_numba_dfun() + state = np.ones((1, n)) + ds, _ = f(state, np.zeros((1, n))) + np.testing.assert_allclose(ds[0], -1.0 / 50.0, atol=1e-12) From a5ba031c170c8e08f6fd3d128a78d29e035088ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:27 +0200 Subject: [PATCH 31/48] test(tests): add delayed coupling spec and equivalence tests --- tests/test_dsl_delayed.py | 167 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 167 insertions(+) create mode 100644 tests/test_dsl_delayed.py diff --git a/tests/test_dsl_delayed.py b/tests/test_dsl_delayed.py new file mode 100644 index 0000000..8efabb8 --- /dev/null +++ b/tests/test_dsl_delayed.py @@ -0,0 +1,167 @@ +"""Tests for `kind="delayed"` coupling. + +Covers: + - DSL spec validation (build_model accepts kind="delayed", rejects mixed kinds). + - get_jacobian raises NotImplementedError for delayed kernels. + - End-to-end equivalence: with effectively-zero delays (huge speed), + HistoryDelays + DSL kind="delayed" produces the same trajectory as + HistoryNoDelays + DSL kind="linear" on the same spec. + - Sanity: with non-trivial delays, the trajectory differs. +""" +import numpy as np +import pytest + +from neuronumba.simulator.connectivity import Connectivity +from neuronumba.simulator.history import HistoryDelays, HistoryNoDelays +from neuronumba.simulator.integrators.euler import EulerDeterministic +from neuronumba.simulator.monitors import RawSubSample +from neuronumba.simulator.simulator import Simulator +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +def _delayed_hopf_spec(name): + return ModelSpec( + name=name, + state_vars=[StateVar("x", initial=0.1), StateVar("y", initial=0.1)], + coupling_vars=[ + CouplingVar("x", kind="delayed"), + CouplingVar("y", kind="delayed"), + ], + observables=[], + parameters=[ + Parameter("a", default=-0.5), + Parameter("omega", default=0.3), + Parameter("I_external", default=0.0), + ], + equations=""" + d_x = (a - x*x - y*y) * x - omega * y + coupling.x + I_external + d_y = (a - x*x - y*y) * y + omega * x + coupling.y + """, + ) + + +def _linear_hopf_spec(name): + """Same spec but with kind='linear' — for the zero-delay equivalence test.""" + return ModelSpec( + name=name, + state_vars=[StateVar("x", initial=0.1), StateVar("y", initial=0.1)], + coupling_vars=[ + CouplingVar("x", kind="linear"), + CouplingVar("y", kind="linear"), + ], + observables=[], + parameters=[ + Parameter("a", default=-0.5), + Parameter("omega", default=0.3), + Parameter("I_external", default=0.0), + ], + equations=""" + d_x = (a - x*x - y*y) * x - omega * y + coupling.x + I_external + d_y = (a - x*x - y*y) * y + omega * x + coupling.y + """, + ) + + +def _run(model, weights, lengths, speed, history, t_max=50.0, dt=0.1): + np.random.seed(0) + con = Connectivity(weights=weights, lengths=lengths, speed=speed) + integrator = EulerDeterministic(dt=dt) + monitor = RawSubSample(period=1.0, monitor_vars=model.get_var_info(["x"])) + sim = Simulator( + connectivity=con, model=model, history=history, + integrator=integrator, monitors=[monitor], + ) + sim.run(0, t_max) + return monitor.data("x") + + +def test_kind_delayed_builds(): + Cls = build_model(_delayed_hopf_spec("DelayedDSL_Build")) + assert "x" in Cls._coupling_var_names + assert Cls._coupling_var_kinds == ["delayed", "delayed"] + + +def test_kind_delayed_mixed_with_linear_rejected(): + spec = ModelSpec( + name="MixedKindsDSL", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[ + CouplingVar("x", kind="delayed"), + ], + observables=[], + parameters=[Parameter("a", default=1.0)], + equations="d_x = -x * a + coupling.x", + ) + # Single delayed cv is fine. + build_model(spec) + + # Now make it mixed: add a linear coupling on a second state var. + mixed = ModelSpec( + name="MixedKindsDSL2", + state_vars=[StateVar("x", initial=0.0), StateVar("y", initial=0.0)], + coupling_vars=[ + CouplingVar("x", kind="delayed"), + CouplingVar("y", kind="linear"), + ], + observables=[], + parameters=[Parameter("a", default=1.0)], + equations=""" + d_x = -x * a + coupling.x + d_y = -y + coupling.y + """, + ) + with pytest.raises(NotImplementedError, match="Mixing 'delayed'"): + build_model(mixed) + + +def test_jacobian_raises_for_delayed_coupling(weights): + Cls = build_model(_delayed_hopf_spec("DelayedDSL_Jac")) + m = Cls(g=0.5).configure(weights=weights) + state = np.zeros((2, weights.shape[0])) + with pytest.raises(NotImplementedError, match="not supported for 'delayed'"): + m.get_jacobian(state) + + +def test_zero_delay_equivalent_to_linear(weights): + """`speed=1e6` makes all delays round to zero. HistoryDelays + delayed-DSL + must then match HistoryNoDelays + linear-DSL, bit-for-bit.""" + n = weights.shape[0] + rng = np.random.default_rng(0) + lengths = rng.random((n, n)) * 10.0 + 1.0 + huge_speed = 1e6 # delays << dt → all i_delays round to 0 + + DelCls = build_model(_delayed_hopf_spec("DelayedDSL_ZeroDelay")) + LinCls = build_model(_linear_hopf_spec("LinearDSL_ZeroDelay")) + + delayed_traj = _run( + DelCls(g=0.5), weights, lengths, huge_speed, HistoryDelays(), + ) + linear_traj = _run( + LinCls(g=0.5), weights, lengths, huge_speed, HistoryNoDelays(), + ) + + np.testing.assert_allclose(delayed_traj, linear_traj, atol=1e-12, rtol=1e-12) + + +def test_nontrivial_delay_changes_trajectory(weights): + """With finite speed, the delayed trajectory must differ from no-delay.""" + n = weights.shape[0] + rng = np.random.default_rng(0) + lengths = rng.random((n, n)) * 10.0 + 1.0 + + DelCls = build_model(_delayed_hopf_spec("DelayedDSL_Real")) + + # speed=1.0 → delays of order 10 ms with our random lengths; with dt=0.1 ms + # that's ~100 steps of delay. Definitely not zero. + fast = _run( + DelCls(g=0.5), weights, lengths, speed=1e6, history=HistoryDelays(), + ) + slow = _run( + DelCls(g=0.5), weights, lengths, speed=1.0, history=HistoryDelays(), + ) + + # The two runs must be DIFFERENT — coupling timing changed. + diff = np.abs(fast - slow).max() + assert diff > 1e-6, f"expected delay to perturb trajectory, max|Δ|={diff}" From 35cc5344866e2a7085dc048f07165e31edd1f258 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:31 +0200 Subject: [PATCH 32/48] test(tests): add helper function callable from dfun tests --- tests/test_dsl_helpers.py | 142 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 tests/test_dsl_helpers.py diff --git a/tests/test_dsl_helpers.py b/tests/test_dsl_helpers.py new file mode 100644 index 0000000..040ba15 --- /dev/null +++ b/tests/test_dsl_helpers.py @@ -0,0 +1,142 @@ +"""Tests for `helpers=[...]` — user @nb.njit functions usable from DSL equations. + +This is the v0.2 unlock for Zerlaut-class models that need custom subroutines +(`erfc_approx`, threshold functions) the DSL itself doesn't provide. +""" +import numpy as np +import numba as nb +import pytest + +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +# A pair of @nb.njit helpers we'll reuse across tests. + +@nb.njit(nb.f8[:](nb.f8[:]), cache=False) +def soft_clip(x): + """Saturating nonlinearity, scalar→scalar in numba (vectorized over arr).""" + return np.tanh(x) + + +@nb.njit(nb.f8[:](nb.f8[:], nb.f8[:]), cache=False) +def gated_sum(a, b): + return a + 0.5 * b + + +def _spec_with_helpers(helpers, equations="d_x = -x + soft_clip(x) + coupling.x"): + return ModelSpec( + name="HelpersModel", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[Parameter("a", default=1.0)], + equations=equations, + helpers=helpers, + ) + + +def test_helper_callable_from_dfun(): + """The dfun should compile and produce sensible output when calling soft_clip.""" + Cls = build_model(_spec_with_helpers([soft_clip])) + n = 4 + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + + f = m.get_numba_dfun() + state = np.array([[0.5, 1.0, -0.5, 2.0]]) + coupling = np.zeros((1, n)) + ds, _ = f(state, coupling.copy()) + expected = -state[0] + np.tanh(state[0]) + np.testing.assert_allclose(ds[0], expected, atol=1e-12) + + +def test_helper_unused_is_silently_dropped(): + """Listing a helper that isn't referenced in equations is a no-op (not an error).""" + Cls = build_model(_spec_with_helpers([soft_clip, gated_sum])) + n = 3 + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + f = m.get_numba_dfun() + state = np.array([[0.0, 0.5, 1.0]]) + ds, _ = f(state, np.zeros((1, n))) + np.testing.assert_allclose(ds[0], -state[0] + np.tanh(state[0]), atol=1e-12) + + +def test_two_helpers_referenced(): + """Multiple helpers in one equation.""" + spec = _spec_with_helpers( + [soft_clip, gated_sum], + equations="d_x = -x + gated_sum(x, soft_clip(x)) + coupling.x", + ) + Cls = build_model(spec) + n = 3 + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + f = m.get_numba_dfun() + state = np.array([[0.5, 1.0, -0.5]]) + ds, _ = f(state, np.zeros((1, n))) + expected = -state[0] + state[0] + 0.5 * np.tanh(state[0]) + np.testing.assert_allclose(ds[0], expected, atol=1e-12) + + +def test_helper_unknown_name_in_equation_raises(): + """Calling a name that isn't a helper, param, or np func should still error.""" + with pytest.raises(ValueError, match="unknown identifier"): + build_model(_spec_with_helpers( + [soft_clip], + equations="d_x = -x + not_a_helper(x) + coupling.x", + )) + + +def test_helper_name_collision_with_param_rejected(): + @nb.njit + def a(z): # collides with parameter "a" + return z + + with pytest.raises(ValueError, match="clash with"): + build_model(_spec_with_helpers([a])) + + +def test_helper_name_collision_with_state_var_rejected(): + @nb.njit + def x(z): # collides with state var "x" + return z + + with pytest.raises(ValueError, match="clash with"): + build_model(_spec_with_helpers([x])) + + +def test_helper_name_collision_with_np_func_rejected(): + @nb.njit + def exp(z): # collides with whitelisted np.exp + return z + + with pytest.raises(ValueError, match="clash with"): + build_model(_spec_with_helpers([exp])) + + +def test_helper_duplicate_names_rejected(): + @nb.njit + def helper(z): + return z + + # Same helper twice — same __name__, same identity, but the validator + # objects regardless because the source-generation can't distinguish. + with pytest.raises(ValueError, match="must be unique"): + build_model(_spec_with_helpers([helper, helper])) + + +def test_helper_lambda_rejected(): + """Lambdas have __name__ == ''; we reject them upfront.""" + f = lambda z: z # noqa: E731 + f.__name__ = "" # simulate a missing/empty name + with pytest.raises(ValueError, match="no __name__"): + build_model(_spec_with_helpers([f])) + + +def test_helpers_not_listed_for_specs_without_them(): + """Specs with no helpers continue to work exactly as before.""" + Cls = build_model(_spec_with_helpers([], equations="d_x = -x * a + coupling.x")) + m = Cls(g=0.0).configure(weights=np.zeros((4, 4))) + f = m.get_numba_dfun() + ds, _ = f(np.array([[0.5, 1.0, -0.5, 2.0]]), np.zeros((1, 4))) + np.testing.assert_allclose(ds[0], -np.array([0.5, 1.0, -0.5, 2.0]), atol=1e-12) From 0c243d994961f1c7bcfb8680baa19f1b2cee1db2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:35 +0200 Subject: [PATCH 33/48] test(tests): add numerical jacobian shape and correctness tests --- tests/test_dsl_jacobian.py | 150 +++++++++++++++++++++++++++++++++++++ 1 file changed, 150 insertions(+) create mode 100644 tests/test_dsl_jacobian.py diff --git a/tests/test_dsl_jacobian.py b/tests/test_dsl_jacobian.py new file mode 100644 index 0000000..5bd780c --- /dev/null +++ b/tests/test_dsl_jacobian.py @@ -0,0 +1,150 @@ +"""Tests for the DSL's numerical Jacobian. + +The DSL builds the network Jacobian in two pieces: + 1. Local partials of the dfun w.r.t. each state var and each coupling var, + via centered finite differences. + 2. Closed-form linearization of the coupling kernel (linear / diffusive). + +These tests verify both pieces by: + - Comparing the structured DSL Jacobian against a "naive" full-network + finite-difference reference that doesn't exploit locality (catches + assembly bugs). + - Comparing the DSL Jacobian for Hopf at the origin against the + hand-written analytical Jacobian (catches both math and assembly bugs). +""" +import numpy as np +import pytest + + +def _naive_full_network_fd_jacobian(model, state, eps=1e-6): + """Reference: brute-force finite difference treating dfun(state)+coupling + as one big function of the flat state. Slow, but free of structural + assumptions.""" + nsv, N = state.shape + c_vars = list(model.c_vars) + dfun = model.get_numba_dfun() + coupling_factory = model.get_numba_coupling() + + def total(flat_state): + s = flat_state.reshape(nsv, N) + c = coupling_factory(s[c_vars, :].copy()) + ds, _ = dfun(s.copy(), c) + return ds.ravel() + + flat = state.ravel().copy() + base = total(flat) + J = np.empty((flat.size, flat.size)) + for k in range(flat.size): + x_plus = flat.copy(); x_plus[k] += eps + x_minus = flat.copy(); x_minus[k] -= eps + J[:, k] = (total(x_plus) - total(x_minus)) / (2.0 * eps) + return J + + +def test_jacobian_shape(weights, n_rois, hopf_dsl_cls): + m = hopf_dsl_cls(g=0.5).configure(weights=weights) + nsv = m.n_state_vars + state = np.zeros((nsv, n_rois)) + J = m.get_jacobian(state) + assert J.shape == (nsv * n_rois, nsv * n_rois) + + +def test_jacobian_rejects_wrong_shape(weights, hopf_dsl_cls): + m = hopf_dsl_cls(g=0.5).configure(weights=weights) + with pytest.raises(ValueError, match="state must have shape"): + m.get_jacobian(np.zeros((3, 5))) # wrong shape + + +def test_jacobian_hopf_at_origin_matches_analytic(weights, hopf_dsl_cls): + """DSL Jacobian at the origin matches the closed-form derivation. + + Hopf dfun at the origin (x=y=0): + d(d_x)/dx = a + d(coupling.x)/dx (the cubic terms vanish) + d(d_x)/dy = -omega + d(d_y)/dx = +omega + d(d_y)/dy = same as d(d_x)/dx + Diffusive coupling linearization: d(coupling.x[i])/dx[j] = g*(W[i,j] - delta_ij*ink[i]). + + Note: we deliberately build the reference here rather than calling + `Hopf.get_jacobian`. The hand-written method has an internal factor-of-2pi + inconsistency between its dfun (uses `omega`) and its analytical Jacobian + (uses `2*pi*omega`), so it can't serve as a faithful oracle for any + Jacobian that's consistent with the dfun. + """ + N = weights.shape[0] + g = 0.5 + a = -0.5 + omega = 0.3 + + # Diffusive-coupling linearization for Hopf's coupling vars. + ink = weights.T.sum(axis=1) + C = g * (weights - np.diag(ink)) + + # Closed-form blocks. + block_diag = a * np.eye(N) + C # (x->x) and (y->y) + block_off_xy = -omega * np.eye(N) # (x<-y) cross + block_off_yx = +omega * np.eye(N) # (y<-x) cross + J_expected = np.block([ + [block_diag, block_off_xy], + [block_off_yx, block_diag], + ]) + + m_dsl = hopf_dsl_cls(g=g, a=a, omega=omega).configure(weights=weights) + J_dsl = m_dsl.get_jacobian(np.zeros((2, N))) + + np.testing.assert_allclose(J_dsl, J_expected, atol=1e-7, rtol=1e-7) + + +def test_jacobian_structured_matches_naive_fd_hopf(weights, hopf_dsl_cls): + """Structured DSL Jacobian == brute-force network-wide FD (Hopf, random state).""" + m = hopf_dsl_cls(g=0.5).configure(weights=weights) + rng = np.random.default_rng(123) + state = rng.standard_normal((2, weights.shape[0])) * 0.05 # small to keep regime stable + + J_struct = m.get_jacobian(state) + J_naive = _naive_full_network_fd_jacobian(m, state) + np.testing.assert_allclose(J_struct, J_naive, atol=1e-7, rtol=1e-7) + + +def test_jacobian_structured_matches_naive_fd_deco(weights, deco_dsl_cls): + """Same check on Deco2014 — exercises linear coupling and observables.""" + m = deco_dsl_cls(g=0.5).configure(weights=weights) + rng = np.random.default_rng(7) + # Stay inside the model's bounds (S_e, S_i in [0, 1]). + state = rng.uniform(0.1, 0.5, size=(2, weights.shape[0])) + + J_struct = m.get_jacobian(state) + J_naive = _naive_full_network_fd_jacobian(m, state) + # Deco's sigmoid has steep regions; bump tolerance slightly. + np.testing.assert_allclose(J_struct, J_naive, atol=1e-5, rtol=1e-5) + + +def test_jacobian_structured_matches_naive_fd_naskar(weights, naskar_dsl_cls): + """Three-state-var sanity check: Naskar2021 (S_e, S_i, J).""" + m = naskar_dsl_cls(g=0.5).configure(weights=weights) + rng = np.random.default_rng(11) + state = np.empty((3, weights.shape[0])) + state[0] = rng.uniform(0.1, 0.5, size=weights.shape[0]) # S_e + state[1] = rng.uniform(0.1, 0.5, size=weights.shape[0]) # S_i + state[2] = rng.uniform(0.5, 1.5, size=weights.shape[0]) # J + + J_struct = m.get_jacobian(state) + J_naive = _naive_full_network_fd_jacobian(m, state) + np.testing.assert_allclose(J_struct, J_naive, atol=1e-5, rtol=1e-5) + + +def test_jacobian_recomputes_after_param_change(weights, hopf_dsl_cls): + """Changing `a` and reconfiguring must change the Jacobian's diagonal.""" + m = hopf_dsl_cls(g=0.5).configure(weights=weights) + state = np.zeros((2, weights.shape[0])) + J1 = m.get_jacobian(state) + + m.a = 0.5 + m.configure(weights=weights) + J2 = m.get_jacobian(state) + + # The diagonal of the (x, x) block contains a + (coupling diagonal). + # Changing a from -0.5 to +0.5 shifts those entries by exactly +1.0. + N = weights.shape[0] + diag_diff = np.diag(J2[:N, :N]) - np.diag(J1[:N, :N]) + np.testing.assert_allclose(diag_diff, 1.0, atol=1e-7) From 8d691a6eefc06ff3fe272a669df63409a40a1362 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:39 +0200 Subject: [PATCH 34/48] test(tests): add matrix-shaped dependent parameter closure tests --- tests/test_dsl_matrix_params.py | 159 ++++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 tests/test_dsl_matrix_params.py diff --git a/tests/test_dsl_matrix_params.py b/tests/test_dsl_matrix_params.py new file mode 100644 index 0000000..e2f3788 --- /dev/null +++ b/tests/test_dsl_matrix_params.py @@ -0,0 +1,159 @@ +"""Tests for matrix-shaped dependent parameters. + +After the closure-capture refactor, there's no special "matrix" marker — a +formula that returns a 2D ndarray simply produces a matrix-valued parameter. +The DSL captures it via `model.` like any other parameter; numba's +type inference handles `A @ x` style expressions in the dfun. + +Use case: OrnsteinUhlenbeck-style models where the dynamics need a full +(n_rois, n_rois) operator derived from `weights`/`g`/`tau`. +""" +import numpy as np +import numba as nb + +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +def test_matrix_param_computed_and_visible_in_dfun(): + """Build a model whose dfun reads a 2D matrix via closure capture.""" + n = 4 + spec = ModelSpec( + name="MatrixModel", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[ + Parameter("scale", default=2.0), + # A is the (n_rois, n_rois) identity matrix scaled by `scale`. + Parameter("A", formula="scale * np.eye(weights.shape[0])"), + ], + equations="d_x = -x + A @ x + coupling.x", + ) + Cls = build_model(spec) + W = np.zeros((n, n)) + m = Cls(g=0.0).configure(weights=W) + + # The matrix lives on `self`. + assert m.A.shape == (n, n) + np.testing.assert_allclose(m.A, 2.0 * np.eye(n)) + + # The dfun should compile and use the matrix correctly. + f = m.get_numba_dfun() + state = np.array([[0.5, 1.0, -0.5, 2.0]]) + ds, _ = f(state, np.zeros((1, n))) + # d_x = -x + 2.0 * I @ x = -x + 2*x = x (in this trivial case) + np.testing.assert_allclose(ds[0], state[0], atol=1e-12) + + +def test_matrix_param_lives_on_instance(): + """A 2D dependent param is just an instance attribute — no special storage.""" + n = 3 + spec = ModelSpec( + name="MatrixOnInstance", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[ + Parameter("alpha", default=1.0), + Parameter("M", formula="np.eye(weights.shape[0])"), + ], + equations="d_x = alpha * (M @ x) + coupling.x", + ) + Cls = build_model(spec) + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + + assert isinstance(m.M, np.ndarray) + assert m.M.shape == (n, n) + np.testing.assert_allclose(m.M, np.eye(n)) + + +def test_matrix_param_recomputes_on_configure(): + """Like other dependents, matrix params re-evaluate when configure() runs.""" + n = 4 + spec = ModelSpec( + name="MatrixModel_Recompute", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[ + Parameter("k", default=1.0), + Parameter("A", formula="k * np.eye(weights.shape[0])"), + ], + equations="d_x = -(A @ x) + coupling.x", + ) + Cls = build_model(spec) + W = np.zeros((n, n)) + m = Cls(g=0.0).configure(weights=W) + np.testing.assert_allclose(m.A, np.eye(n)) + + m.k = 5.0 + m.configure(weights=W) + np.testing.assert_allclose(m.A, 5.0 * np.eye(n)) + + +def test_matrix_param_OU_like(): + """Smoke test of the OU-style use case: A derived from weights/g/tau. + + OU dynamics: dx/dt = -A x with A = (-W * g + diag(rowsum)) / tau. + """ + n = 5 + rng = np.random.default_rng(0) + W = rng.random((n, n)); W = (W + W.T) / 2; np.fill_diagonal(W, 0) + + spec = ModelSpec( + name="OULike", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[ + Parameter("tau", default=10.0), + Parameter( + "A", + formula="(-g * weights + np.diag(weights.sum(axis=1))) / tau", + ), + ], + # Ignore the linear-coupling output here; the matrix A captures the + # full network operator. + equations="d_x = -(A @ x)", + ) + Cls = build_model(spec) + m = Cls(g=0.5).configure(weights=W) + + expected_A = (-0.5 * W + np.diag(W.sum(axis=1))) / 10.0 + np.testing.assert_allclose(m.A, expected_A, atol=1e-12) + + f = m.get_numba_dfun() + state = rng.standard_normal((1, n)) * 0.1 + ds, _ = f(state, np.zeros((1, n))) + expected = -(expected_A @ state[0]) + np.testing.assert_allclose(ds[0], expected, atol=1e-10) + + +def test_matrix_param_can_combine_with_helpers(): + """Helpers and matrix params should compose without interference.""" + n = 4 + + @nb.njit(nb.f8[:](nb.f8[:]), cache=False) + def squash(x): + return np.tanh(x) + + spec = ModelSpec( + name="MatrixHelperCombo", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[ + Parameter("k", default=2.0), + Parameter("A", formula="k * np.eye(weights.shape[0])"), + ], + equations="d_x = squash(A @ x) + coupling.x", + helpers=[squash], + ) + Cls = build_model(spec) + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + f = m.get_numba_dfun() + state = np.array([[0.1, 0.2, 0.3, 0.4]]) + ds, _ = f(state, np.zeros((1, n))) + np.testing.assert_allclose(ds[0], np.tanh(2.0 * state[0]), atol=1e-12) From 92b8cca8590e581d834d5714e2368c5326080a5c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 19:09:43 +0200 Subject: [PATCH 35/48] test(tests): add on_configure escape hatch and fic use case tests --- tests/test_dsl_on_configure.py | 164 +++++++++++++++++++++++++++++++++ 1 file changed, 164 insertions(+) create mode 100644 tests/test_dsl_on_configure.py diff --git a/tests/test_dsl_on_configure.py b/tests/test_dsl_on_configure.py new file mode 100644 index 0000000..421aaa9 --- /dev/null +++ b/tests/test_dsl_on_configure.py @@ -0,0 +1,164 @@ +"""Tests for the `on_configure` imperative escape hatch. + +Covers: + - The callback fires on configure() and sees current weights/g. + - Mutations land in self.m (the parameter matrix the dfun reads). + - The callback runs AFTER dependent-parameter evaluation. + - Re-configure re-runs the callback (matches the dependents protocol). + - Real FIC use case: FICHerzog2022 computing J for a Deco-style spec. +""" +import numpy as np +import pytest + +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) + + +def _basic_spec(parameters, on_configure=None, equations="d_x = -x * a + coupling.x"): + return ModelSpec( + name="ConfModel", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=parameters, + equations=equations, + on_configure=on_configure, + ) + + +def test_on_configure_fires_with_self(): + captured = {} + def hook(self): + captured["weights_shape"] = self.weights.shape + captured["g"] = self.g + + spec = _basic_spec([Parameter("a", default=1.0)], on_configure=hook) + Cls = build_model(spec) + n = 4 + Cls(g=0.7).configure(weights=np.eye(n)) + assert captured["weights_shape"] == (n, n) + assert captured["g"] == 0.7 + + +def test_on_configure_mutation_persists_on_instance(): + """Setting self. in the hook persists on the instance and is + visible to the dfun closure on the next get_numba_dfun() call.""" + def hook(self): + # Compute J from weights/g — toy version of FIC. + self.J = 2.0 * self.g * self.weights.sum(axis=0) + 1.0 + + spec = _basic_spec( + [Parameter("a", default=1.0), Parameter("J", default=0.0)], + on_configure=hook, + ) + Cls = build_model(spec) + n = 5 + W = np.ones((n, n)) * 0.5 + np.fill_diagonal(W, 0) + m = Cls(g=0.3).configure(weights=W) + + expected_J = 2.0 * 0.3 * W.sum(axis=0) + 1.0 + np.testing.assert_allclose(m.J, expected_J) + # Confirm the dfun closure sees the hook-set value (build compiles cleanly). + m.get_numba_dfun() + + +def test_on_configure_runs_after_dependents(): + """The hook can read computed dependent params (they're set first).""" + captured = {} + def hook(self): + captured["b_at_hook_time"] = float(np.atleast_1d(self.b)[0]) + + spec = _basic_spec( + [ + Parameter("a", default=3.0), + Parameter("b", formula="a * 100.0"), # dependent + ], + on_configure=hook, + ) + Cls = build_model(spec) + Cls(g=0.0).configure(weights=np.zeros((2, 2))) + assert captured["b_at_hook_time"] == 300.0 + + +def test_on_configure_overrides_dependents(): + """Whatever the hook sets is final — it runs AFTER dependents.""" + def hook(self): + self.b = 999.0 + + spec = _basic_spec( + [ + Parameter("a", default=1.0), + Parameter("b", formula="a * 100.0"), + ], + on_configure=hook, + ) + Cls = build_model(spec) + m = Cls(g=0.0).configure(weights=np.zeros((2, 2))) + np.testing.assert_allclose(m.b, 999.0) + + +def test_on_configure_reruns_on_reconfigure(): + """Same protocol as dependent params: every configure() re-runs the hook.""" + counter = {"n": 0} + def hook(self): + counter["n"] += 1 + self.J = float(counter["n"]) * 10.0 + + spec = _basic_spec( + [Parameter("a", default=1.0), Parameter("J", default=0.0)], + on_configure=hook, + ) + Cls = build_model(spec) + W = np.zeros((3, 3)) + m = Cls(g=0.0).configure(weights=W) + assert counter["n"] == 1 + np.testing.assert_allclose(m.J, 10.0) + + m.configure(weights=W) + assert counter["n"] == 2 + np.testing.assert_allclose(m.J, 20.0) + + +def test_on_configure_with_fic_herzog(): + """Canonical use case: imperative FIC computation that doesn't fit a formula. + + `FICHerzog2022.compute_J(sc, g)` returns a per-region J. We stash the + result on the model and verify the dfun sees it via self.m. + """ + from neuronumba.fitting.fic.fic import FICHerzog2022 + + def auto_fic(self): + self.J = FICHerzog2022().compute_J(self.weights, self.g) + + spec = ModelSpec( + name="MiniDeco", + state_vars=[StateVar("S_e", initial=0.001, bounds=(0.0, 1.0))], + coupling_vars=[CouplingVar("S_e", kind="linear")], + observables=[], + parameters=[ + Parameter("tau_e", default=100.0), + Parameter("J", default=0.0), # placeholder; FIC overwrites + ], + equations="d_S_e = -S_e / tau_e + J * coupling.S_e", + on_configure=auto_fic, + ) + Cls = build_model(spec) + + rng = np.random.default_rng(0) + W = rng.random((6, 6)); W = (W + W.T) / 2; np.fill_diagonal(W, 0) + m = Cls(g=0.4).configure(weights=W) + + expected_J = 0.75 * 0.4 * W.sum(axis=0) + 1.0 # FICHerzog2022 closed form + np.testing.assert_allclose(m.J, expected_J) + # Confirm the dfun closure sees the hook-set value (build compiles cleanly). + m.get_numba_dfun() + + +def test_on_configure_none_keeps_old_behavior(): + """When the hook is unset, configure() behaves exactly as before.""" + spec = _basic_spec([Parameter("a", default=2.0)]) # on_configure=None + Cls = build_model(spec) + m = Cls(g=0.5).configure(weights=np.eye(3)) + np.testing.assert_allclose(m.a, 2.0) From a9c6a71f5f611d526d4943eec5acdacbd6841d04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 21:42:30 +0200 Subject: [PATCH 36/48] docs(dsl_models): update readme with loc counts and builder api docs --- examples/dsl_models/README.md | 43 ++++++++++++++++++++++++++--------- 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/examples/dsl_models/README.md b/examples/dsl_models/README.md index bdce7c7..2cce097 100644 --- a/examples/dsl_models/README.md +++ b/examples/dsl_models/README.md @@ -6,10 +6,14 @@ hand-written counterpart in `src/neuronumba/simulator/models/`. | Spec | Hand-written | LOC ratio | |---|---|---| -| [hopf_dsl.py](hopf_dsl.py) | `hopf.py` | ~30 vs 152 | -| [deco2014_dsl.py](deco2014_dsl.py) | `deco2014.py` | ~50 vs 416 | -| [naskar2021_dsl.py](naskar2021_dsl.py) | `naskar2021.py` | ~60 vs 104 | -| [montbrio_dsl.py](montbrio_dsl.py) | `montbrio.py` | ~70 vs 160 | +| [hopf_dsl.py](hopf_dsl.py) | `hopf.py` | 38 vs 152 | +| [deco2014_dsl.py](deco2014_dsl.py) | `deco2014.py` | 67 vs 416 | +| [naskar2021_dsl.py](naskar2021_dsl.py) | `naskar2021.py` | 69 vs 104 | +| [montbrio_dsl.py](montbrio_dsl.py) | `montbrio.py` | 88 vs 160 | + +There is also [compare_simulations.py](compare_simulations.py): a runnable +script that simulates pairs of DSL and hand-written models on the same inputs +and plots their trajectories side-by-side. ## Dual purpose @@ -31,15 +35,16 @@ If you change a DSL spec (e.g. for a new feature you want to add), the equivalence test still has to pass against the hand-written reference. If they genuinely diverge, that's a model change and belongs in a separate PR. -## Models not included +## Models not yet ported -- **Zerlaut** — needs `@njit` helper subroutines (`get_fluct_regime_vars`, - `threshold_func`, `erfc_approx`). The DSL doesn't support that yet — see the - Phase 6 deferred items in `IMPLEMENTATION_PLAN.md`. -- **OrnsteinUhlenbeck** — needs non-scalar matrix parameters (`A` derived from - weights/g/tau). Deferred to v0.2. +- **Zerlaut** (`zerlaut.py`, ~800 LOC) — uses several `@njit` helper subroutines + (`get_fluct_regime_vars`, `threshold_func`, `erfc_approx`). The DSL now + supports user-supplied helpers via `helpers=[...]`, so a port is feasible; + it just hasn't been written yet. -For these, subclass `Model` directly. +The DSL can also express models with non-scalar matrix parameters and +conduction-delay coupling — see `tests/test_dsl_matrix_params.py` and +`tests/test_dsl_delayed.py` for the relevant patterns. ## Usage @@ -52,3 +57,19 @@ HopfDSL = build_model(hopf_spec) m = HopfDSL(g=0.5).configure(weights=W) ``` + +For incremental construction (instead of writing a full `ModelSpec` literal), +use `ModelBuilder`: + +```python +from neuronumba.simulator.models.dsl import ModelBuilder + +HopfDSL = (ModelBuilder("Hopf") + .add_state("x").add_state("y") + .add_coupling("x", kind="diffusive") + .add_param("a", default=-0.5) + .add_param("omega", default=0.3) + .add_equation("d_x = (a - x*x - y*y)*x - omega*y + coupling.x") + .add_equation("d_y = (a - x*x - y*y)*y + omega*x") + .build()) +``` From 7f7a4503df73275f68704a22a7d640fe36fdfd6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 21:42:35 +0200 Subject: [PATCH 37/48] feat(dsl): update public api to export modelbuilder class --- src/neuronumba/simulator/models/dsl/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/neuronumba/simulator/models/dsl/__init__.py b/src/neuronumba/simulator/models/dsl/__init__.py index a2e9bef..f3ebc8b 100644 --- a/src/neuronumba/simulator/models/dsl/__init__.py +++ b/src/neuronumba/simulator/models/dsl/__init__.py @@ -8,15 +8,18 @@ Public API: ModelSpec, StateVar, CouplingVar, Parameter + ModelBuilder build_model, dump_generated, get_source_file get_cache_dir, cleanup_cache """ from .builder import build_model, dump_generated, get_source_file from .materialize import cleanup_cache, get_cache_dir +from .model_builder import ModelBuilder from .spec import CouplingVar, ModelSpec, Parameter, StateVar __all__ = [ "CouplingVar", + "ModelBuilder", "ModelSpec", "Parameter", "StateVar", From c29ab686d8e9a9d3192730b9445c9a47c6d1be9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 21:42:38 +0200 Subject: [PATCH 38/48] feat(dsl): add fluent incremental builder for modelspec assembly --- .../simulator/models/dsl/model_builder.py | 146 ++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 src/neuronumba/simulator/models/dsl/model_builder.py diff --git a/src/neuronumba/simulator/models/dsl/model_builder.py b/src/neuronumba/simulator/models/dsl/model_builder.py new file mode 100644 index 0000000..5d22a68 --- /dev/null +++ b/src/neuronumba/simulator/models/dsl/model_builder.py @@ -0,0 +1,146 @@ +""" +Incremental builder for DSL ModelSpec. + +`ModelBuilder` lets you assemble a model piece by piece — adding state vars, +coupling vars, parameters, observables, helpers, and equations one call at a +time — instead of constructing a full `ModelSpec` literal up front. Every +mutating method returns `self`, so calls can be chained or run imperatively. + +Two terminal methods: + - `.spec()` → returns the assembled `ModelSpec` (no compilation). + - `.build()` → compiles via `build_model(spec)` and returns the model class. +""" +from __future__ import annotations + +from typing import Callable, Dict, List, Optional, Tuple + +from .builder import build_model +from .spec import CouplingVar, ModelSpec, Parameter, StateVar + + +class ModelBuilder: + """Fluent, incremental builder for `ModelSpec`. + + Example: + Cls = (ModelBuilder("Hopf") + .add_state("x", initial=0.0) + .add_state("y", initial=0.0) + .add_coupling("x", kind="diffusive") + .add_param("a", default=-0.5) + .add_param("omega", default=0.3) + .add_equation("d_x = (a - x*x - y*y)*x - omega*y + coupling.x") + .add_equation("d_y = (a - x*x - y*y)*y + omega*x") + .build()) + """ + + def __init__(self, name: str): + if not isinstance(name, str) or not name: + raise ValueError("ModelBuilder requires a non-empty `name` string") + self._name: str = name + self._states: List[StateVar] = [] + self._couplings: List[CouplingVar] = [] + self._params: List[Parameter] = [] + self._observables: List[str] = [] + self._helpers: List[Callable] = [] + self._equations: List[str] = [] + self._initial_overrides: Dict[str, float] = {} + self._on_configure: Optional[Callable] = None + + def add_state( + self, + name: str, + initial: float = 0.0, + bounds: Optional[Tuple[float, float]] = None, + ) -> "ModelBuilder": + if any(s.name == name for s in self._states): + raise ValueError(f"Duplicate state variable: {name!r}") + self._states.append(StateVar(name=name, initial=initial, bounds=bounds)) + return self + + def add_coupling(self, name: str, kind: str = "linear") -> "ModelBuilder": + if any(c.name == name for c in self._couplings): + raise ValueError(f"Duplicate coupling variable: {name!r}") + self._couplings.append(CouplingVar(name=name, kind=kind)) + return self + + def add_param( + self, + name: str, + default: Optional[float] = None, + formula: Optional[str] = None, + required: bool = False, + doc: str = "", + ) -> "ModelBuilder": + if any(p.name == name for p in self._params): + raise ValueError(f"Duplicate parameter: {name!r}") + self._params.append(Parameter( + name=name, + default=default, + formula=formula, + required=required, + doc=doc, + )) + return self + + def add_observable(self, name: str) -> "ModelBuilder": + if name in self._observables: + raise ValueError(f"Duplicate observable: {name!r}") + self._observables.append(name) + return self + + def add_helper(self, fn: Callable) -> "ModelBuilder": + if not callable(fn): + raise TypeError("add_helper expects a callable") + if any(getattr(h, "__name__", None) == getattr(fn, "__name__", None) + for h in self._helpers): + raise ValueError(f"Helper {getattr(fn, '__name__', fn)!r} already added") + self._helpers.append(fn) + return self + + def add_equation(self, line: str) -> "ModelBuilder": + """Append one line to the equations block.""" + if not isinstance(line, str) or not line.strip(): + raise ValueError("add_equation expects a non-empty string") + self._equations.append(line.rstrip()) + return self + + def set_equations(self, source: str) -> "ModelBuilder": + """Replace the equations block wholesale (for multi-line literals).""" + if not isinstance(source, str): + raise TypeError("set_equations expects a string") + self._equations = [source] + return self + + def override_initial(self, name: str, value: float) -> "ModelBuilder": + self._initial_overrides[name] = float(value) + return self + + def on_configure(self, fn: Callable) -> "ModelBuilder": + if not callable(fn): + raise TypeError("on_configure expects a callable") + self._on_configure = fn + return self + + def spec(self) -> ModelSpec: + if not self._states: + raise ValueError( + "Model has no state variables; call add_state() at least once" + ) + if not self._equations: + raise ValueError( + "Model has no equations; call add_equation() or set_equations()" + ) + return ModelSpec( + name=self._name, + state_vars=list(self._states), + coupling_vars=list(self._couplings), + observables=list(self._observables), + parameters=list(self._params), + equations="\n".join(self._equations), + initial_state_overrides=dict(self._initial_overrides), + on_configure=self._on_configure, + helpers=list(self._helpers), + ) + + def build(self): + return build_model(self.spec()) From 0a8a47145722aa190333ffb97ca04066e62d6e32 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 21:42:42 +0200 Subject: [PATCH 39/48] test(tests): add modelbuilder chaining and imperative style tests --- tests/test_dsl_builder.py | 232 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 232 insertions(+) create mode 100644 tests/test_dsl_builder.py diff --git a/tests/test_dsl_builder.py b/tests/test_dsl_builder.py new file mode 100644 index 0000000..2b6ef23 --- /dev/null +++ b/tests/test_dsl_builder.py @@ -0,0 +1,232 @@ +"""Tests for `ModelBuilder` — incremental, fluent construction of DSL models.""" +import numpy as np +import numba as nb +import pytest + +from neuronumba.simulator.models.dsl import ( + ModelBuilder, ModelSpec, build_model, +) + + +def test_chainable_build_matches_modelspec(): + """A chained builder produces the same compiled model as a literal ModelSpec.""" + Cls_b = (ModelBuilder("Hopf") + .add_state("x", initial=0.0) + .add_state("y", initial=0.0) + .add_coupling("x", kind="diffusive") + .add_param("a", default=-0.5) + .add_param("omega", default=0.3) + .add_equation("d_x = (a - x*x - y*y)*x - omega*y + coupling.x") + .add_equation("d_y = (a - x*x - y*y)*y + omega*x") + .build()) + + n = 4 + W = np.zeros((n, n)) + m_b = Cls_b(g=0.0).configure(weights=W) + f_b = m_b.get_numba_dfun() + + state = np.array([[0.1, 0.2, -0.3, 0.4], + [0.5, -0.1, 0.2, 0.0]]) + coup = np.zeros((1, n)) + ds_b, _ = f_b(state, coup) + + # Hand-derived expected derivatives. + a, omega = -0.5, 0.3 + r2 = state[0]**2 + state[1]**2 + expected_dx = (a - r2) * state[0] - omega * state[1] + expected_dy = (a - r2) * state[1] + omega * state[0] + np.testing.assert_allclose(ds_b[0], expected_dx, atol=1e-12) + np.testing.assert_allclose(ds_b[1], expected_dy, atol=1e-12) + + +def test_imperative_style_works(): + """Same builder used imperatively (no chaining) produces an equivalent class.""" + b = ModelBuilder("Imp") + b.add_state("x", initial=0.0) + b.add_coupling("x", kind="linear") + b.add_param("k", default=2.0) + b.add_equation("d_x = -k * x + coupling.x") + Cls = b.build() + + m = Cls(g=0.0).configure(weights=np.zeros((3, 3))) + f = m.get_numba_dfun() + state = np.array([[1.0, 2.0, 3.0]]) + ds, _ = f(state, np.zeros((1, 3))) + np.testing.assert_allclose(ds[0], -2.0 * state[0], atol=1e-12) + + +def test_spec_returns_modelspec_without_compiling(): + """`.spec()` should return a ModelSpec without going through build_model.""" + spec = (ModelBuilder("S") + .add_state("x") + .add_coupling("x") + .add_param("a", default=1.0) + .add_equation("d_x = -a*x + coupling.x") + .spec()) + assert isinstance(spec, ModelSpec) + assert spec.name == "S" + assert len(spec.state_vars) == 1 + assert len(spec.parameters) == 1 + assert "d_x = -a*x + coupling.x" in spec.equations + + +def test_set_equations_replaces_block(): + """`.set_equations(...)` replaces any previously-added equation lines.""" + b = (ModelBuilder("S") + .add_state("x") + .add_coupling("x") + .add_param("a", default=1.0) + .add_equation("d_x = a*x") # will be discarded + .set_equations("d_x = -a*x + coupling.x")) + spec = b.spec() + assert spec.equations == "d_x = -a*x + coupling.x" + + +def test_helper_and_on_configure_wire_through(): + """`add_helper` and `on_configure` propagate to the spec.""" + @nb.njit(nb.f8[:](nb.f8[:]), cache=False) + def squash(x): + return np.tanh(x) + + captured = {"hit": False} + def hook(self): + captured["hit"] = True + + Cls = (ModelBuilder("WithExtras") + .add_state("x", initial=0.0) + .add_coupling("x") + .add_param("k", default=1.5) + .add_helper(squash) + .on_configure(hook) + .add_equation("d_x = squash(k * x) + coupling.x") + .build()) + + n = 3 + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + assert captured["hit"] is True + f = m.get_numba_dfun() + state = np.array([[0.1, 0.2, 0.3]]) + ds, _ = f(state, np.zeros((1, n))) + np.testing.assert_allclose(ds[0], np.tanh(1.5 * state[0]), atol=1e-12) + + +def test_add_observable_wires_through(): + """`add_observable` is propagated; an unassigned observable is rejected at build.""" + # Happy path: observable assigned in equations. + Cls = (ModelBuilder("Obs") + .add_state("x", initial=0.0) + .add_coupling("x") + .add_param("a", default=1.0) + .add_observable("r") + .add_equation("r = a * x") + .add_equation("d_x = -r + coupling.x") + .build()) + m = Cls(g=0.0).configure(weights=np.zeros((2, 2))) + f = m.get_numba_dfun() + state = np.array([[0.5, -0.5]]) + ds, obs = f(state, np.zeros((1, 2))) + np.testing.assert_allclose(ds[0], -1.0 * state[0], atol=1e-12) + np.testing.assert_allclose(obs[0], state[0], atol=1e-12) + + # Sad path: observable declared but never assigned -> build fails. + bad = (ModelBuilder("ObsBad") + .add_state("x", initial=0.0) + .add_coupling("x") + .add_param("a", default=1.0) + .add_observable("r") + .add_equation("d_x = -a * x + coupling.x")) + with pytest.raises(ValueError, match="never assigned"): + bad.build() + + +def test_dependent_param_via_builder(): + """Dependent params (`formula=...`) are added the same way as independents.""" + Cls = (ModelBuilder("Dep") + .add_state("x") + .add_coupling("x") + .add_param("a", default=2.0) + .add_param("b", default=3.0) + .add_param("c", formula="a + b") + .add_equation("d_x = -x / c + coupling.x") + .build()) + m = Cls(g=0.0).configure(weights=np.zeros((2, 2))) + np.testing.assert_allclose(m.c, 5.0) + + +def test_override_initial_propagates(): + Cls = (ModelBuilder("Init") + .add_state("x", initial=0.0) + .add_coupling("x") + .add_param("a", default=1.0) + .add_equation("d_x = -a*x + coupling.x") + .override_initial("x", 0.7) + .build()) + m = Cls(g=0.0).configure(weights=np.zeros((2, 2))) + s0 = m.initial_state(n_rois=2) + np.testing.assert_allclose(s0, np.full((1, 2), 0.7)) + + +def test_duplicate_names_rejected(): + b = ModelBuilder("Dup").add_state("x") + with pytest.raises(ValueError, match="Duplicate state"): + b.add_state("x") + + b2 = ModelBuilder("Dup").add_param("a", default=1.0) + with pytest.raises(ValueError, match="Duplicate parameter"): + b2.add_param("a", default=2.0) + + b3 = ModelBuilder("Dup").add_coupling("x") + with pytest.raises(ValueError, match="Duplicate coupling"): + b3.add_coupling("x") + + +def test_empty_builder_rejected_at_spec(): + with pytest.raises(ValueError, match="no state variables"): + ModelBuilder("Empty").spec() + + with pytest.raises(ValueError, match="no equations"): + ModelBuilder("NoEq").add_state("x").spec() + + +def test_invalid_inputs_raise(): + with pytest.raises(ValueError): + ModelBuilder("") + b = ModelBuilder("X").add_state("x").add_coupling("x").add_param("a", default=1.0) + with pytest.raises(ValueError): + b.add_equation(" ") + with pytest.raises(TypeError): + b.add_helper(42) + with pytest.raises(TypeError): + b.on_configure("not-callable") + + +def test_builder_equivalent_to_handwritten_modelspec(): + """A builder-built class should be byte-for-byte interchangeable with one + constructed from a literal ModelSpec.""" + spec_literal = ModelSpec( + name="Equiv", + state_vars=[__import__("neuronumba.simulator.models.dsl", fromlist=["StateVar"]).StateVar("x", initial=0.0)], + coupling_vars=[__import__("neuronumba.simulator.models.dsl", fromlist=["CouplingVar"]).CouplingVar("x", kind="linear")], + observables=[], + parameters=[__import__("neuronumba.simulator.models.dsl", fromlist=["Parameter"]).Parameter("a", default=2.5)], + equations="d_x = -a * x + coupling.x", + ) + Cls_literal = build_model(spec_literal) + + Cls_builder = (ModelBuilder("Equiv") + .add_state("x", initial=0.0) + .add_coupling("x", kind="linear") + .add_param("a", default=2.5) + .add_equation("d_x = -a * x + coupling.x") + .build()) + + n = 4 + W = np.zeros((n, n)) + m_lit = Cls_literal(g=0.0).configure(weights=W) + m_bld = Cls_builder(g=0.0).configure(weights=W) + + state = np.array([[0.5, -0.1, 0.3, 0.2]]) + coup = np.zeros((1, n)) + ds_lit, _ = m_lit.get_numba_dfun()(state, coup) + ds_bld, _ = m_bld.get_numba_dfun()(state, coup) + np.testing.assert_allclose(ds_lit, ds_bld, atol=0.0) From 04f4c1cc1fc3927078e450f1ab32ab14ad0063aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:20 +0200 Subject: [PATCH 40/48] docs(dsl_models): update readme with zerlaut port and feature coverage --- examples/dsl_models/README.md | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/examples/dsl_models/README.md b/examples/dsl_models/README.md index 2cce097..daafeaa 100644 --- a/examples/dsl_models/README.md +++ b/examples/dsl_models/README.md @@ -10,6 +10,7 @@ hand-written counterpart in `src/neuronumba/simulator/models/`. | [deco2014_dsl.py](deco2014_dsl.py) | `deco2014.py` | 67 vs 416 | | [naskar2021_dsl.py](naskar2021_dsl.py) | `naskar2021.py` | 69 vs 104 | | [montbrio_dsl.py](montbrio_dsl.py) | `montbrio.py` | 88 vs 160 | +| [zerlaut_dsl.py](zerlaut_dsl.py) (1st + 2nd order) | `zerlaut.py` | ~320 vs 800 | There is also [compare_simulations.py](compare_simulations.py): a runnable script that simulates pairs of DSL and hand-written models on the same inputs @@ -35,16 +36,24 @@ If you change a DSL spec (e.g. for a new feature you want to add), the equivalence test still has to pass against the hand-written reference. If they genuinely diverge, that's a model change and belongs in a separate PR. -## Models not yet ported +## DSL feature coverage -- **Zerlaut** (`zerlaut.py`, ~800 LOC) — uses several `@njit` helper subroutines - (`get_fluct_regime_vars`, `threshold_func`, `erfc_approx`). The DSL now - supports user-supplied helpers via `helpers=[...]`, so a port is feasible; - it just hasn't been written yet. +All neuronumba models in `src/neuronumba/simulator/models/` are now expressed +in the DSL. The Zerlaut port additionally exercised: + +- Tuple-unpacking from a multi-return helper (`mu_V, sigma_V, T_V = + get_fluct_regime_vars(...)`). +- Long argument lists (TF takes 21 positional arguments — the DSL just + transcribes them; numba treats this exactly the same as the hand-written + code). +- Helpers that compose other helpers (`TF` calls `get_fluct_regime_vars` + and `threshold_func` and `erfc_approx`). +- 1D-array parameters as defaults (`P_e`, `P_i` are length-10 polynomial + coefficients). The DSL can also express models with non-scalar matrix parameters and conduction-delay coupling — see `tests/test_dsl_matrix_params.py` and -`tests/test_dsl_delayed.py` for the relevant patterns. +`tests/test_dsl_delayed.py`. ## Usage From 4688a8c5e5b4ef3fd68e2dea86f5bbda2a49f5db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:24 +0200 Subject: [PATCH 41/48] feat(dsl_models): update exports to include zerlaut dsl models --- examples/dsl_models/__init__.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/examples/dsl_models/__init__.py b/examples/dsl_models/__init__.py index abceb92..59ff9d7 100644 --- a/examples/dsl_models/__init__.py +++ b/examples/dsl_models/__init__.py @@ -4,7 +4,7 @@ purposes simultaneously: 1. **User-facing examples.** Read them to learn how to express a model in the - DSL. Each file is self-contained and small (~30-60 LOC vs 100-400 LOC for + DSL. Each file is self-contained and small (~30-300 LOC vs 100-800 LOC for the hand-written equivalents). 2. **Regression test references.** `tests/test_dsl_equivalence.py` imports @@ -12,19 +12,23 @@ hand-written ones bit-by-bit. If you change a hand-written model (e.g. a numerical bug fix), the corresponding spec here must follow, or the equivalence test will fail. - -Models out of scope for this set: - - Zerlaut: needs @njit helper subroutines (deferred to v0.2) - - OrnsteinUhlenbeck: needs non-scalar matrix parameters (deferred to v0.2) """ from .deco2014_dsl import deco_spec, Deco2014DSL from .hopf_dsl import hopf_spec, HopfDSL from .montbrio_dsl import montbrio_spec, MontbrioDSL from .naskar2021_dsl import naskar_spec, Naskar2021DSL +from .zerlaut_dsl import ( + zerlaut_first_order_spec, + zerlaut_second_order_spec, + ZerlautAdaptationFirstOrderDSL, + ZerlautAdaptationSecondOrderDSL, +) __all__ = [ "Deco2014DSL", "deco_spec", "HopfDSL", "hopf_spec", "MontbrioDSL", "montbrio_spec", "Naskar2021DSL", "naskar_spec", + "ZerlautAdaptationFirstOrderDSL", "zerlaut_first_order_spec", + "ZerlautAdaptationSecondOrderDSL", "zerlaut_second_order_spec", ] From fee5dfa960124c86f6b3ad254fc4af245cca2f39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:28 +0200 Subject: [PATCH 42/48] feat(dsl): update rewriter to support tuple-unpacking assignments --- .../simulator/models/dsl/rewriter.py | 47 ++++++++++++++----- 1 file changed, 35 insertions(+), 12 deletions(-) diff --git a/src/neuronumba/simulator/models/dsl/rewriter.py b/src/neuronumba/simulator/models/dsl/rewriter.py index 1b4ba96..5f841c6 100644 --- a/src/neuronumba/simulator/models/dsl/rewriter.py +++ b/src/neuronumba/simulator/models/dsl/rewriter.py @@ -48,22 +48,45 @@ def __init__( # --- assignment LHS handling ------------------------------------------ def visit_Assign(self, node: ast.Assign): - # Targets: only simple names allowed (e.g. `Ie = ...` or `dSe = ...`). - if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name): + # Single-target assigns: either `name = ...` (intermediate / derivative) + # or `a, b, c = helper(...)` (tuple-unpacking from a multi-return helper). + # Chained assignment (`a = b = ...`) and starred unpacking are rejected. + if len(node.targets) != 1: self.errors.append( - f"Line {node.lineno}: only single-name assignments are supported" + f"Line {node.lineno}: only single-target assignments are supported" + ) + return node + + tgt = node.targets[0] + if isinstance(tgt, ast.Name): + target = tgt.id + # If the target is `d_`, mark it as a derivative (no rewriting + # needed — we collect derivatives at emit time). Anything else is an + # intermediate (or an observable). + if not ( + target.startswith("d_") + and target[2:] in self.state_index + ): + self.intermediates.add(target) + elif isinstance(tgt, ast.Tuple): + # `mu_V, sigma_V, T_V = get_fluct_regime_vars(...)` — every name in + # the tuple becomes an intermediate. Nested tuples and starred + # targets are rejected for simplicity. + for elt in tgt.elts: + if not isinstance(elt, ast.Name): + self.errors.append( + f"Line {node.lineno}: tuple-unpacking targets must be " + f"plain names; got {ast.dump(elt)}" + ) + return node + self.intermediates.add(elt.id) + else: + self.errors.append( + f"Line {node.lineno}: assignment target must be a name or a " + f"tuple of names; got {type(tgt).__name__}" ) return node - target = node.targets[0].id - # If the target is `d_`, mark it as a derivative (no rewriting needed — - # we collect derivatives at emit time and assemble the return value ourselves). - # Anything else is an intermediate (or an observable). - if not ( - target.startswith("d_") - and target[2:] in self.state_index - ): - self.intermediates.add(target) # Recurse into the RHS only. node.value = self.visit(node.value) return node From 668667b3bd60567402ba948ec05832986c307f40 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:32 +0200 Subject: [PATCH 43/48] feat(models): update zerlaut dfun to hoist m_aux out of jit closure --- src/neuronumba/simulator/models/zerlaut.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/neuronumba/simulator/models/zerlaut.py b/src/neuronumba/simulator/models/zerlaut.py index c0e6477..cb3b081 100644 --- a/src/neuronumba/simulator/models/zerlaut.py +++ b/src/neuronumba/simulator/models/zerlaut.py @@ -346,13 +346,19 @@ def get_numba_dfun(self): m_aux = self.m_aux.copy() P = self.P P_aux = self.P_aux + # Hoist m_aux access into the factory: m_aux is a numba typed Dict and + # cannot be lowered as a freevar inside the JIT'd dfun. P_e/P_i are + # plain ndarrays and capture cleanly via closure (this matches what + # ZerlautAdaptationSecondOrder.get_numba_dfun does). + P_e = m_aux[np.intp(P_aux.P_e)] + P_i = m_aux[np.intp(P_aux.P_i)] @nb.njit(cache=NUMBA_CACHE, fastmath=NUMBA_FASTMATH, nogil=NUMBA_NOGIL) def ZerlautAdaptationFirstOrder_dfun(state: NDA_f8_2d, coupling: NDA_f8_2d): """ Numba-compiled function to compute the derivatives of the state variables and observables for the Zerlaut adaptation first order model. - + :param state: 2D array of state variables (shape: n_state_vars x n_nodes) :param coupling: 2D array of coupling values (shape: n_coupling_vars x n_nodes) :return: 2D array of derivatives (shape: n_state_vars x n_nodes) @@ -362,7 +368,7 @@ def ZerlautAdaptationFirstOrder_dfun(state: NDA_f8_2d, coupling: NDA_f8_2d): inv_T = 1.0 / T_val weight_noise_val = m[np.intp(P.weight_noise)] K_ext_e_val = m[np.intp(P.K_ext_e)] - + # Precompute adaptation parameters tau_w_e_val = m[np.intp(P.tau_w_e)] tau_w_i_val = m[np.intp(P.tau_w_i)] @@ -372,10 +378,8 @@ def ZerlautAdaptationFirstOrder_dfun(state: NDA_f8_2d, coupling: NDA_f8_2d): a_i_val = m[np.intp(P.a_i)] E_L_e_val = m[np.intp(P.E_L_e)] E_L_i_val = m[np.intp(P.E_L_i)] - - # Precompute TF parameters - P_e = m_aux[np.intp(P_aux.P_e)] - P_i = m_aux[np.intp(P_aux.P_i)] + + # Precompute TF parameters (P_e, P_i closed over from factory) Q_e_val = m[np.intp(P.Q_e)] tau_e_val = m[np.intp(P.tau_e)] E_e_val = m[np.intp(P.E_e)] @@ -453,6 +457,8 @@ def ZerlautAdaptationFirstOrder_dfun(state: NDA_f8_2d, coupling: NDA_f8_2d): return np.stack((dE, dI, dWe, dWi, dod)), np.empty((1,1)) + return ZerlautAdaptationFirstOrder_dfun + class ZerlautAdaptationSecondOrder(ZerlautAdaptationFirstOrder): From 55276b060c796c914b58ed69446bca66fe61f49e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:35 +0200 Subject: [PATCH 44/48] test(tests): update conftest to add zerlaut spec and class fixtures --- tests/conftest.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index 489dca1..305b195 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -24,10 +24,14 @@ deco_spec as _deco_spec, naskar_spec as _naskar_spec, montbrio_spec as _montbrio_spec, + zerlaut_first_order_spec as _zerlaut_1o_spec, + zerlaut_second_order_spec as _zerlaut_2o_spec, HopfDSL as _HopfDSL, Deco2014DSL as _Deco2014DSL, Naskar2021DSL as _Naskar2021DSL, MontbrioDSL as _MontbrioDSL, + ZerlautAdaptationFirstOrderDSL as _ZerlautAdaptationFirstOrderDSL, + ZerlautAdaptationSecondOrderDSL as _ZerlautAdaptationSecondOrderDSL, ) @@ -73,6 +77,16 @@ def montbrio_spec(): return _montbrio_spec +@pytest.fixture(scope="session") +def zerlaut_1o_spec(): + return _zerlaut_1o_spec + + +@pytest.fixture(scope="session") +def zerlaut_2o_spec(): + return _zerlaut_2o_spec + + # ----- materialized model classes -------------------------------------------- # Re-export the built classes from `examples/dsl_models/` directly. Tests # exercise the exact same class objects a user would import. @@ -95,3 +109,13 @@ def naskar_dsl_cls(): @pytest.fixture(scope="session") def montbrio_dsl_cls(): return _MontbrioDSL + + +@pytest.fixture(scope="session") +def zerlaut_1o_dsl_cls(): + return _ZerlautAdaptationFirstOrderDSL + + +@pytest.fixture(scope="session") +def zerlaut_2o_dsl_cls(): + return _ZerlautAdaptationSecondOrderDSL From 804111347ebd94d054470bf9dfcb59a8cfb8d48a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:39 +0200 Subject: [PATCH 45/48] test(tests): update equivalence tests to cover zerlaut models --- tests/test_dsl_equivalence.py | 27 ++++++++++++++++++++++++++- 1 file changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/test_dsl_equivalence.py b/tests/test_dsl_equivalence.py index d4671b1..3bc1bfc 100644 --- a/tests/test_dsl_equivalence.py +++ b/tests/test_dsl_equivalence.py @@ -6,7 +6,10 @@ """ import numpy as np -from neuronumba.simulator.models import Hopf, Deco2014, Naskar2021, Montbrio +from neuronumba.simulator.models import ( + Hopf, Deco2014, Naskar2021, Montbrio, + ZerlautAdaptationFirstOrder, ZerlautAdaptationSecondOrder, +) def _compare_dfun(m_ref, m_dsl, n_rois, rng, atol=1e-12): @@ -66,3 +69,25 @@ def test_equivalence_montbrio(weights, n_rois, montbrio_dsl_cls): m_dsl = montbrio_dsl_cls(g=0.5).configure(weights=weights) _compare_coupling(m_ref, m_dsl, n_rois, rng) _compare_dfun(m_ref, m_dsl, n_rois, rng, atol=1e-10) + + +def test_equivalence_zerlaut_first_order(weights, n_rois, zerlaut_1o_dsl_cls): + rng = np.random.default_rng(42) + m_ref = ZerlautAdaptationFirstOrder(g=0.5).configure(weights=weights) + m_dsl = zerlaut_1o_dsl_cls(g=0.5).configure(weights=weights) + _compare_coupling(m_ref, m_dsl, n_rois, rng) + _compare_dfun(m_ref, m_dsl, n_rois, rng) + + +def test_equivalence_zerlaut_second_order(weights, n_rois, zerlaut_2o_dsl_cls): + rng = np.random.default_rng(42) + m_ref = ZerlautAdaptationSecondOrder(g=0.5).configure(weights=weights) + m_dsl = zerlaut_2o_dsl_cls(g=0.5).configure(weights=weights) + _compare_coupling(m_ref, m_dsl, n_rois, rng) + # 2nd-order's 12 chained TF calls + finite-diff derivatives sit at the FP + # noise floor (~1e-12 absolute on dE/dI). The hand-written code keeps + # inv_T as a 1D array (slice of `self.m`), while the DSL closes over it + # as a scalar; numba's fastmath emits slightly different FMA sequences + # for `arr * scalar` vs `arr * arr-of-constants`, so the lowest ~1 ulp + # diverges per FLOP. 1e-11 covers that without hiding real bugs. + _compare_dfun(m_ref, m_dsl, n_rois, rng, atol=1e-11) From 54bab4b418ecca3e70b203db0e5b1fca8508b59d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:43 +0200 Subject: [PATCH 46/48] test(tests): update helpers tests to cover tuple-unpacking cases --- tests/test_dsl_helpers.py | 94 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) diff --git a/tests/test_dsl_helpers.py b/tests/test_dsl_helpers.py index 040ba15..2e7ddd9 100644 --- a/tests/test_dsl_helpers.py +++ b/tests/test_dsl_helpers.py @@ -140,3 +140,97 @@ def test_helpers_not_listed_for_specs_without_them(): f = m.get_numba_dfun() ds, _ = f(np.array([[0.5, 1.0, -0.5, 2.0]]), np.zeros((1, 4))) np.testing.assert_allclose(ds[0], -np.array([0.5, 1.0, -0.5, 2.0]), atol=1e-12) + + +# --- tuple-unpacking from helpers (Zerlaut-class enabler) ------------------ + +@nb.njit(nb.types.UniTuple(nb.f8[:], 3)(nb.f8[:]), cache=False) +def split3(x): + """Returns three arrays derived from x — exercises tuple unpacking.""" + return x, 2.0 * x, x * x + + +def test_helper_tuple_unpacking(): + """A helper returning a tuple should unpack into multiple intermediates.""" + spec = ModelSpec( + name="TupleUnpack", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[Parameter("a", default=1.0)], + equations=""" + u, v, w = split3(x) + d_x = -a * (u + v + w) + coupling.x + """, + helpers=[split3], + ) + Cls = build_model(spec) + n = 4 + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + f = m.get_numba_dfun() + state = np.array([[0.5, 1.0, -0.5, 2.0]]) + ds, _ = f(state, np.zeros((1, n))) + # u + v + w = x + 2x + x^2 = 3x + x^2; d_x = -1 * (3x + x^2) + expected = -1.0 * (3.0 * state[0] + state[0] ** 2) + np.testing.assert_allclose(ds[0], expected, atol=1e-12) + + +def test_helper_tuple_unpacking_targets_become_intermediates(): + """Names introduced via tuple-unpacking should be usable in subsequent equations.""" + spec = ModelSpec( + name="TupleUnpackChain", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[Parameter("a", default=2.0)], + equations=""" + u, v, w = split3(x) + combined = u * v + w + d_x = -a * combined + coupling.x + """, + helpers=[split3], + ) + Cls = build_model(spec) + n = 3 + m = Cls(g=0.0).configure(weights=np.zeros((n, n))) + f = m.get_numba_dfun() + state = np.array([[0.1, 0.5, -0.3]]) + ds, _ = f(state, np.zeros((1, n))) + # combined = x * 2x + x^2 = 3 x^2; d_x = -2 * 3 x^2 = -6 x^2 + expected = -6.0 * state[0] ** 2 + np.testing.assert_allclose(ds[0], expected, atol=1e-12) + + +def test_helper_tuple_unpacking_chained_assignment_rejected(): + """`a = b = expr` is not supported (multiple targets).""" + spec = ModelSpec( + name="ChainBad", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[Parameter("a", default=1.0)], + equations=""" + u = v = a * x + d_x = -u + coupling.x + """, + ) + with pytest.raises(ValueError, match="single-target"): + build_model(spec) + + +def test_helper_tuple_unpacking_starred_target_rejected(): + """Starred targets (`a, *b = ...`) are not supported.""" + spec = ModelSpec( + name="StarBad", + state_vars=[StateVar("x", initial=0.0)], + coupling_vars=[CouplingVar("x", kind="linear")], + observables=[], + parameters=[Parameter("a", default=1.0)], + equations=""" + u, *rest = split3(x) + d_x = -a * u + coupling.x + """, + helpers=[split3], + ) + with pytest.raises(ValueError, match="tuple-unpacking targets must be"): + build_model(spec) From 4fba90886f803d8dfa81edfabdb96ced3f98ca8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:19:47 +0200 Subject: [PATCH 47/48] feat(dsl_models): add zerlaut dsl spec with first and second order --- examples/dsl_models/zerlaut_dsl.py | 318 +++++++++++++++++++++++++++++ 1 file changed, 318 insertions(+) create mode 100644 examples/dsl_models/zerlaut_dsl.py diff --git a/examples/dsl_models/zerlaut_dsl.py b/examples/dsl_models/zerlaut_dsl.py new file mode 100644 index 0000000..71cc722 --- /dev/null +++ b/examples/dsl_models/zerlaut_dsl.py @@ -0,0 +1,318 @@ +"""Zerlaut adaptation models, expressed in the neuronumba DSL. + +DSL equivalents of: + - `neuronumba.simulator.models.zerlaut.ZerlautAdaptationFirstOrder` + - `neuronumba.simulator.models.zerlaut.ZerlautAdaptationSecondOrder` + +The math is identical. The four njit subroutines used inside the dfun +(`get_fluct_regime_vars`, `threshold_func`, `TF`, `erfc_approx`) are imported +verbatim from the hand-written module so both versions share the exact same +helper bytecode — only the dfun itself is regenerated from the spec. + +Two patterns the DSL needs (and now provides): + - Tuple unpacking from a multi-return helper: + ``mu_V, sigma_V, T_V = get_fluct_regime_vars(...)`` + - Long argument lists (TF takes 21 positional args). The DSL just transcribes + them; numba treats this exactly the same as the hand-written code. + +The hand-written 1st- and 2nd-order dfuns clamp inputs via fancy-index +assignment after `np.where(... < 0)`. The DSL spec uses ``np.maximum(x, 0.0)`` +which is bit-equivalent for the values these tensors actually take (K_ext_e, +K_ext_i are non-negative by construction). + +`tests/test_dsl_equivalence.py::test_equivalence_zerlaut_*` compares the +DSL-built dfuns to the hand-written ones on random inputs at 1e-12. +""" +import numpy as np + +from neuronumba.simulator.models.dsl import ( + ModelSpec, StateVar, CouplingVar, Parameter, build_model, +) +# Helpers come straight from the hand-written file: bit-identical bytecode. +from neuronumba.simulator.models.zerlaut import ( + get_fluct_regime_vars, threshold_func, TF, +) + + +# ---- shared parameter list -------------------------------------------------- + +_P_e_DEFAULT = np.array([ + -0.04983106, 0.005063550882777035, -0.023470121807314552, + 0.0022951513725067503, -0.0004105302652029825, 0.010547051343547399, + -0.03659252821136933, 0.007437487505797858, 0.001265064721846073, + -0.04072161294490446, +]) +_P_i_DEFAULT = np.array([ + -0.05149122024209484, 0.004003689190271077, -0.008352013668528155, + 0.0002414237992765705, -0.0005070645080016026, 0.0014345394104282397, + -0.014686689498949967, 0.004502706285435741, 0.0028472190352532454, + -0.015357804594594548, +]) + + +def _common_parameters(): + """The 32 parameters shared by both 1st- and 2nd-order Zerlaut models.""" + return [ + # Cellular properties + Parameter("g_L", default=10.0), + Parameter("E_L_e", default=-65.0), + Parameter("E_L_i", default=-65.0), + Parameter("C_m", default=200.0), + Parameter("b_e", default=60.0), + Parameter("a_e", default=4.0), + Parameter("b_i", default=0.0), + Parameter("a_i", default=0.0), + Parameter("tau_w_e", default=500.0), + Parameter("tau_w_i", default=1.0), + # Synaptic properties + Parameter("E_e", default=0.0), + Parameter("E_i", default=-80.0), + Parameter("Q_e", default=1.5), + Parameter("Q_i", default=5.0), + Parameter("tau_e", default=5.0), + Parameter("tau_i", default=5.0), + # Network composition + Parameter("N_tot", default=10000.0), + Parameter("p_connect_e", default=0.05), + Parameter("p_connect_i", default=0.05), + Parameter("gi", default=0.2), + Parameter("K_ext_e", default=400.0), + Parameter("K_ext_i", default=0.0), + # Time scale + Parameter("T", default=20.0), + # Polynomial coefficients (length-10 1D arrays) + Parameter("P_e", default=_P_e_DEFAULT), + Parameter("P_i", default=_P_i_DEFAULT), + # External drives & noise + Parameter("external_input_ex_ex", default=0.0), + Parameter("external_input_ex_in", default=0.0), + Parameter("external_input_in_ex", default=0.0), + Parameter("external_input_in_in", default=0.0), + Parameter("tau_OU", default=5.0), + Parameter("weight_noise", default=10.5), + Parameter("S_i", default=1.0), + # Factory-level constant: 1/T precomputed at configure() time so the + # dfun multiplies by inv_T (matches the hand-written code's hoisting + # exactly and is bit-equivalent to it). + Parameter("inv_T", formula="1.0 / T"), + ] + + +# ---- First-order model ------------------------------------------------------ + +zerlaut_first_order_spec = ModelSpec( + name="ZerlautAdaptationFirstOrderDSL", + state_vars=[ + StateVar("E", initial=0.001), + StateVar("I", initial=0.001), + StateVar("W_e", initial=0.001), + StateVar("W_i", initial=0.001), + StateVar("ou_drift", initial=0.001), + ], + coupling_vars=[CouplingVar("E", kind="linear")], + observables=[], + parameters=_common_parameters(), + helpers=[get_fluct_regime_vars, threshold_func, TF], + equations=""" + # External firing rate (clamped to non-negative; bit-equivalent to the + # hand-written `np.where(Fe_ext * K_ext_e < 0); Fe_ext[idx] = 0` since + # K_ext_e is non-negative). + Fe_ext_raw = coupling.E + weight_noise * ou_drift + Fe_ext = np.maximum(Fe_ext_raw, 0.0) + + # Population-specific external inputs. Fi_ext = 0 in the hand-written + # code, so Fi_ext_ex / Fi_ext_in collapse to the corresponding ext_*_in + # parameters. + Fe_ext_ex = Fe_ext + external_input_ex_ex + Fi_ext_ex = external_input_ex_in + Fe_ext_in = Fe_ext + external_input_in_ex + Fi_ext_in = external_input_in_in + + # Excitatory firing rate: dE = (TF_e - E) / T + TF_e = TF(E, I, Fe_ext_ex, Fi_ext_ex, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + d_E = (TF_e - E) * inv_T + + # Inhibitory firing rate: dI = (TF_i - I) / T + TF_i = TF(E, I, Fe_ext_in, Fi_ext_in, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + d_I = (TF_i - I) * inv_T + + # Excitatory adaptation + mu_V_e, sigma_V_e, T_V_e = get_fluct_regime_vars( + E, I, Fe_ext_ex, Fi_ext_ex, W_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, E_L_e, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + d_W_e = -W_e / tau_w_e + b_e * E + a_e * (mu_V_e - E_L_e) / tau_w_e + + # Inhibitory adaptation + mu_V_i, sigma_V_i, T_V_i = get_fluct_regime_vars( + E, I, Fe_ext_in, Fi_ext_in, W_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, E_L_i, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + d_W_i = -W_i / tau_w_i + b_i * I + a_i * (mu_V_i - E_L_i) / tau_w_i + + # OU drift relaxation (noise injection happens in the integrator) + d_ou_drift = -ou_drift / tau_OU + """, +) + +ZerlautAdaptationFirstOrderDSL = build_model(zerlaut_first_order_spec) + + +# ---- Second-order model ----------------------------------------------------- + +zerlaut_second_order_spec = ModelSpec( + name="ZerlautAdaptationSecondOrderDSL", + state_vars=[ + StateVar("E", initial=0.0), + StateVar("I", initial=0.0), + StateVar("C_ee", initial=0.0), + StateVar("C_ei", initial=0.0), + StateVar("C_ii", initial=0.0), + StateVar("W_e", initial=0.0), + StateVar("W_i", initial=0.0), + StateVar("ou_drift", initial=0.0), + ], + coupling_vars=[CouplingVar("E", kind="linear")], + observables=[], + parameters=_common_parameters() + [ + # 2nd-order-specific factory-level precomputes (mirror the hand-written + # constants computed at the top of get_numba_dfun). + Parameter("N_e", formula="N_tot * (1.0 - gi)"), + Parameter("N_i", formula="N_tot * gi"), + ], + helpers=[get_fluct_regime_vars, threshold_func, TF], + equations=""" + # Finite-difference step constants (factory-equivalent literals; numba + # constant-folds these). + df = 1e-7 + df_1e3 = df * 1e3 + df_2_1e3 = 2 * df * 1e3 + df_1e3_sq = df_1e3 * df_1e3 + inv_df_2_1e3 = 1.0 / df_2_1e3 + inv_df_1e3_sq = 1.0 / df_1e3_sq + + # External firing-rate inputs (clamped to non-negative). + noise_input = weight_noise * ou_drift + E_input_excitatory_raw = coupling.E + external_input_ex_ex + noise_input + E_input_excitatory = np.maximum(E_input_excitatory_raw, 0.0) + E_input_inhibitory_raw = S_i * coupling.E + external_input_in_ex + noise_input + E_input_inhibitory = np.maximum(E_input_inhibitory_raw, 0.0) + I_input_excitatory = external_input_ex_in + I_input_inhibitory = external_input_in_in + + # Base TF values (one for each population). + _TF_e = TF(E, I, E_input_excitatory, I_input_excitatory, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + _TF_i = TF(E, I, E_input_inhibitory, I_input_inhibitory, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + + # Perturbed TF values for centered-FD derivatives. + TF_e_plus_df = TF(E + df, I, E_input_excitatory, I_input_excitatory, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_e_minus_df = TF(E - df, I, E_input_excitatory, I_input_excitatory, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_e_I_plus_df = TF(E, I + df, E_input_excitatory, I_input_excitatory, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_e_I_minus_df = TF(E, I - df, E_input_excitatory, I_input_excitatory, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + + TF_i_E_plus_df = TF(E + df, I, E_input_inhibitory, I_input_inhibitory, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_i_E_minus_df = TF(E - df, I, E_input_inhibitory, I_input_inhibitory, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_i_plus_df = TF(E, I + df, E_input_inhibitory, I_input_inhibitory, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_i_minus_df = TF(E, I - df, E_input_inhibitory, I_input_inhibitory, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + + # First derivatives. + diff_fe_TF_e = (TF_e_plus_df - TF_e_minus_df) * inv_df_2_1e3 + diff_fe_TF_i = (TF_i_E_plus_df - TF_i_E_minus_df) * inv_df_2_1e3 + diff_fi_TF_e = (TF_e_I_plus_df - TF_e_I_minus_df) * inv_df_2_1e3 + diff_fi_TF_i = (TF_i_plus_df - TF_i_minus_df) * inv_df_2_1e3 + + # Second derivatives (single-variable). + diff2_fe_fe_e = (TF_e_plus_df - 2.0 * _TF_e + TF_e_minus_df) * inv_df_1e3_sq + diff2_fi_fi_e = (TF_e_I_plus_df - 2.0 * _TF_e + TF_e_I_minus_df) * inv_df_1e3_sq + diff2_fe_fe_i = (TF_i_E_plus_df - 2.0 * _TF_i + TF_i_E_minus_df) * inv_df_1e3_sq + diff2_fi_fi_i = (TF_i_plus_df - 2.0 * _TF_i + TF_i_minus_df) * inv_df_1e3_sq + + # Mixed derivatives (need both-corner perturbations). + TF_e_both_plus = TF(E + df, I + df, E_input_excitatory, I_input_excitatory, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_e_both_minus = TF(E - df, I - df, E_input_excitatory, I_input_excitatory, W_e, P_e, E_L_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_i_both_plus = TF(E + df, I + df, E_input_inhibitory, I_input_inhibitory, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + TF_i_both_minus = TF(E - df, I - df, E_input_inhibitory, I_input_inhibitory, W_i, P_i, E_L_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, N_tot, p_connect_e, p_connect_i, gi, K_ext_e, K_ext_i) + + diff2_fe_fi_e = (TF_e_both_plus - TF_e_plus_df - TF_e_I_plus_df + 2.0*_TF_e - TF_e_minus_df - TF_e_I_minus_df + TF_e_both_minus) * (0.25 * inv_df_1e3_sq) + diff2_fe_fi_i = (TF_i_both_plus - TF_i_E_plus_df - TF_i_plus_df + 2.0*_TF_i - TF_i_E_minus_df - TF_i_minus_df + TF_i_both_minus) * (0.25 * inv_df_1e3_sq) + + # Firing-rate equations. + half_inv_T = 0.5 * inv_T + d_E = (_TF_e - E + half_inv_T * ( + C_ee * diff2_fe_fe_e + C_ei * diff2_fe_fi_e + C_ei * diff2_fe_fi_e + C_ii * diff2_fi_fi_e + )) * inv_T + d_I = (_TF_i - I + half_inv_T * ( + C_ee * diff2_fe_fe_i + C_ei * diff2_fe_fi_i + C_ei * diff2_fe_fi_i + C_ii * diff2_fi_fi_i + )) * inv_T + + # Covariance equations. + TF_e_diff = _TF_e - E + TF_i_diff = _TF_i - I + d_C_ee = (_TF_e * (inv_T - _TF_e) / N_e + TF_e_diff * TF_e_diff + + 2.0 * (C_ee * diff_fe_TF_e + C_ei * diff_fi_TF_e) - 2.0 * C_ee) * inv_T + d_C_ei = (TF_e_diff * TF_i_diff + + C_ee * diff_fe_TF_e + C_ei * diff_fe_TF_i + C_ei * diff_fi_TF_e + C_ii * diff_fi_TF_i - + 2.0 * C_ei) * inv_T + d_C_ii = (_TF_i * (inv_T - _TF_i) / N_i + TF_i_diff * TF_i_diff + + 2.0 * (C_ii * diff_fi_TF_i + C_ei * diff_fe_TF_i) - 2.0 * C_ii) * inv_T + + # Adaptation (excitatory). + mu_V_e, sigma_V_e, T_V_e = get_fluct_regime_vars( + E, I, E_input_excitatory, I_input_excitatory, W_e, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, E_L_e, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + d_W_e = -W_e / tau_w_e + b_e * E + a_e * (mu_V_e - E_L_e) / tau_w_e + + # Adaptation (inhibitory). + mu_V_i, sigma_V_i, T_V_i = get_fluct_regime_vars( + E, I, E_input_inhibitory, I_input_inhibitory, W_i, + Q_e, tau_e, E_e, Q_i, tau_i, E_i, + g_L, C_m, E_L_i, N_tot, p_connect_e, p_connect_i, gi, + K_ext_e, K_ext_i) + d_W_i = -W_i / tau_w_i + b_i * I + a_i * (mu_V_i - E_L_i) / tau_w_i + + # OU drift relaxation. + d_ou_drift = -ou_drift / tau_OU + """, +) + +ZerlautAdaptationSecondOrderDSL = build_model(zerlaut_second_order_spec) From 9da14e3f32fc63336b87f4b5f9d558cf4f6dfb0f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ignacio=20Mart=C3=ADn?= Date: Thu, 7 May 2026 22:20:59 +0200 Subject: [PATCH 48/48] docs(dsl_models): update readme with per-spec usage guide --- examples/dsl_models/README.md | 40 +++++++++++++++++++---------------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/examples/dsl_models/README.md b/examples/dsl_models/README.md index daafeaa..f0a9870 100644 --- a/examples/dsl_models/README.md +++ b/examples/dsl_models/README.md @@ -36,24 +36,28 @@ If you change a DSL spec (e.g. for a new feature you want to add), the equivalence test still has to pass against the hand-written reference. If they genuinely diverge, that's a model change and belongs in a separate PR. -## DSL feature coverage - -All neuronumba models in `src/neuronumba/simulator/models/` are now expressed -in the DSL. The Zerlaut port additionally exercised: - -- Tuple-unpacking from a multi-return helper (`mu_V, sigma_V, T_V = - get_fluct_regime_vars(...)`). -- Long argument lists (TF takes 21 positional arguments — the DSL just - transcribes them; numba treats this exactly the same as the hand-written - code). -- Helpers that compose other helpers (`TF` calls `get_fluct_regime_vars` - and `threshold_func` and `erfc_approx`). -- 1D-array parameters as defaults (`P_e`, `P_i` are length-10 polynomial - coefficients). - -The DSL can also express models with non-scalar matrix parameters and -conduction-delay coupling — see `tests/test_dsl_matrix_params.py` and -`tests/test_dsl_delayed.py`. +## What each spec illustrates + +Pick the spec that's closest to what you want to build: + +- **`hopf_dsl.py`** — minimal 2-variable oscillator. Shows the basic shape + of `state_vars` / `coupling_vars` / `parameters` / `equations` and the + `kind="diffusive"` coupling kernel. +- **`deco2014_dsl.py`** — single-coupling-var rate model with declared + observables (`Ie`, `re`) and an EPS-fix using `np.where` in the equations. +- **`naskar2021_dsl.py`** — extends Deco-style dynamics with a third state + variable that itself follows a slow plasticity rule. +- **`montbrio_dsl.py`** — 6-variable Montbrio model with cross-coupling and + recurrent dependent parameters (`Parameter("J_N_ee", formula="...")`). +- **`zerlaut_dsl.py`** — uses user-supplied `@nb.njit` helpers + (`get_fluct_regime_vars`, `threshold_func`, `TF`, `erfc_approx`) imported + from `neuronumba.simulator.models.zerlaut`. Demonstrates tuple-unpacking + from a multi-return helper, helper composition (`TF` calls the others + internally), and 1D-array parameters as defaults (length-10 polynomial + coefficient vectors). + +For matrix-valued parameters and conduction-delay coupling, see +`tests/test_dsl_matrix_params.py` and `tests/test_dsl_delayed.py`. ## Usage