README.md

metadata

license: cc-by-4.0
library_name: pytorch
tags:
  - neuroscience
  - electrophysiology
  - neuropixels
  - unit-level-encoder
  - international-brain-lab

Ephys Atlas unit-level interpolation model (2026_W42)

Predicts which kinds of neurons you expect to record at a brain position, and their expected waveform phenotype. Trained by the International Brain Laboratory on spike-sorted units of the Ephys Atlas release 2026_W41 (56497 training units).

The model has four stages:

1. a multimodal autoencoder embeds each unit's multi-channel waveform, 3D autocorrelogram and population coupling (waveform, acg, stpc) into a 60-d latent -- its phenotype;
2. a hierarchical Gaussian mixture over the standardised latents: 30 putative cell types, each split into at most 10 sub-types (300 in all). Assignment is top-down -- the most likely type, then the most likely of its sub-types -- so a unit's sub-type always lies inside its type. The types and sub-types are global;
3. the sub-type weights depend on the molecular and cellular context (PCA volumes of the molecular and cellular atlases) at the position: one network predicts the type weights, a second the shares of the sub-types within each type;
4. a members readout maps the model back to 10 interpretable waveform features: depolarisation_slope, recovery_slope, repolarisation_slope, spatial_spread_um, tip_val, spike_width_secs, predepolarisation_width_secs, spike_amplitude, peak_to_trough_ratio_log, polarity. Each sub-type is represented by its training units, from anywhere in the brain, weighted by their posterior of the sub-type.

The input is a position. predict takes x, y, z (IBL/Allen frame, metres) and returns the expected phenotype there: the local sub-type weights times each sub-type's expected features. mixture_weights returns the local putative cell-type composition itself.

Quickstart

import pandas as pd
from ephysatlas import load_pretrained

model = load_pretrained("int-brain-lab/ea-encoder-unit", revision="2026_W42")
df = pd.read_parquet("example/positions_sample.parquet")   # positions (x, y, z)
out = model.predict(df)                                     # one pred_<feature> column each
types = model.mixture_weights(df, level="type")             # putative cell-type composition
subtypes = model.mixture_weights(df)                        # ... and of the sub-types
draws = model.sample(df, n_samples=16)                      # training units drawn at each position
model.selftest()                                            # reproduces the shipped outputs

Units you recorded yourself can be embedded with model.encode(waveform, acg, stpc) (arrays shaped as in config.json), then assigned to a putative cell type and sub-type with model.assign(z).

What ships here

• autoencoder.pt, shared_latent_scaler.joblib -- the phenotype encoder and its latent scaler.
• types_gmm.joblib, subtypes_gmm.joblib, subtype_parent.npy -- the types, the sub-types and the type of each sub-type.
• context_transform.joblib, context_weights_types.pt, context_weights_subtypes.pt -- context -> type and sub-type weights.
• members.npz -- the phenotype features of the training units and their sub-type posteriors (the readout; part of the model, like the stored points of any nearest-neighbour model).
• agea_vol_pca.npy, merfish_vol_pca.npy -- the molecular and cellular context volumes, the same frozen volumes as the channel-level model int-brain-lab/ea-encoder-channel of release 2026_W41.
• split.json, config.json, preprocessing/unit_stats.npz, results/ -- the probe split, the training configuration, the training-data statistics and held-out evaluation summaries.

The recorded per-unit dataset (waveforms, ACGs, positions) is not here: it stays on IBL S3 and is only needed to retrain the model or to reproduce the atlas-wide analyses.

First run downloads the Allen volume. Sampling the context constructs an AllenAtlas, which fetches the Allen CCF volume (public, no account, a few hundred MB) once per machine.

Reproducibility

Pin the revision. revision="2026_W42" is an immutable tag; omitting it resolves to main, which tracks the currently recommended model. ephysatlas_model.json records the training-time environment and random seed; model.selftest() verifies your install reproduces the shipped output (no IBL account needed).

Citation

Please cite the International Brain Laboratory Ephys Atlas. Model id 2026_W42_unit, vintage 2026_W42.