Point-neuron networks with brainpy.state

brainpy.state provides point-neuron and synapse building blocks for stateful, JAX-native spiking network simulations. It is designed to work with brainstate transformations and brainunit quantities, so the same model can keep explicit units, reusable states, and efficient loop execution.

In this tutorial, you will learn how to:

• create a leaky integrate-and-fire population with brainpy.state.LIF;

• initialize and inspect stateful modules;

• read the four parts of a projection: comm, syn, out, and post;

• build a compact excitatory-inhibitory network with AlignPostProj;

• run the network and visualize spikes as a raster plot.

Imports and simulation clock

Most brainpy.state models are unit-aware. Set the global time step once before creating and running the model.

import numpy as np
import matplotlib.pyplot as plt

import brainunit as u
import brainstate
import braintools
import brainpy.state

brainstate.environ.set(dt=0.1 * u.ms)
dt = brainstate.environ.get_dt()
print("dt =", dt)

1. Create a LIF population

brainpy.state.LIF is a stateful neuron population. The population keeps membrane potential and spike state internally, while the update call advances the dynamics by one time step.

The example below creates a small population of independent LIF neurons. The parameters use physical units from brainunit, which makes the equations easier to read and helps catch dimension mistakes.

num_neurons = 16

lif = brainpy.state.LIF(
    num_neurons,
    V_rest=-60.0 * u.mV,
    V_th=-50.0 * u.mV,
    V_reset=-65.0 * u.mV,
    tau=10.0 * u.ms,
    V_initializer=braintools.init.Normal(-60.0, 3.0, unit=u.mV),
    spk_reset="soft",
)

_ = brainstate.nn.init_all_states(lif)
lif

2. Drive the LIF population

A model update is just a Python function. brainstate.transform.for_loop runs that function over a sequence and records the returned values. The state of lif is carried from one iteration to the next.

def run_lif_population(inputs):
    lif.reset_state()
    indices = np.arange(len(inputs))

    def step(i, current):
        with brainstate.environ.context(i=i):
            spike = lif(current)
            return lif.V.value, spike

    voltages, spikes = brainstate.transform.for_loop(step, indices, inputs)
    return voltages, spikes

n_steps = 600
constant_current = 3.0 * u.mA
inputs = np.ones(n_steps) * constant_current
voltages, spikes = run_lif_population(inputs)

times = np.arange(n_steps) * dt
print("voltages:", voltages.shape)
print("spikes:", spikes.shape)

plt.figure(figsize=(8, 4))
plt.plot(times, voltages[:, 0], label="neuron 0")
plt.plot(times, voltages[:, 1], label="neuron 1", alpha=0.8)
plt.axhline(-50.0 * u.mV, color="tab:red", ls="--", lw=0.8, label="threshold")
plt.xlabel("time")
plt.ylabel("membrane potential")
plt.title("LIF membrane potentials under constant current")
plt.legend(loc="best")
plt.tight_layout()
plt.show()

3. Projection anatomy

A projection connects a presynaptic event stream to a postsynaptic target. In brainpy.state, projections are built from explicit pieces:

• comm: communication or connectivity, such as brainstate.nn.EventFixedProb;

• syn: synaptic dynamics, such as brainpy.state.Expon;

• out: output conversion, such as brainpy.state.CUBA for current-based output;

• post: the postsynaptic population that receives the output.

This decomposition lets you swap the connectivity, synapse, or output model without rewriting the whole network. It also makes the central projection question explicit: where should the synaptic state live?

In a naive implementation, every realized connection can own its own synaptic state. For a sparse recurrent network, that means the state size scales with the number of connected pairs. This is often much larger than the number of neurons, and it becomes expensive when the same synaptic dynamics are repeated across many similar connections.

brainpy.state avoids this by aligning synaptic state to a neuron dimension instead of a connection dimension. The projection still computes the same synaptic drive, but the temporal state is stored either on the presynaptic side or on the postsynaptic side.

Figure: AlignPre vs AlignPost. The synapse-dynamics block sits before the connection matrix on the left (state aligned to the presynaptic population, , works for all synapses, natural for one-to-many) and after it on the right (state aligned to the postsynaptic population, , exponential-family synapses, natural for many-to-one). (Orange = presynaptic, blue = postsynaptic, green = connection matrix.)

AlignPre aligns synaptic state with the presynaptic population. Conceptually, each presynaptic neuron maintains its own trace, such as an exponentially decaying conductance or a more complex synaptic variable. After that trace is updated, comm distributes it through the connectivity and weights to the targets. This is useful when one source projects to many targets, or when the synapse model has dynamics that must be evaluated before the signal is communicated. Because the synapse state is attached to presynaptic neurons, memory scales with the presynaptic population size rather than the number of connections.

AlignPost aligns synaptic state with the postsynaptic population. Here, presynaptic events are first communicated and accumulated for each target neuron. The synapse dynamics then evolve on the postsynaptic side, with one state per postsynaptic neuron for each projection. This is useful when many presynaptic sources converge onto the same target population. For exponential-family synapses, the summed incoming events can be merged before the synaptic state update, so the result is exact while memory scales with the postsynaptic population size.

A practical way to choose between them is to ask where repeated work can be shared:

• use AlignPre when one presynaptic source is reused across many targets, or when a general synapse model is naturally updated before communication;

• use AlignPost when many inputs converge onto the same target population, especially for exponential current- or conductance-based synapses;

• keep the four roles separate: comm moves spikes or traces across connectivity, syn stores and updates temporal state, out converts synaptic variables into current or conductance, and post receives the result.

The E/I network below uses AlignPostProj because both excitatory and inhibitory presynaptic slices project back to the same LIF population. Each projection can therefore keep its synaptic state aligned to the target population size, while the two projections still use different weights and time constants.

4. Build an excitatory-inhibitory network

The network below contains one mixed LIF population. The first n_exc neurons are excitatory and the remaining n_inh neurons are inhibitory. Two post-aligned projections send spikes from each subpopulation back to the full population.

AlignPostProj is a good default when several presynaptic sources converge onto the same postsynaptic group, because the synaptic state is aligned to the target population.

class EINet(brainstate.nn.Module):
    """A compact recurrent E/I network built from brainpy.state components."""

    def __init__(self, n_exc, n_inh, prob, exc_weight, inh_weight):
        super().__init__()
        self.n_exc = n_exc
        self.n_inh = n_inh
        self.num = n_exc + n_inh

        self.neurons = brainpy.state.LIF(
            self.num,
            V_rest=-52.0 * u.mV,
            V_th=-50.0 * u.mV,
            V_reset=-60.0 * u.mV,
            tau=10.0 * u.ms,
            V_initializer=braintools.init.Normal(-60.0, 5.0, unit=u.mV),
            spk_reset="soft",
        )

        self.exc = brainpy.state.AlignPostProj(
            comm=brainstate.nn.EventFixedProb(n_exc, self.num, prob, exc_weight),
            syn=brainpy.state.Expon.desc(self.num, tau=2.0 * u.ms),
            out=brainpy.state.CUBA.desc(),
            post=self.neurons,
        )
        self.inh = brainpy.state.AlignPostProj(
            comm=brainstate.nn.EventFixedProb(n_inh, self.num, prob, inh_weight),
            syn=brainpy.state.Expon.desc(self.num, tau=5.0 * u.ms),
            out=brainpy.state.CUBA.desc(),
            post=self.neurons,
        )

    def update(self, background_current):
        previous_spikes = self.neurons.get_spike() != 0.0
        self.exc(previous_spikes[:self.n_exc])
        self.inh(previous_spikes[self.n_exc:])
        self.neurons(background_current)
        return self.neurons.get_spike()

5. Initialize and inspect the network

The connection weights are scaled by population size. This keeps the total recurrent input in a reasonable range when the connection probability or population size changes.

n_exc = 320
n_inh = 80
prob = 0.05

exc_weight = 1.0 / u.math.sqrt(prob * n_exc) * u.mS
inh_weight = -1.2 / u.math.sqrt(prob * n_inh) * u.mS
background_current = 3.0 * u.mA

net = EINet(n_exc, n_inh, prob, exc_weight, inh_weight)
_ = brainstate.nn.init_all_states(net)
net

6. Run the recurrent simulation

Each time step uses the previous spikes to update recurrent synaptic input, then advances the neuron population. The returned array is a time-by-neuron spike matrix.

duration = 300.0 * u.ms
times = u.math.arange(0.0 * u.ms, duration, brainstate.environ.get_dt())

spike_history = brainstate.transform.for_loop(
    lambda t: net.update(background_current),
    times,
    pbar=brainstate.transform.ProgressBar(20),
)

print("spike_history:", spike_history.shape)

t_indices, neuron_indices = u.math.where(spike_history)

plt.figure(figsize=(8, 4))
plt.scatter(times[t_indices], neuron_indices, s=1, color="black")
plt.axhline(n_exc - 0.5, color="tab:red", lw=0.8, label="E/I boundary")
plt.xlabel("time")
plt.ylabel("neuron index")
plt.title("Spike raster from a brainpy.state E/I network")
plt.legend(loc="upper right")
plt.tight_layout()
plt.show()

7. What to change next

Try these small modifications to understand the components:

1. Change the excitatory tau in brainpy.state.Expon.desc(...) from 2.0 * u.ms to 5.0 * u.ms.

2. Increase prob and reduce the weights to compare dense weak connectivity with sparse strong connectivity.

3. Change CUBA.desc() to a conductance-based output in a model where the postsynaptic target exposes the voltage state needed by that output.

4. Split the excitatory group into two projections with different time constants, then compare the raster.

Because the network is stateful, rebuild or reset the states before each independent experiment.

Troubleshooting

• No spikes: increase background_current, increase exc_weight, or simulate for a longer duration.

• Runaway activity: reduce exc_weight, strengthen inh_weight, or shorten the excitatory synaptic time constant.

• Shape errors: make sure EventFixedProb(n_pre, n_post, ...) matches the presynaptic slice and postsynaptic population. In this tutorial, the synaptic state uses self.num because AlignPostProj aligns it to the target population.

• Unit errors: keep time constants in u.ms, voltages in u.mV, currents in u.mA, and conductance weights in u.mS.

• Unexpected repeated results: call brainstate.nn.init_all_states(model) after constructing a model, and reset states between independent runs when needed.