Skip to content

Repository files navigation

BINN — Biologically informed Neural Networks

Pathway-aware neural networks for gene expression classification and regression. BINN allows learned connections from gene → pathway to be exploited from ontology datasets (Reactome and GO), with support for multi-modal ensembling.

Encoders are implemented as scvi-tools modules bringing PyTorch Lightning training loops (checkpointing, early stopping, GPU/multi-GPU handling), a consistent model.train() / .predict() / .get_latent_representation() API shared with the wider scvi-tools ecosystem.

Install

pip install -e .

Requires Python ≥3.10. Core dependencies are declared in pyproject.toml.

The three stages

BINN pipelines always follow the same shape: preprocess data → build a pathway map → train a model.

1. Preprocessing — tabular data → graph

Raw expression data (a cells/samples × genes DataFrame or AnnData) is turned into a PyTorch Geometric Data object: a kNN graph over samples/cells, with expression as node features and labels as y.

from binn.data import prepare_anndata, anndata_to_graph_data
from torch_geometric.transforms import RandomNodeSplit

# `merged_df`: rows = samples, columns = genes + metadata (e.g. "response")
# `pathway_map`: from stage 2 below — genes present here are kept, the rest dropped
adata, map_df = prepare_anndata(
    data=merged_df,
    obs_vars=["response"],          # metadata columns to carry into adata.obs
    map=pathway_map,
)

data = anndata_to_graph_data(
    adata=adata,
    group="response",               # obs column used as the classification label
    knn=3,                          # neighbours per node in the similarity graph
    transform=RandomNodeSplit(split="random", num_val=200, num_test=314),
)

For multi-modality data (e.g. RNA + protein), use dataframes_to_mudata / mudata_to_graph_data instead — they run the same steps per modality and return one Data object per modality, keyed by name.

2. Pathway creation — Reactome hierarchy → mask map

The pathway map is what drives pruning: it tells the encoder which genes feed into which pathways, and which pathways roll up into which parent pathways.

from binn.data import ReactomeNetwork, get_map

reactome_net = ReactomeNetwork(
    reactome_base_dir="train_data/reactome",
    relations_file_name="ReactomePathwaysRelation.txt",
    pathway_names="ReactomePathways.txt",
    pathway_genes="ReactomePathways.gmt",
)

pathway_map = get_map(reactome_net, n_levels=3)
# DataFrame: layer0 (genes) ... layerN (top-level pathways),
# one row per gene→pathway ancestry path

Any hierarchy source works as long as it produces a DataFrame with the same layer0 … layerN shape — swap in a custom ontology and everything downstream (masking, encoder construction) is unaffected.

3. Model training — base encoders and ensembles

Hyperparameters is the single config object threaded through every training entry point:

from binn.learn import Hyperparameters

config = Hyperparameters(
    num_node_features=data.num_node_features,
    num_classes=int(data.y.unique().numel()),
    lr=1e-3,
    epochs=400,
    patience=200,
    save_dir="weights",
)

Base model — one pathway-pruned encoder ("ANN", "GCN", or "GAT"):

from binn.train import train_graph_model

model = train_graph_model("GCN", data, map_df, config, adata=adata)
preds = model.predict(mask=data.test_mask)

Ensemble — one encoder per modality, fused with a VCDN head. Staged training (encoders first, fusion head second) is faster and more stable; end-to-end lets encoders co-adapt to the fusion head but needs more careful tuning:

from binn.train import train_ensemble_model, train_ensemble_model_e2e

names = {"rna": "GCN", "protein": "GAT"}   # {modality: architecture}

ensemble = train_ensemble_model(
    names, graph_datas, map_dfs, config,   # graph_datas/map_dfs keyed by modality
    vcdn_hidden_dim=128,
)
# or, jointly:
ensemble = train_ensemble_model_e2e(names, graph_datas, map_dfs, config)

Donor-level ensemble — pools cell-level encoder outputs per donor via learned attention before classification (splits on donors, not cells, to avoid leakage):

from binn.train import train_donor_attention_ensemble

model = train_donor_attention_ensemble(names, graph_datas, map_dfs, config)

Interpreting results

from binn.learn import plot_confusion_matrix, plot_pca_latent, plot_last_layer_weights

plot_confusion_matrix(model, data, data.test_mask)
plot_pca_latent(model, data)
plot_last_layer_weights(model)

model.get_latent_representation() returns the last encoder block's activations. For per-layer activations, call the encoder directly — model.module.encoder(x, edge_index) returns (output, [block0_out, block1_out, ...]).

About

Training Lib for BINN

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages