BioBatchNet

A dual-encoder VAE framework for batch effect correction in single-cell RNA-seq and Imaging Mass Cytometry (IMC) data.

Installation

pip install biobatchnet

For development:

git clone https://github.com/UoM-HealthAI/BioBatchNet
cd BioBatchNet
pip install -e .

Quick Start

import scanpy as sc
from biobatchnet import correct_batch_effects

adata = sc.read_h5ad('your_data.h5ad')

bio_emb, batch_emb = correct_batch_effects(
    adata,
    batch_key='BATCH',
    data_type='imc',    # 'imc' or 'seq'
)

adata.obsm['X_biobatchnet'] = bio_emb

# Visualize
sc.pp.neighbors(adata, use_rep='X_biobatchnet')
sc.tl.umap(adata)
sc.pl.umap(adata, color=['BATCH', 'celltype'])

See tutorial.ipynb for an interactive walkthrough.

Full Parameters

bio_emb, batch_emb = correct_batch_effects(
    adata,                      # AnnData object
    batch_key='BATCH',          # obs column for batch labels
    cell_type_key='celltype',   # optional: obs column for cell types
    data_type='imc',            # 'imc' or 'seq' (seq auto-applies HVG + normalize + log1p)
    latent_dim=20,              # latent space dimension
    epochs=100,                 # max training epochs
    lr=1e-4,                    # learning rate
    batch_size=128,             # training batch size
    device='auto',              # 'auto', 'cuda', or 'cpu'
    loss_weights=None,          # optional: override default loss weights
)

Custom Loss Weights

Pass a dict to loss_weights to override defaults:

bio_emb, batch_emb = correct_batch_effects(
    adata,
    batch_key='BATCH',
    data_type='imc',
    loss_weights={
        'discriminator': 0.1,   # batch mixing strength (lower = more mixing)
        'kl_bio': 0.01,         # bio encoder regularization
    },
)

Config-based Training

For reproducible experiments using preset hyperparameters:

python -m biobatchnet.train --config <preset> --data <path_to_h5ad>

Available Presets

Override Parameters

python -m biobatchnet.train --config immucan --data data.h5ad \
    --loss.discriminator 0.1 \
    --loss.kl_bio 0.01 \
    --trainer.epochs 50 \
    --seed 42

Data Format

Input: AnnData object (.h5ad file) with:

• adata.X: expression matrix (cells x features), dense or sparse
• adata.obs[batch_key]: batch labels (string or integer)
• adata.obs[cell_type_key] (optional): cell type annotations

Note: For scRNA-seq data (data_type='seq'), preprocessing is applied automatically: highly variable gene selection (top 2000), normalization (target sum 1e4), and log1p transform. IMC data is used as-is.

Output:

• bio_embeddings: (n_cells, latent_dim) — batch-corrected representations
• batch_embeddings: (n_cells, latent_dim) — batch-specific information

Data

• IMC datasets: Bodenmiller Group IMC datasets
• scRNA-seq datasets: Google Drive

Citation

If you use BioBatchNet in your research, please cite:

Liu H, Zhang S, Mao S, et al. BioBatchNet: A Dual-Encoder Framework for Robust Batch Effect
Correction in Imaging Mass Cytometry[J]. bioRxiv, 2025: 2025.03.15.643447.

License

MIT License