Pathway-aware neural networks for single-cell gene expression classification. Encodes cells as nodes in a kNN graph, constrains network weights to Reactome pathway hierarchies via binary pruning masks, and supports multi-modality ensemble fusion.
binn/
api.py — binn.BINN: two-persona facade (bio-informed one-liner + computational drop-down)
data/ — Raw files → AnnData → PyG Data objects (Reactome + GO)
pathways.py — Pathway map → binary pruning masks (the biological prior)
nn/ — Encoder architectures + scvi-tools wrappers + donor ensembles
learn/ — Hyperparameters / TrainConfig + interpretability plots
train/ — High-level training entry points (incl. lightweight fit())
data_utils.py — Prostate-paper-specific pipeline helpers (load CSVs, build graphs)
main.py — Reference example: binn.BINN(...) one-liner + manual computational-persona path
citeseq_run.py no longer exists — its role is now split between binn/train/benchmarks.py
(binn_donor_cv*, pseudobulk baselines) and the trials/citeseq/ scripts.
Reactome TSV/GMT files
↓ ReactomeNetwork → get_map()
Pathway map DataFrame (genes × hierarchy levels)
↓ prepare_anndata()
AnnData (.X = expression, .obs = metadata, genes filtered to common set)
↓ anndata_to_graph_data()
PyG Data (x, edge_index, y, train_mask, val_mask, test_mask)
↓ Encoder.generate_masks() [called inside Encoder.__init__]
Binary pruning masks (one tensor per layer transition)
↓ GraphModelFactory.create()
GraphModel (encoder pruned to pathway topology)
↓ GraphModel.train()
Trained weights → predict() / get_latent_representation()
Pathway map — DataFrame from get_map(reactome_net, n_levels)
- Columns:
layer0(genes) …layerN(top pathways), one row per gene ancestry path - Passed to every
Encodersubclass; drivesgenerate_masks()andunitslist
PyG Data object — produced by anndata_to_graph_data()
x (num_cells, num_genes) float32, node features
edge_index (2, num_edges) COO, kNN cell similarity graph
y (num_cells,) int, label-encoded class
train_mask / val_mask / test_mask bool, added by RandomNodeSplit transform
Hyperparameters — binn/learn/hparams.py, threaded through everything
Hyperparameters(
num_node_features, num_classes,
lr, w_decay, dropout, epochs, patience, min_delta,
method, # "ANN" | "GCNConv" | "GATConv"
heads, # attention heads for GAT
bias, save_dir, device
)
hp.for_method("GCNConv") # shallow copy with updated method field
hp.save_path("GCN") # returns save_dir/GCN.ptBoth hierarchies flatten to the same layer0…layerN map shape via shared helpers in
binn/data/_hierarchy.py (add_edges, complete_network, get_nodes_at_level,
get_layers_from_net, get_map) — reactome.get_map and go.get_map are literally the same
function object, so anywhere a Reactome map works, a GO map works too.
| Class/Function | What it does |
|---|---|
Reactome |
Loads three flat files (hierarchy TSV, pathway names TSV, GMT) into DataFrames |
ReactomeNetwork |
Wraps Reactome into a NetworkX DAG; adds synthetic root node |
ReactomeNetwork.get_completed_network(n_levels) |
Pads short branches with _copy dummy nodes so all paths have equal depth |
ReactomeNetwork.get_layers(n_levels) |
Returns list of level dicts; final dict maps pathways → gene symbols |
get_map(reactome_net, n_levels) |
Flattens hierarchy into wide DataFrame (genes × levels); strips _copy artifacts |
GoNetwork (binn.data.go, also binn.data.GoNetwork) |
GO analogue of ReactomeNetwork, one namespace (BP or MF) per instance |
build_go_network(go_base_dir, namespace, n_levels) / build_go_map(...) |
GO analogues of build_reactome_network() / get_map(...); GO data files aren't bundled, so go_base_dir is required |
build_merged_go_map(go_base_dir, ...) |
Concatenates BP + MF maps for a MULGONET-style dual-ontology encoder |
resolve_pathway_map(spec, n_levels, go_namespace, go_base_dir) (binn.data.andata) |
"reactome" / "go"/"go:<namespace>" / a pre-built map DataFrame → a layer0..layerN map; the one place binn.BINN resolves pathway= |
| Function | What it does |
|---|---|
prepare_anndata(data, obs_vars, map) |
Intersects gene columns with map layer0; returns (AnnData, filtered_map) (old name prepareAnnData still works, deprecated). Note: the returned AnnData is built from a bare array, so its var_names are placeholder integers, not gene symbols — recover the real column order via merged_df.columns.intersection(map['layer0']), the same computation binn.api._common_genes uses. |
anndata_to_graph_data(adata, group, adjMat, knn, transform, knn_fn) |
Builds PyG Data; constructs kNN graph via graphtools if no adjMat provided. Without knn_fn, the kNN graph is built from adata.obs columns other than group, not from adata.X — usually not what you want; pass knn_fn=lambda ad: ad.X (or a PCA projection) explicitly. binn.BINN defaults to this for you. |
dataframes_to_mudata(dataframes, obs_vars, pathway_map) |
Multi-modality wrapper; calls prepare_anndata per modality |
mudata_to_graph_data(mdata, group, knn, ...) |
Calls anndata_to_graph_data per modality slot; returns {modality: Data} |
All three encoders inherit from Encoder:
Encoder.__init__(map, args, bias)
→ self.units list of layer widths (unique values per map column)
→ self.masks list of binary float tensors, one per layer transition
→ self.layers nn.Sequential (subclass-specific)
Encoder.forward(x, edge_index=None, *, return_activations=False)
→ logits (a single tensor — behaves like a normal nn.Module)
→ (logits, [per_block_activations]) when return_activations=True
dispatches x-only to Linear/BN layers, (x, edge_index) to GNN layers
ANN ignores edge_index, so it can be called as model(x)
| Class | Layer type | Pruning target | Notes |
|---|---|---|---|
ANN |
PyG Linear + BN |
linear.weight |
Pure MLP, no graph |
GCN |
GCNConv + BN |
conv.lin.weight |
Standard message passing |
GAT |
GATConv(concat=True) + BN |
conv.lin.weight |
Multi-head; mask expanded with torch.block_diag for heads > 1 |
Pruning: torch.nn.utils.prune.custom_from_mask(layer, "weight", mask) — zeros and freezes weights outside pathway connections. Applied once at init; masks are not updated during training.
Persistence gotcha: torch_geometric.nn.Linear (used by ANN directly, and internally as GCNConv.lin/GATConv.lin) overrides _save_to_state_dict to always emit a plain "weight" key, which then fails load_state_dict on a freshly-pruned module expecting weight_orig/weight_mask. Use encoder.load_pruned_state_dict(state_dict) instead of encoder.load_state_dict(state_dict) for any pruned encoder — it handles this transparently and is what binn.api.BINN.save/.load() and EnsembleModel.load() use internally. PoolingDonorEnsemble (donor ensembles) has the same method, delegating each sub-encoder's slice to it.
VCDN — two-layer MLP fusion head: Linear → ReLU → Dropout → Linear. No pathway structure. Takes concatenated encoder logits as input.
GraphModelFactory Registry; call register() once at module load
.create(name, ...) → GraphModel (encoder + scvi wrapper)
.create_encoder(name) → nn.Module (raw encoder, no adata)
.get_train_kwargs(name) (e.g. accelerator="cpu" for GAT)
GraphEncoderModule(BaseModuleClass) Wraps encoder; owns loss fn
GraphTrainingPlan(LightningModule) Full-batch training; reads graph tensors directly
(bypasses scvi mini-batch pipeline)
GraphModel(BaseModelClass) User-facing API
.train(max_epochs, ...) Dummy DataLoader drives Lightning epoch loop
.predict(x, edge_index, mask) Returns softmax probs
.get_latent_representation(...) Last-layer activations before classifier
GAT uses accelerator="cpu" because GATConv's attention mechanism uses scatter-based softmax over neighborhoods (torch_scatter), which is not fully implemented in Apple's MPS backend. On Linux/CUDA this is a no-op penalty — remove it if running on CUDA only.
EnsembleModel — inference + persistence for trained ensembles
save(dir): writesmanifest.json+{modality}.ptper encoder +vcdn.ptload(dir, graph_datas): reconstructs from manifest usingGraphModelFactory.create_encoder
EnsembleTrainingPlan — end-to-end Lightning module; separate Adam param groups for encoders vs. VCDN head.
EncoderRegistry builds plain nn.Module encoders from a registered name ("ANN"/"GCN"/"GAT",
or any custom class registered via EncoderRegistry.register(...)) with no scvi-tools /
PyTorch Lightning imports. GraphModelFactory (in scvi_models.py) subclasses it and adds
.create(), which wraps a built encoder in a GraphModel. Anything that only needs a raw
encoder — donor ensembles, binn.BINN's custom-architecture= path — imports from here, not
from scvi_models, so building/training a pure-PyTorch model never pulls in scvi-tools.
For donor-level classification (cells → pooled donor representations). PoolingDonorEnsemble
(_donor_pooling.py) is the shared "encode cells → pool per donor → fuse" scaffold — it owns
the per-donor pooling loop, the missing-donor zero-vector + presence-mask handling, and
load_pruned_state_dict; DonorAttentionEnsemble and DonorMomentEnsemble are thin subclasses
differing only in pool module and pooled width:
AttentionPool(dim) Learned softmax attention over cells per donor → dim
MomentPool(dim) Parameter-free: concat [mean, var] of cell embeddings → 2*dim
DonorAttentionEnsemble / DonorMomentEnsemble Encoders → pool per modality → fusion MLP
.forward(graph_datas, donor_idx, num_donors) → (num_donors, num_classes)
build_donor_attention_ensemble(names, graph_datas, map_dfs, config, *, num_classes, embed_dim, ...)
build_donor_moment_ensemble(...) same signature — construct without training
build_donor_metadata(obs_adatas, donor_col, label_col)
→ (sorted_donors, donor_to_idx, y_donors, donor_idx_per_modality)
The fusion MLP input is concat(pooled_embeddings, presence_mask) where presence_mask is 1.0 if a donor appears in that modality, 0.0 if not — handles missing modalities gracefully.
| Function | Strategy | Uses Lightning? |
|---|---|---|
train_graph_model(name, graph_data, map_df, config, adata) |
Single modality | Yes |
train_ensemble_model(names, ...) |
Staged: encoders first, then VCDN | Encoders yes, VCDN no |
train_ensemble_model_e2e(names, ...) |
Joint: all weights simultaneously | Yes (EarlyStopping + TensorBoard) |
train_donor_attention_ensemble(names, ...) |
Donor-level attention + fusion | No (manual loop) |
train_donor_moment_ensemble(names, ...) |
Donor-level moment (mean/var) pooling + fusion | No (manual loop) |
Staged vs. E2E trade-off: staged is faster and more stable (each encoder converges independently); E2E allows encoders to co-adapt to the fusion head but is harder to tune (separate lr for encoders, vcdn_lr for fusion).
Donor training splits on donors (not cells) to prevent data leakage, then copies cell-level encoder outputs to donor level via attention or moment pooling. Both train_donor_*_ensemble functions are thin parameterizations of the shared core in _donor_common.py (run_donor_training) — stratified donor split with fallback, GAT/MPS device handling, class-weighted loss, and the manual early-stopping loop live there once, not duplicated per pooling strategy. Both also accept train_edge_index for inductive (leakage-free) training subgraphs.
benchmarks.py — non-BINN baselines (pseudobulk_lr_benchmark, pseudobulk_xgb_benchmark, per_lineage_lr_benchmark, cell_composition_benchmark) plus BINN donor-CV wrappers (binn_donor_cv, binn_donor_cv_inductive, binn_donor_cv_moments), all sharing one generic K-fold runner (_run_donor_cv_generic) parameterized by which train_donor_*_ensemble function and whether to build an inductive training subgraph per fold.
load_and_merge_tables(data_dir) → pd.DataFrame (CNA + response + clinical CSVs)
prepare_graph_data(merged_df, obs_vars, reactome_net, data_dir) → (Data, AnnData, map_df)Thin, dataset-specific glue (main.py's main()/ensemble_main()/e2e_main() use it) — not a
general pipeline entry point. For a general Reactome/GO + AnnData/DataFrame → PyG pipeline use
binn.data.prepare_graph_data_from_map() directly (what binn.BINN itself is built on), or
binn.data.andata.knn_pca_features(adata, n_components) for a PCA-projection knn_fn.
BINN has two user personas (see project README/description) served by three layers, each built directly on top of the one below — a computational user can always drop down a level.
1. binn.BINN facade (binn/api.py) — bio-informed persona, ~3 lines from data to interpretation:
from binn import BINN
model = BINN(adata, pathway="reactome", label_col="condition") # or a dict/MuData of modalities
model.fit()
report = model.interpret() # pathway importance + a Sankey figuredata's shape (single AnnData/DataFrame vs. dict/MuData of modalities) plus donor_col= decide
the backend (GraphModel / EnsembleModel / a donor-pooled ensemble) — see binn/api.py's
module docstring for the full dispatch table. architecture= also accepts an Encoder subclass
(not just a registry name string), so a computational user can plug in a custom architecture and
still get data plumbing, training orchestration, and interpretability for free. model.model is
always the underlying GraphModel/EnsembleModel/donor ensemble, for dropping to layer 2 or 3.
2. PyTorch-native (binn.nn.models + binn.train.loop) — computational persona, explicit kwargs, no config object, forward returns logits:
from binn import GCN, fit # ANN | GCN | GAT
model = GCN(pathway_map, num_classes=5, dropout=0.1, heads=2) # heads: GAT only
fit(model, data, lr=0.05, epochs=1000, patience=20) # or write your own loop
probs = model.predict_proba(data.x, data.edge_index, mask=data.test_mask)fit() (binn/train/loop.py) is a torch-only full-batch loop with real early stopping — it imports neither scvi-tools nor Lightning. TrainConfig (binn/learn/hparams.py) holds the optimisation knobs; architecture goes to the constructor. binn.nn.registry.EncoderRegistry builds encoders by name without importing scvi-tools either.
3. scvi-tools / factory (still fully supported, what layers 1 and 2's ensemble/donor paths are built on) — Hyperparameters + GraphModelFactory.create(...) → GraphModel.train(). Used by all the binn.train.train_* entry points. Encoders still accept a Hyperparameters object as args=.
Top-level names (from binn import GCN, fit, TrainConfig, pathway_masks, BINN) are lazily
imported (PEP 562): binn.nn / binn.train stay cheap for the encoder-only path, but importing
BINN does pull in scvi-tools/Lightning (its default single-modality backend is GraphModel).
Adding a new encoder: subclass Encoder, build self.layers as nn.Sequential from self.units/self.masks/self.num_classes/self.bias, and register it with binn.nn.registry.EncoderRegistry.register("Name", YourClass, args_method="YourMethod") (no scvi-tools import) or GraphModelFactory.register(...) (same effect — they share one registration table). Direct construction (YourClass(pathway_map, num_classes=...)) works without registration; so does BINN(data, ..., architecture=YourClass), which registers it under an auto-generated name the first time it's used.
Custom pathway hierarchy: three options. (a) Produce the same pathway map DataFrame shape (columns layer0…layerN) and pass it to any encoder, or to BINN(data, pathway=your_map_df, ...). (b) Build the pruning masks yourself via binn.pathways.pathway_masks(df) / reactome_masks(net, n_levels) / go_masks(go_net, n_levels) and construct with GCN.from_masks(masks, num_classes=5) — no DataFrame needed. (c) binn.data.andata.resolve_pathway_map("reactome" | "go"/"go:<namespace>" | a map DataFrame, ...) is the single dispatch point BINN's pathway= uses; call it directly if you want the resolved map without going through the facade. binn.data.prepare_graph_data_from_map() accepts a pre-built map directly.
Accessing intermediate activations: model.get_latent_representation() returns the last encoder block output. For per-layer activations call encoder.activations(x, edge_index) (or encoder(x, edge_index, return_activations=True)) — returns [block0_out, block1_out, ...].
Persisting a trained model: BINN.save(directory) / BINN.load(directory, data) round-trips any of the three backends (data only needs the same shape/genes as originally used — the exact per-modality pathway maps and gene order come from disk, not from re-deriving them). At the encoder level, never call a pruned encoder's plain load_state_dict — use encoder.load_pruned_state_dict(state_dict) (see the Pruning persistence gotcha under binn/nn/).