ea-encoder-channel / README.md
rpranav's picture
Publish W26 (repackaged from AlonSaguy/ephys-atlas-models)
23ff516 verified
|
Raw
History Blame Contribute Delete
4.64 kB
metadata
license: cc-by-4.0
library_name: pytorch
tags:
  - neuroscience
  - electrophysiology
  - neuropixels
  - spatial-encoding
  - international-brain-lab

Ephys Atlas spatial encoder (2026_W26)

Predicts the electrophysiological feature vector expected at a channel's position -- from that position plus anatomical context and the recorded features of nearby channels on other insertions. Trained by the International Brain Laboratory on the Ephys Atlas feature release 2026_W26.

The input is a position, not features. This is the inverse of the region classifier: that model takes features and returns a region; this one takes a channel's x, y, z (Allen frame, metres) and returns the 41 ephys features it would expect there. Feeding it the feature columns instead of coordinates is the most common mistake -- predict raises and names the coordinates it needs.

What ships in this repo, and why all of it is needed

A published encoder is not the weights alone. predict cannot run without every one of:

  • spatial_encoder.pt -- the trained network.
  • the context PCA volumes (agea_vol_pca.npy, merfish_vol_pca.npy) -- the anatomical context it samples at each position.
  • neighbor_bank.npz -- the neighbour bank: the position, standardised feature vector and insertion id of every training channel. At inference the model gathers the nearest recorded channels from this bank and attends to them, so the bank is part of the model, the way the stored points are part of a k-nearest-neighbours model. It cannot be reconstructed from the weights, which is why it is shipped here rather than recomputed on your machine.

First run downloads the Allen volume. Building the context sampler constructs an AllenAtlas, which fetches the Allen CCF volume from download.alleninstitute.org (public, no account, a few hundred MB) the first time it runs on a machine, then caches it.

Quickstart

import pandas as pd
from ephysatlas import load_pretrained

model = load_pretrained("int-brain-lab/ea-encoder-channel", revision="2026_W26")
df = pd.read_parquet("example/features_sample.parquet")   # channel positions (x, y, z)
out = model.predict(df)                                    # one pred_<feature> column each
print(out.head())

load_pretrained is the entry point for every ephysatlas model, whatever its family -- it reads ephysatlas_model.json and returns the right wrapper. Use it rather than importing a concrete class, so your code keeps working as the package evolves.

predict returns one row per input channel, indexed identically to the input, with a pred_<feature> column per entry in outputs.columns. The pred_ prefix keeps df.join(out) from colliding with the ground-truth feature columns of the same names.

Inputs

  • Indexed by (pid, channel), one row per recording channel.
  • The coordinate columns x, y, z named in ephysatlas_model.json under inputs.columns, in the Allen atlas frame, in metres. The feature columns are not read -- they are the output.

Neighbourhood selection

For each query position the model attends to recorded channels within 500.0 µm, taking the nearest 8 (selection: nearest). This is deterministic: two calls on the same input return identical predictions.

Training used a random subset of the in-radius neighbours; this published predict takes the nearest instead. The two agree exactly whenever a position has at most 8 neighbours in radius, and differ only in dense regions where the training-time subset would itself have varied run to run. Deterministic output is the right contract for a published model.

Limitations

  • Trained on IBL Neuropixels 1.0 recordings in mouse. Transfer to NP2, other species or other rigs is untested.
  • Coverage follows IBL brain-wide-map targeting; predictions far from any recorded channel fall back on anatomical context alone and are correspondingly weaker.

Reproducibility

Pin the revision. revision="2026_W26" is an immutable tag. Omitting revision resolves to main, which tracks whichever model is currently recommended and will change when a new feature vintage is published -- fine for a first look, not for anything you publish or re-run.

ephysatlas_model.json records the training-time environment and random_seed. Verify your install reproduces the shipped output:

model.selftest()

Citation

Please cite the International Brain Laboratory Ephys Atlas. Model id ea-encoder-channel, feature vintage 2026_W26.