Emulation of the Drosophila Fly Brain

Whole-brain leaky integrate-and-fire model of the adult fruit fly, built from the FlyWire connectome (~138k neurons, ~5M synapses). Activate and silence arbitrary neurons; observe downstream spike propagation.

Based on the paper A leaky integrate-and-fire computational model based on the connectome of the entire adult Drosophila brain reveals insights into sensorimotor processing (Shiu et al.).

Usage

With this computational model, one can manipulate the neural activity of a set of Drosophila neurons. The output of the model is the spike times and rates of all affected neurons.

Two types of manipulations are currently implemented:

• Activation: Neurons can be activated at a fixed frequency to model optogenetic activation. This triggers Poisson spiking in the target neurons. Two sets of neurons with distinct frequencies can be defined.
• Silencing: In addition to activation, a different set of neurons can be silenced to model optogenetic silencing. This sets all synaptic connections to and from those neurons to zero.

The entrypoint is main.py, which parses CLI arguments and calls code/benchmark.py -- the central orchestrator that dispatches to framework-specific runners: run_brian2_cuda.py, run_pytorch.py, run_nestgpu.py, and run_genn.py. The optional Brian2GeNN backend lives in run_brian2_genn.py and uses a separate conda environment because Brian2GeNN 1.7.0 pins Brian2<2.6 while Brian2CUDA uses Brian2 2.8.0.

# Run the 5 main-environment frameworks with default durations (0.1s–1000s)
# and trials (1,4,8,16,32)
python main.py

# Specific durations and trial count
python main.py --t_run 0.1 1 10 --n_run 1

# Single framework
python main.py --nestgpu --t_run 1 --n_run 1
python main.py --genn --t_run 1 --n_run 1
python main.py --brian2genn --t_run 1 --n_run 1

# Combine frameworks
python main.py --brian2-cpu --pytorch --t_run 0.1 1 --n_run 1 4 8 16 32

# Five-round Nature-paper benchmark suite
# Uses the March grid: t_run=(0.1,1,10,100), n_run=(1,4,8,16,32), 5 core backends
python main.py --paper --run-label nature_2026_05

# Add Brian2GeNN as the 6th framework from the brain-fly-brian2genn environment
python main.py --brian2genn --paper --run-label nature_2026_05

Results are incrementally saved to data/benchmark-results.csv as each benchmark completes, with separate columns for setup time (loading, compilation) and simulation time (the always-on cost). For repeated paper runs, the CSV keeps the original March rows and appends new rows keyed by run_label and round; the corresponding spike parquet path is recorded in spike_path.

Spike timing exports are written to parquet outside the timed simulation section so file I/O does not contaminate sim_time. GeNN additionally flushes bounded on-device spike-recording windows during long batched runs; that transfer time is tracked as result collection rather than simulation time. A labeled paper run writes partitioned outputs like:

data/results/nature_2026_05/
├── manifest.csv
├── checksums.sha256
├── round_01/
│   ├── brian2cpp_t1.0s_n1.parquet
│   ├── brian2cuda_t1.0s_n1.parquet
│   ├── pytorch_t1.0s_n1.parquet
│   ├── nestgpu_t1.0s_n1.parquet
│   ├── genn_t1.0s_n1.parquet
│   └── brian2genn_t1.0s_n1.parquet
└── round_02/

Each spike parquet has one row per spike. The canonical timing column for new exports is time_ms, with trial, neuron_index, flywire_id, and exp_name. The legacy t column is kept for existing analysis scripts.

The full nature_2026_05 spike parquet bundle is too large for regular Git tracking, so parquet files are intentionally gitignored. The committed metadata files are manifest.csv and checksums.sha256; the full bundle is stored in Google Drive:

https://drive.google.com/drive/folders/1jiSfb5lNfm9gwP0YyyRz5ATIrDpBAcjs

After downloading the Drive folder into data/results/nature_2026_05/, verify the bundle with:

cd data/results/nature_2026_05
sha256sum -c checksums.sha256

Ground truth comparison

Brian2 (CPU) serves as the ground truth for neural accuracy: it implements the canonical LIF model from Shiu et al. (Nature 2024), which achieved 91% prediction accuracy against experimental Drosophila data. Each backend also saves per-neuron spike trains to data/results/, and a comparison script measures how closely the other backends reproduce Brian2's output:

python code/compare_ground_truth.py                  # default: t_run=1s, n_run=1
python code/compare_ground_truth.py --t_run 10 --n_run 4   # longer / averaged
python code/compare_ground_truth.py --run-label nature_2026_05 --round 1

This computes active-neuron overlap (Jaccard), per-neuron firing-rate correlation, and spike-count ratios, and writes structured results to data/ground-truth-comparison.json.

For all-framework pairwise comparisons, including firing-rate parity rows and spike-time matches within a tolerance window, use:

python code/compare_spike_outputs.py \
  --run-label nature_2026_05 \
  --round 1 \
  --output-dir data/results/nature_2026_05/comparisons

This writes pairwise_summary.csv, pairwise_summary.json, parity_rates.csv, and missing_inputs.json. The pairwise summary has one row per framework pair and t_run/n_run combination.

For paper-support parity files comparing one backend against Brian2 CPU across all five labeled rounds, use:

python code/compare_backend_to_brian2.py \
  --run-label nature_2026_05 \
  --backend brian2genn \
  --output-dir data/results/nature_2026_05/comparisons

This writes <backend>_vs_brian2_rate_summary.csv/json, <backend>_vs_brian2_rate_parity.csv, and <backend>_vs_brian2_missing_inputs.json. Add --include-timing only for smaller targeted checks where greedy spike-time matching is scientifically useful and computationally reasonable.

Installation

Conda environment

The brain-fly conda environment provides everything needed to run the Brian2, Brian2CUDA, PyTorch, NEST GPU, and GeNN backends (including CUDA-enabled PyTorch and PyGeNN):

conda env create -f environment.yml
conda activate brain-fly

On Ubuntu/WSL, PyGeNN's source build also needs the system pkg-config binary and libffi headers:

sudo apt-get install -y pkg-config libffi-dev

GeNN

The --genn backend uses PyGeNN 5.4.0 with the CUDA backend. It implements Brian2-style Poisson activation into membrane voltage, delayed sparse recurrent synapses, GeNN batching for n_run, and the same parquet spike schema as the other benchmark runners.

Large batched GeNN runs cap the on-device spike recording buffer with GENN_RECORDING_WINDOW_MAX_SLOTS (default: 800000). This preserves full spike timing exports while avoiding CUDA out-of-memory errors for large n_run * t_run combinations.

If PyGeNN was not installed when the conda environment was created, install it inside brain-fly with:

export CUDA_PATH=/usr/local/cuda-12.5
export CUDA_HOME=$CUDA_PATH
export PATH=$CUDA_PATH/bin:$PATH
pip install https://github.com/genn-team/genn/archive/refs/tags/5.4.0.zip

Brian2GeNN

The --brian2genn backend uses Brian2GeNN 1.7.0 as a Brian2 standalone device targeting GeNN/CUDA. It is intentionally isolated from the main brain-fly environment because Brian2GeNN pins Brian2<2.6, which conflicts with Brian2CUDA's Brian2 2.8.0 requirement.

Create the environment with:

conda env create -f environment-brian2genn.yml
conda activate brain-fly-brian2genn

Brian2GeNN 1.7.0 expects GeNN 4.x command-line scripts such as genn-buildmodel.sh. If they are not already installed, place GeNN 4.9.0 at ~/.local/src/genn-4.9.0 or set BRIAN2GENN_GENN_PATH/GENN_PATH to your GeNN 4.x source tree:

export CUDA_PATH=/usr/local/cuda-12.5
export CUDA_HOME=$CUDA_PATH
export BRIAN2GENN_GENN_PATH=$HOME/.local/src/genn-4.9.0
export GENN_PATH=$BRIAN2GENN_GENN_PATH
export PATH=$GENN_PATH/bin:$CUDA_PATH/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_PATH/lib64:$LD_LIBRARY_PATH

For scientific comparability, the Brian2GeNN runner exports the same per-spike parquet schema as the other backends and uses the same upstream Poisson drive. Brian2GeNN cannot run this model's independent trials as a true GeNN batch in the way the direct --genn backend can, so n_run>1 is implemented as independent build/run trials with deterministic per-trial C RNG seeds. The sim_time column records GeNN executable time; build_time records the Brian2GeNN code generation/compilation overhead.

NEST GPU

NEST GPU requires a separate build from source with a custom neuron model (user_m1). This is only needed if you want to use the --nestgpu backend.

Prerequisites:

• NVIDIA CUDA Toolkit (12.x) — follow the official installation guide.
• CMake — sudo apt install cmake (or see cmake.org).

Steps:

1. Clone NEST GPU:

git clone https://github.com/nest/nest-gpu

2. Copy the custom source files into the NEST GPU tree. You must replace /path/to/nest-gpu with your own local path:

cp scripts/nestgpu_source_files/src/user_m1.{h,cu}    /path/to/nest-gpu/src/
cp scripts/nestgpu_source_files/pythonlib/nestgpu.py   /path/to/nest-gpu/pythonlib/

The patched nestgpu.py fixes weight array initialization (lines 2225-2227).

3. Build and install (set -DCMAKE_CUDA_ARCHITECTURES to match your GPU, e.g. 89 for RTX 4070):

cmake -DCMAKE_CUDA_ARCHITECTURES=89 \
      -DCMAKE_INSTALL_PREFIX=$HOME/.nest-gpu-build \
      /path/to/nest-gpu
make -j$(nproc) && make install

For a full setup from a fresh Windows machine (WSL2 + CUDA + Miniconda), see scripts/setup_WSL_CUDA.sh.

---

Frameworks

Framework	Backend	Status
Brian2	C++ standalone (multi-core CPU)	ready
Brian2CUDA	CUDA standalone (GPU)	ready
PyTorch	CUDA (GPU)	ready
NEST GPU	CUDA (GPU, custom user_m1 neuron)	ready
GeNN	CUDA (GPU, PyGeNN 5.4.0)	ready
Brian2GeNN	Brian2GeNN 1.7.0 / GeNN CUDA	ready, separate env

All six frameworks share the same data, model parameters, spike-output schema, and folder structure. The five main backends run from brain-fly plus a system-level NEST GPU install; Brian2GeNN runs from brain-fly-brian2genn because of its Brian2 version pin.

Quickstart

# Create the conda environment (includes CUDA-enabled PyTorch)
conda env create -f environment.yml
conda activate brain-fly

# Run a 1-second benchmark on the five main-environment backends
python main.py --t_run 1 --n_run 1 --no_log_file

# Specific backends (combinable)
python main.py --brian2-cpu                    # Brian2 CPU only
python main.py --brian2cuda-gpu               # Brian2CUDA GPU only
python main.py --pytorch                      # PyTorch only
python main.py --nestgpu                      # NEST GPU only
python main.py --genn                         # GeNN only
python main.py --brian2genn                   # Brian2GeNN only, from brain-fly-brian2genn
python main.py --pytorch --genn               # PyTorch + GeNN

# Full benchmark suite (all durations, n_run=1,4,8,16,32, five main backends)
python main.py

# Nature-paper suite: five main backends, March parameter grid, 5 rounds
python main.py --paper --run-label nature_2026_05

# Brian2GeNN Nature-paper add-on from the separate brain-fly-brian2genn env
python main.py --brian2genn --paper --run-label nature_2026_05

main.py options

Flag	Description
(default)	Run all: Brian2 (CPU) → Brian2CUDA (GPU) → PyTorch → NEST GPU → GeNN
--brian2-cpu	Brian2 C++ standalone (CPU) only
--brian2cuda-gpu	Brian2CUDA (GPU) only
--pytorch	PyTorch (GPU/CPU) only
--nestgpu	NEST GPU only
--genn	GeNN CUDA backend only
--brian2genn	Brian2GeNN backend only; use the brain-fly-brian2genn environment
--t_run	Simulation duration(s) in seconds, e.g. --t_run 0.1 1 10
--n_run	Number of independent trials, e.g. --n_run 1 4 8 16 32
--paper	Run the paper suite: t_run=[0.1,1,10,100], n_run=[1,4,8,16,32], 5 rounds
--rounds	Repeat the full selected backend/parameter suite N times
--round-start	First round number to write, useful for resuming a labeled run
--run-label	Group repeated spike outputs under data/results/<label>/ and append labeled CSV rows
--log_file FILE	Write log to file (default: data/results/benchmarks.log)
--no_log_file	Console output only

Backend flags are combinable: --brian2-cpu --pytorch runs Brian2 CPU then PyTorch.

Project structure

fly-brain/
├── main.py                     # Entrypoint (benchmark runner CLI)
├── environment.yml             # Conda env definition (brain-fly)
├── environment-brian2genn.yml  # Separate Brian2GeNN env definition
├── code/
│   ├── benchmark.py            # Orchestrator: config, logging, dispatcher
│   ├── run_brian2_cuda.py      # Brian2 / Brian2CUDA benchmark runner
│   ├── run_pytorch.py          # PyTorch benchmark runner (model + utils)
│   ├── run_nestgpu.py          # NEST GPU benchmark runner (subprocess per trial)
│   ├── run_genn.py             # GeNN/PyGeNN benchmark runner
│   ├── compare_ground_truth.py # Compare backends against Brian2 (CPU) ground truth
│   └── paper-brian2/           # Original paper code (not used by benchmarks)
│       ├── model.py            # Core LIF network model (Brian2)
│       ├── utils.py            # Analysis helpers (load_exps, get_rate)
│       ├── example.ipynb       # Tutorial: activation, silencing, rate analysis
│       └── figures.ipynb       # Reproduce paper figures (uses archive 630 data)
├── data/
│   ├── 2025_Completeness_783.csv       # Neuron list (FlyWire v783)
│   ├── 2025_Connectivity_783.parquet   # Synapse connectivity (FlyWire v783)
│   ├── benchmark-results.csv           # Accumulated benchmark timings
│   ├── ground-truth-comparison.json   # Backend accuracy vs Brian2 (CPU)
│   ├── sez_neurons.pickle              # SEZ neuron subset (for figures)
│   ├── weight_coo.pkl                  # Cached sparse weights COO (gitignored)
│   ├── weight_csr.pkl                  # Cached sparse weights CSR (gitignored)
│   ├── archive/
│   │   ├── 2023_Completeness_630.csv   # Legacy v630 data
│   │   └── 2023_Connectivity_630.parquet
└── scripts/
    └── setup_WSL_CUDA.sh       # WSL2 + CUDA + Miniconda setup

Data

The model uses FlyWire connectome data version 783 (public release). Legacy version 630 data is kept in data/archive/ for paper figure reproduction.

File	Description	Size
2025_Completeness_783.csv	Neuron IDs and metadata	3.2 MB
2025_Connectivity_783.parquet	Pre/post-synaptic indices + weights	97 MB
weight_coo.pkl	Sparse weight matrix (COO), auto-generated by PyTorch	~288 MB
weight_csr.pkl	Sparse weight matrix (CSR), auto-generated by PyTorch	~289 MB

Architecture per framework

	Brian2 / Brian2CUDA	PyTorch	NEST GPU
Build step	C++ / CUDA codegen + compile	None (eager mode)	None
Trial parallelism	Sequential (device.run)	Batched (batch_size=n_run)	Subprocess per trial (cannot reset in-process)
Weight format	Brian2 Synapses object	Sparse CSR tensor	Array-based Connect
Neuron model	Brian2 equations	Custom nn.Module classes	Custom CUDA kernel (user_m1)
Timestep	0.1 ms	0.1 ms	0.1 ms

System requirements

• Linux (tested on Ubuntu 22.04 under WSL2 on Windows 11)
• NVIDIA GPU with CUDA 12.x (tested on RTX 4070)
• Miniconda / Anaconda
• NEST GPU compiled from source (for --nestgpu backend)
• scripts/setup_WSL_CUDA.sh documents the full setup from a fresh Windows machine