Imported from chaobrain/BrainX-skill (brainx-display-cases/creative-experiment-verification/06-seizure-recruitment/run1/.agents/skills/brainpy-state/SKILL.md). Install upstream with npx skills add chaobrain/BrainX-skill --skill brainpy-state. Copyright stays with the author.

BrainPy-State

Purpose and boundary

Use this skill for the native brainpy.state modeling path: compose point neurons, synapses, communication operators, synaptic outputs, projections, inputs, plasticity, and readouts into simulations or trainable spiking neural networks.

Route NEST-compatible work to references/nest-compatible/nest-workflow.md, and online-learning execution to BrainTrace.

Open references/brainPy(legacy)/brainpy legacy workflow.md only when the workspace already contains legacy brainpy code, contains no BrainUnit quantities, unit-aware APIs, and when the user explicitly asks to preserve old BrainPy. Use native BrainPy-State for new project is HIGHLY RECOMMENDED; use BrainPy-State with BrainCell for all cellular modeling.

Treat that workflow as the only first-level legacy reference. Do not open or route directly to any other Markdown under references/brainPy(legacy)/; let the workflow select those second-level references.

Underlying principle of BrainPy-State

BrainPy-State composes unit-aware, stateful neuron and synapse Dynamics into projections that separate comm, syn, out, and post; BrainState transforms run the same Module graph for simulation or surrogate-gradient training. Align synaptic State to a neuron dimension instead of each connection, reduce that state from per-synapse to per-neuron, which makes simulation memory-efficient.

API structure

1. Choose and run neuron dynamics

A neuron call advances one dt, mutates its registered dynamical State, and exposes the current spike through get_spike(); initialize before every independent rollout, set V_initializer whenever startup voltage affects the result, and lower the complete time loop through brainstate.transform.

import brainpy
import brainstate
import braintools
import brainunit as u

with brainstate.environ.context(dt=0.1 * u.ms):
    neuron = brainpy.state.LIFRef(
        1,
        R=1.0 * u.ohm,
        tau=20.0 * u.ms,
        V_rest=-60.0 * u.mV,
        V_th=-50.0 * u.mV,
        V_reset=-60.0 * u.mV,
        tau_ref=5.0 * u.ms,
        V_initializer=braintools.init.Constant(-60.0 * u.mV),
    )
    brainstate.nn.init_all_states(neuron)
    times = u.math.arange(
        0.0 * u.ms,
        200.0 * u.ms,
        brainstate.environ.get_dt(),
    )

    def step(t):
        with brainstate.environ.context(t=t):
            neuron(25.0 * u.mA)
            return neuron.V.value, neuron.get_spike()

    voltages, spikes = brainstate.transform.for_loop(step, times)

assert voltages.shape == (2000, 1)
assert spikes.shape == (2000, 1)

Open references/braintools/input-current.md when a direct current protocol needs timed sections, pulses, waveforms, stochastic processes, or composition. Generate the complete time-major protocol once under the rollout dt and pass it as an input to for_loop; do not rebuild named sections or pulses from time predicates inside the model step. Open references/array-creation.md when constructing specialized unit-aware current ranges, grids, filled arrays, template-shaped arrays, or matrix patterns. Open references/component-selection.md when choosing among documented neuron, input, synapse, output, plasticity, or readout variants, for the decision boundary and complete category list. Open references/braintools/brainstate-control-flow-patterns.md when the rollout needs explicit carry, branching, or checkpointed control flow.

2. Compose synapses and projections

A projection deposits current into post before the postsynaptic neuron integrates the step, while its alignment determines which neuron dimension owns synaptic State and which kinetics remain exact; when recurrent excitation, mutual inhibition, or spike-driven connectivity is a point-neuron mechanism, represent it with explicit projections instead of a hand-written population-rate feedback State, and route aggregate rate coupling to BrainMass.

import brainpy
import brainstate
import braintools
import brainunit as u


class TwoPop(brainstate.nn.Module):
    def __init__(self, n_pre=20, n_post=10):
        super().__init__()
        self.pre = brainpy.state.LIFRef(
            n_pre,
            tau=20.0 * u.ms,
            tau_ref=5.0 * u.ms,
            V_rest=-60.0 * u.mV,
            V_th=-50.0 * u.mV,
            V_reset=-60.0 * u.mV,
            V_initializer=braintools.init.Constant(-60.0 * u.mV),
        )
        self.post = brainpy.state.LIFRef(
            n_post,
            tau=20.0 * u.ms,
            tau_ref=5.0 * u.ms,
            V_rest=-60.0 * u.mV,
            V_th=-50.0 * u.mV,
            V_reset=-60.0 * u.mV,
            V_initializer=braintools.init.Constant(-60.0 * u.mV),
        )
        self.proj = brainpy.state.AlignPostProj(
            comm=brainstate.nn.EventFixedProb(
                n_pre,
                n_post,
                conn_num=0.5,
                conn_weight=0.5 * u.mS,
            ),
            syn=brainpy.state.Expon.desc(n_post, tau=5.0 * u.ms),
            out=brainpy.state.COBA.desc(E=0.0 * u.mV),
            post=self.post,
        )

    def update(self, t, drive):
        with brainstate.environ.context(t=t):
            pre_spikes = self.pre.get_spike() != 0.0
            self.proj(pre_spikes)
            self.pre(drive)
            self.post(0.0 * u.mA)
            return self.post.V.value, self.proj.syn.g.value


with brainstate.environ.context(dt=0.1 * u.ms):
    net = TwoPop()
    brainstate.nn.init_all_states(net)
    times = u.math.arange(0.0 * u.ms, 200.0 * u.ms, brainstate.environ.get_dt())
    post_voltage, conductance = brainstate.transform.for_loop(
        lambda t: net.update(t, 30.0 * u.mA),
        times,
    )

assert post_voltage.shape == (2000, 10)
assert conductance.shape == (2000, 10)

Open references/braintools/connectivity.md first when geometry, distance, degree, modules, or another named topology determines which neuron pairs connect; generate the topology before choosing its BrainEvent representation. Open references/projection-patterns.md before implementing recurrent excitation, mutual inhibition, or other spike-driven connectivity; it owns AlignPre versus AlignPost, communication, delays, direct/delta input, short-term plasticity, and electrical coupling. Open skills/brainevent/references/scripts/coba_ei_teaching.py when a BrainPy network should use BrainEvent for efficient event-driven communication; it keeps BrainPy neuron and synapse dynamics while replacing the communication step with BinaryArray @ connectivity over fixed-degree, CSR, or dense storage.

3. Train a spiking network

Training reuses the simulation graph: place a surrogate on each trained spiking nonlinearity, unroll with a differentiable transform loop, differentiate only ParamState, reset dynamical State for each independent batch, and update inside a compiled train step.

# Given a stateful net whose update() returns per-step logits [batch, classes].
params = net.states(brainstate.ParamState)
optimizer = braintools.optim.Adam(lr=3e-3)
optimizer.register_trainable_weights(params)

def loss_fn():
    logits = brainstate.transform.for_loop(net.update, inputs)
    logits = u.math.mean(logits, axis=0)
    return braintools.metric.softmax_cross_entropy_with_integer_labels(
        logits,
        labels,
    ).mean()

@brainstate.transform.jit
def train_step():
    brainstate.nn.init_all_states(net, batch_size=batch_size)
    grads, loss = brainstate.transform.grad(
        loss_fn,
        params,
        return_value=True,
    )()
    optimizer.update(grads)
    return loss

loss = train_step()

Choose training variations from the first-level references below; open only the smallest reference that owns the decision. Use references/scripts/training-snn.py when a complete runnable training workflow is required.

Reference routing

Open only the smallest reference that owns the decision.

Training variation choices

Modeling variation choices

Application script examples