feat[event]: Add MultiPopulationRecurrentLIF for per-population time constants

[dreliq9]

Summary

RecurrentLIF and the EventProp variants currently share a single LIFParameters across all populations in a connected recurrent graph. Heterogeneous τ across populations — fast PV inhibition coupled to slower pyramidal cells, multi-cell-type biology-faithful circuits, trainable per-population time constants for surrogate-gradient SNN training — is broadly useful but currently inexpressible without bypassing the event-based engine.

This PR adds MultiPopulationRecurrentLIF (one LIFParameters per population) and a backing lif_exponential_flow_vec helper. Submitted as draft for early API feedback; the EventProp adjoint variant follows once this surface is approved.

Motivation

Heterogeneous τ across populations underpins several research lines that map cleanly to jaxsnn / BSS-2:

• E/I balance and PING/ING gamma networks — fast PV interneurons (τ ≈ 3–8 ms) coupled to slower pyramidal cells (τ ≈ 15–25 ms). The asymmetry is the mechanism (van Vreeswijk & Sompolinsky 1996; Whittington & Traub 2003).
• Trainable per-population τ improves SNN learning — Perez-Nieves et al. 2021 (Nature Communications, "Neural heterogeneity promotes robust learning") show heterogeneous τ as a performance feature on standard benchmarks; Yin et al. 2021 and Habashy et al. 2024 extend the line. A natural fit for jaxsnn's EventProp pipeline once the API supports it.
• Multi-cell-type biology-faithful circuits — cerebellum (granule 5 ms, pontine 15 ms, Purkinje 20 ms, DCN 30 ms, IO 50 ms), thalamocortical loops (TC 10 ms vs TRN 5 ms), hippocampus, basal ganglia.

The gap is jaxsnn-specific (verified empirically):

• RecurrentLIF over five populations with five distinct τ fails inside lif_exponential_flow (flow.py:25) with ValueError: All input arrays must have the same shape. The kernel is closed over scalar params.tau_mem and vmap'd with in_axes=(0, None) (params broadcast, not indexed); ttfs_solver is partial'd with scalar tau_mem at trace time.
• The same workload runs unmodified in Brian2, PyNN with NEST/Brian2 backend, and py-spinnaker2's Brian2 backend (each handles per-neuron tau as a declared parameter, not a closed-over scalar). The gap is specific to monolithic-kernel event-based libraries.

What this PR adds

lif_exponential_flow_vec (src/pyjaxsnn/jaxsnn/event/flow.py)

def lif_exponential_flow_vec(params: LIFParameters):
    """Vec'd dynamics flow for per-neuron tau_mem and tau_syn."""

Builds a per-neuron [N, 2, 2] kernel via jax.vmap and applies expm per neuron. Drop-in for callers maintaining the per-neuron-batched LIFState shape; no outer vmap required. Existing scalar-params lif_exponential_flow is unchanged.

MultiPopulationRecurrentLIF (src/pyjaxsnn/jaxsnn/event/modules/leaky_integrate_and_fire.py)

init_fn, apply_fn = MultiPopulationRecurrentLIF(
    layers=[64, 16],
    n_spikes=128,
    t_max=80e-3,
    params_per_population=[
        LIFParameters(tau_mem=5e-3,  tau_syn=5e-3,  v_th=1.0),  # fast pop
        LIFParameters(tau_mem=20e-3, tau_syn=20e-3, v_th=1.0),  # slow pop
    ],
    mean=[1.0, 1.0],
    std=[0.0, 0.0],
)

Internally concatenates per-population params into per-neuron arrays of shape (sum(layers),), threads them through lif_exponential_flow_vec (dynamics) and a per-neuron-vmapped ttfs_solver (spike timing). transition_with_recurrence is reused via partial; broadcasting handles the v_reset field correctly.

Validation rules (all raise ValueError at construction)

1. layers non-empty
2. each n in layers is a positive integer (isinstance(n, (int, np.integer)))
3. len(params_per_population) == len(layers)
4. each pop's tau_mem > 0 and tau_syn > 0
5. each pop's tau_mem / tau_syn ratio is ≈ 1.0 or ≈ 2.0 within 1e-6 rel-tol — analytical TTFS solver constraint inherited from event/root/ttfs.py
6. each pop's v_th == 1.0 within 1e-6 rel-tol — step.py:85 hard-codes evolved_neuron_state.V >= 1. and the TTFS solver's v_th cannot override it; off-1.0 silently desynchronises the two layers

(Rule 6 may also apply to the existing LIF / RecurrentLIF factories — happy to surface as a separate cleanup PR if useful.)

Tests

tests/sw/event/test_multi_population_recurrent_lif.py — 6 tests, all passing locally:

1. Construction with per-population params (happy path)

2. Length-mismatch validation

3. v_th-not-1.0 validation

4. Non-integer-layer validation

5. Heterogeneous-dynamics observable (load-bearing). Two populations (τ = 5 ms fast + τ = 20 ms slow) compared against RecurrentLIF baselines:

   • HETERO: pop A first-spike 3.35 ms, pop B 6.24 ms
   • all-FAST baseline: pop A 3.35 ms, pop B 4.07 ms
   • all-SLOW baseline: pop A 8.66 ms, pop B 11.55 ms

   Pop B in HETERO sits strictly between all-fast (4.07 ms) and all-slow (11.55 ms) — slow τ integrating an early kick from fast pop A. The test asserts both inequalities; catches a class of regression where per-population τ silently flattens.

6. Backward-compat smoke for unchanged RecurrentLIF

Backward compatibility

Zero regression. The 12 existing tests in test_lif, test_flow, test_lif_vs_recurrenteventproplif, test_lif_vs_eventproplif all pass unchanged. No existing factory (LIF, RecurrentLIF, EventPropLIF, RecurrentEventPropLIF, HardwareLIF, HardwareRecurrentLIF) is touched.

Scope and follow-ups

This PR (v1):

• Per-population homogeneity (one LIFParameters per population, not per-neuron).
• τ-ratio constrained to 1 or 2 per population — inherited from the analytical TTFS solver's domain.

Natural follow-ups:

• MultiPopulationRecurrentEventPropLIF — the EventProp adjoint variant. Mechanically parallel to the existing RecurrentEventPropLIF pattern; threads adjoint_lif_exponential_flow through the same per-population concatenation. The EventProp adjoint is linear in 1/tau_mem, so per-population rescales cleanly under heterogeneous τ. Happy to draft this in a follow-up PR once the v1 surface is approved.
• Per-neuron heterogeneity within a population — arbitrary per-neuron τ, not just per-population. Requires switching the solver from analytical (ttfs_solver) to Newton — newton_solver already exists in event/root/newton.py but isn't currently wired into the factories. The combination unlocks the trainable-per-neuron-τ research line cited above.
• Cerebellum-class workloads (τ_syn = 1 ms, τ_mem ∈ {5, 15, 20, 30, 50}) — ratios outside {1, 2}, so the Newton substitution above is the gating change. v1 already enables many simpler scenarios: fast-PV + slow-PC E/I networks, balanced attractors, two-time-scale reservoirs.

I'd appreciate a sanity check on the v2 direction (Newton solver substitution + the EventProp adjoint variant) before starting on it.

Incidental finding

While bootstrapping a dev environment I hit ImportError: No module named 'nir.data_ir' from src/pyjaxsnn/jaxsnn/event/from_nir_data.py against nir 1.0.7 from PyPI. Looks like upstream HEAD references an unreleased nir API. Worked around locally; happy to file as a separate issue if not already known internally.

---

Happy to iterate on naming, API surface, or v2 design.