neuroSush

Brain-inspired computing on PyTorch, validated against its mathematics.

neuroSush simulates spiking neural networks (leaky, exponential and adaptive integrate-and-fire neurons with dendritic compartments and delays, STDP-family plasticity with dopamine modulation) on a CPU or GPU, in batches. It implements the Thousand Brains and HTM models (sparse distributed representations, spatial pooler, temporal memory, grid cells, active dendrites, voting columns) and hierarchical predictive coding, and it joins the two worlds: a spiking layer that provably computes the temporal memory and learns the same sequences. Every model is tested against what it claims: exact solutions, closed-form probabilities and expectations, and the results of its papers.

Design

• Pure math, thin behaviors. Every equation is a small function that returns new tensors and has tests with hand-computed values; behaviors only hold state and call them.
• Explicit execution order. Each behavior class declares when it runs in a step; there is no global priority table and no special initialization order.
• Reproducible. All randomness comes from the network's seeded generator.
• One dependency. torch only at runtime.
• Safe serialization. Structure specs load through a registry of classes, never eval.
• Validated. Every equation has tests with hand-computed values, and every model is checked against its mathematics (see Validation); examples that learn run in CI.

Install

pip install neurosush

Every GitHub release also carries the wheel and the source archive. From source, with CPU PyTorch (on macOS use pip install torch instead):

python -m venv .venv && source .venv/bin/activate
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install -e ".[dev]"

Python 3.10 or newer.

Quickstart

Twenty Poisson inputs drive two LIF neurons through plastic synapses:

import itertools

import torch

from neurosush.core.network import Network, NeuronGroup, SynapseGroup
from neurosush.data import spike_frames
from neurosush.encoding import rate_poisson
from neurosush.neurons.axon import Axon
from neurosush.neurons.dendrite import DendriteIntegration, DendriteStructure
from neurosush.neurons.inputs import SpikeInput
from neurosush.neurons.models import LIF, Fire
from neurosush.synapses.constraints import WeightClip
from neurosush.synapses.currents import DenseInput
from neurosush.synapses.init import WeightInit
from neurosush.synapses.plasticity import STDP
from neurosush.synapses.traces import SpikeGather, Traces

net = Network(seed=0)
train = rate_poisson(torch.full((20,), 0.3), 100, generator=net.generator)
source = NeuronGroup(net, 20, [SpikeInput(itertools.cycle(spike_frames([(train, None)]))), Axon()])
target = NeuronGroup(
    net,
    2,
    [
        DendriteStructure(),  # routes synaptic currents to the soma
        DendriteIntegration(),
        LIF(tau=10.0, threshold=-55.0, v_reset=-70.0, v_rest=-65.0),
        Fire(),
        Axon(),
    ],
)
synapse = SynapseGroup(
    net,
    source,
    target,
    [
        WeightInit(mode="uniform", scale=0.5),
        DenseInput(coef=5.0),
        SpikeGather(),
        Traces(tau_pre=10.0),
        STDP(a_plus=0.02, a_minus=0.01, bound="soft"),
        WeightClip(w_min=0.0, w_max=1.0),
    ],
)

spikes = 0
for _ in range(200):
    net.step()
    spikes += int(target.spikes.sum())
print(f"output spikes: {spikes}, mean weight: {synapse.weights.mean():.3f}")

examples/two_patterns.py goes further: two output neurons with k-winners-take-all and homeostasis learn to respond to one input pattern each.

Concepts

• A Network owns neuron groups, synapse groups, the time step dt, the dtype, the device and a seeded random generator.
• A NeuronGroup has a shape (depth, height, width) (an int n means (1, 1, n)), and a SynapseGroup connects two groups and targets one dendritic compartment of the destination: proximal (drives the soma), distal or apical (prime it).
• A Behavior holds state and dynamics. Behaviors run in this order in every step (and are initialized in the same order):

• Spikes travel through SpikeGather: a synaptic input reads syn.pre_spike, which SpikeGather fills from the source's Axon, so every synapse group with an input needs both; building without them is an error, not a silent network.
• Delays are integers in steps: src_delay is axonal (read from the source's Axon history), dst_delay is dendritic (arrival in the destination's DendriteStructure).
• Units: every time constant shares the unit of dt (milliseconds by convention).
• Signs: weights are magnitudes; NeuronGroup(..., inhibitory=True) makes its outgoing currents negative.
• Batches: Network(batch_size=B) simulates B samples side by side. State tensors get shape (B, size); weights and thresholds stay shared, and learning uses the batch mean. Feed it with spike_frames(samples, batch_size=B). On a GPU a batch costs about as much as one sample, so throughput grows almost linearly with B.

More in docs/ARCHITECTURE.md.

Recording and checkpoints

A Recorder copies attributes of its host at the end of every interval-th step. A checkpoint saves the complete state of an initialized network (iteration, random generator, every tensor and delay buffer, and behavior state such as homeostasis), and loading it into a network built the same way continues the run exactly:

import tempfile
from pathlib import Path

import torch

from neurosush import checkpoint
from neurosush.core.network import Network, NeuronGroup
from neurosush.neurons.competition import InherentNoise
from neurosush.neurons.dendrite import DendriteIntegration, DendriteStructure
from neurosush.neurons.models import LIF, Fire
from neurosush.recording import Recorder


def build():
    net = Network(seed=0)
    NeuronGroup(
        net,
        3,
        [
            DendriteStructure(),
            DendriteIntegration(),
            LIF(tau=10.0, threshold=-55.0, v_reset=-70.0, v_rest=-65.0),
            InherentNoise(scale=8.0),  # random drive from the network generator
            Fire(),
            Recorder("v", "spikes", interval=5),
        ],
    )
    return net


net = build()
net.run(50)
recorder = net.groups[0].behaviors[-1]
print(recorder.get("v").shape, recorder.steps[:3])  # torch.Size([10, 3]) [5, 10, 15]

path = Path(tempfile.mkdtemp()) / "net.pt"
checkpoint.save(net, path)
net.run(20)

resumed = build()
checkpoint.load(resumed, path)
resumed.run(20)
assert torch.equal(resumed.groups[0].v, net.groups[0].v)  # continues exactly

Loading uses torch.load(weights_only=True), so a checkpoint cannot run code. Input streams (SpikeInput) are not part of it: resume them yourself.

Modules

Structures and specs

A column is plain data. Building it twice gives two independent columns; JSON only names registered classes, so loading it never runs code.

from neurosush.core.network import Network
from neurosush.structure.spec import (
    BehaviorSpec,
    ColumnSpec,
    GroupSpec,
    LayerSpec,
    SynapseSpec,
    build_column,
    from_json,
    to_json,
)

lif = BehaviorSpec("LIF", {"tau": 10.0, "threshold": -55.0, "v_reset": -70.0, "v_rest": -65.0})
neurons = (
    BehaviorSpec("DendriteStructure"),
    BehaviorSpec("DendriteIntegration"),
    lif,
    BehaviorSpec("Fire"),
    BehaviorSpec("Axon"),
)
dense = (
    BehaviorSpec("WeightInit", {"mode": "uniform"}),
    BehaviorSpec("DenseInput"),
    BehaviorSpec("SpikeGather"),
)
layer = LayerSpec(
    groups={"exc": GroupSpec(8, neurons), "inh": GroupSpec(2, neurons, inhibitory=True)},
    synapses=(SynapseSpec("exc", "inh", dense), SynapseSpec("inh", "exc", dense)),
    outputs={"out": ("exc",)},
)
spec = ColumnSpec(
    layers={"L4": layer, "L23": layer},
    synapses=(SynapseSpec("L4.exc", "L23.exc", dense),),  # "layer.group" inside a column
    inputs={"in": "L4.in"},
    outputs={"out": "L23.out"},
)

assert from_json(to_json(spec)) == spec
net = Network(seed=0)
column = build_column(net, "C1", spec)
net.run(10)
print([group.name for group in column.output_port("out")])  # ['C1.L23.exc']

Thousand Brains models

neurosush.htm implements the algorithms of hierarchical temporal memory and the Thousand Brains theory as tensor code, each checked against the mathematics of its paper:

import torch

from neurosush.htm.encoders import CategoryEncoder
from neurosush.htm.sdr import match_probability
from neurosush.htm.temporal_memory import TemporalMemory

# chance that a random 40-of-2048 SDR shares 20 or more bits with a stored one
print(f"{match_probability(2048, 40, 40, 20):.1e}")

a, b, c, d = CategoryEncoder(256, 12, 4, seed=0).encode(torch.arange(4))
tm = TemporalMemory(
    256,
    8,
    activation_threshold=8,
    min_threshold=6,
    initial_permanence=0.51,
    max_new_synapses=12,
    max_synapses_per_segment=16,
)
for _ in range(3):
    tm.reset()
    for symbol in (a, b, c, d):
        tm.compute(symbol)
tm.reset()
tm.compute(a, learn=False)
tm.compute(b, learn=False)
assert torch.equal(tm.predicted_columns(), c)  # after A B it expects C

examples/sequence_prediction.py chains an encoder, the spatial pooler, temporal memory and a classifier on sequences that differ only in their first symbol; examples/object_recognition.py shows voting columns recognizing objects in fewer touches.

Spiking temporal memory

The same sequence memory also runs as a spiking network. ActiveSegments gives every cell distal segments that fire a dendritic spike when enough of their synapses see input within a coincidence window, and hold a plateau that primes the cell below threshold. MinicolumnInhibition lets the first cells of a minicolumn to reach threshold silence the rest. A primed (predicted) cell therefore fires alone, and a minicolumn without one fires all at once (a burst). tests/validation/test_sequence_math.py checks that such a layer activates exactly the cells TemporalMemory activates, element by element, under timing conditions it also checks.

SegmentLearning adds the temporal memory's learning in spike time (reinforcement, growth towards the previous element's winners, new segments in bursting minicolumns, punishment of wrong predictions), and sequence_memory builds the whole layer. It derives the plateau, coincidence window and learning context from the neuron's exact race times and refuses parameters under which the equivalence would fail. On sequences that share their middle it learns the same curves as TemporalMemory and ends in the same state (tests/validation/test_sequence_learning.py, and experiments/sequence_learning.py for the full curves).

A four-element sequence, learned in one repetition (initial_permanence=0.51 makes new synapses connected at once):

import torch

from neurosush.core.behavior import Behavior
from neurosush.core.network import Network, NeuronGroup
from neurosush.core.order import Order
from neurosush.htm.sdr import random_sdr
from neurosush.neurons.axon import Axon
from neurosush.structure.sequence import sequence_memory

COLUMNS, PERIOD, WINDOW = 32, 60, 25
sequence = random_sdr(COLUMNS, 6, batch=(4,), generator=torch.Generator().manual_seed(0))


class Present(Behavior):
    """Drive the columns of element t // PERIOD for the first WINDOW steps of its period."""

    order = Order.FIRE

    def initialize(self, group):
        group.spikes = group.state(False, dtype=torch.bool)

    def forward(self, group):
        t = group.net.iteration - 1
        element, offset = divmod(t % (PERIOD * 6), PERIOD)  # 4 elements, then 2 silent
        on = element < len(sequence) and offset < WINDOW
        group.spikes = sequence[element].clone() if on else group.state(False, dtype=torch.bool)


net = Network(seed=0)
columns = NeuronGroup(net, COLUMNS, [Present(), Axon()])
layer, _ = sequence_memory(
    net, columns, cells_per_column=4, period=PERIOD, window=WINDOW, initial_permanence=0.51
)
for repetition in range(2):
    bursts = []
    for element in range(6):
        fired = torch.zeros(layer.size, dtype=torch.bool)
        for _ in range(PERIOD):
            net.step()
            fired |= layer.spikes
        if element < len(sequence):
            active = fired.view(COLUMNS, 4)[sequence[element]]
            bursts.append(int(active.all(-1).sum()))  # columns that were not predicted
    print(repetition, bursts)  # [6, 6, 6, 6], then [6, 0, 0, 0]: only the first element surprises

Validation

Besides unit tests, tests/validation checks the spiking core against the mathematics it implements. The LIF reproduces the exact solution of its Euler scheme and converges to the continuous solution at first order in dt; interspike intervals match the closed form, and the exponential LIF starts firing exactly at its rheobase R I = theta_rh - v_rest - delta. Traces respond to a spike with exactly (1 - dt/tau)^k; the STDP window is the exponential a_plus c^k or -a_minus c^k, and uncorrelated spike trains drift the weights by the expected amount. Inhibitory STDP settles the firing rate at its target, homeostasis settles the spike count, Poisson spike counts are binomial with geometric intervals, and every spike arrives exactly src_delay + dst_delay + 1 steps after it was fired. Conv, local, lateral and pooling synapses equal their dense synapse-by-synapse definition, for both currents and STDP; reward-modulated STDP matches the closed form of the three-factor rule and solves the distal reward problem of Izhikevich (2007).

Limitations

• Alpha. The API may still change between minor versions.
• Speed. A simulation step is a sequence of small tensor operations driven from Python: a few milliseconds per step on a CPU, whatever the network size up to thousands of neurons. Batches spread that cost over many samples, especially on a GPU. The spatial pooler learns one sample at a time.
• CPU-only HTM models. The neurosush.htm models and SegmentLearning run unbatched on the CPU.
• Checkpoints save a network's state but not the input streams feeding SpikeInput.
• The spiking temporal memory equals the algorithm only under the timing conditions that sequence_timing checks; sequence_memory refuses other parameters.

Citing

If you use neuroSush in research, please cite it with the metadata in CITATION.cff (GitHub shows it under "Cite this repository"), and cite the papers of the models you use; each module names them.

Development

bash scripts/check.sh runs ruff (lint and format check), strict mypy and the full test suite. python benchmarks/suite.py measures throughput (the Benchmarks workflow runs it every week), and the scripts in experiments/ reproduce the numbers behind the validation claims. Tests marked gpu compare CUDA with CPU results and run only where a CUDA device exists. See CONTRIBUTING.md for the workflow and code standards.

Releases

Releases are automatic. A commit on main that sets a new release version with its changelog section triggers the Release workflow: it checks the version and the changelog, runs the full CI, creates the tag and a GitHub release with the wheel and the sdist, and publishes them to PyPI. The steps are in CONTRIBUTING.md; the history is in CHANGELOG.md.

License

MIT, see LICENSE.