Skip to content
Β 
Β 

Repository files navigation

$scE^2TM$: Toward Interpretable Single-Cell Embedding via Topic Modeling

PyPI version License: MIT

The full description of $scE^2TM$ and its application to published single-cell RNA-seq datasets are available in our paper.


πŸ“– Overview

$scE^2TM$ is a framework for interpretable single-cell embedding learning via topic modeling. It integrates external single-cell foundation-model embeddings with gene expression data and learns sparse topic-gene dependencies for improved interpretability.

Schematic overview of $scE^2TM$

(a) To better collaborate information from different modalities, clusters and topic heads are trained based on mutually refined neighborhood information by encouraging consistent clustering assignments of mutual nearest neighbors of corresponding cells across modalities in the embedding space.

(b) ECR clusters gene embeddings $g_j$ (β€’) as samples and topic embeddings $t_k$ (β˜…) as centers with soft assignment $\pi^{*}_{\epsilon,jk}$.

(c) Sparse linear decoders learn topic embeddings, gene embeddings, and sparse topic-gene dependencies during reconstruction, thus ensuring model interpretability.


πŸ”§ Installation

Note: The complete installation process, including environment setup and dependency installation, typically takes around 1–1.5 hours.

1. Create a conda environment

conda create --name scE2TM_env python=3.8.8 -y
conda activate scE2TM_env

2. Install PyTorch

pip install torch==1.9.1+cu102 torchvision==0.10.1+cu102 torchaudio==0.9.1 \
    -f https://download.pytorch.org/whl/torch_stable.html

3. Install scEΒ²TM

Option A: Install from PyPI (recommended for users)

Hardware Command
CPU pip install scE2TM
GPU pip install scE2TM[gpu]

Option B: Install from source (recommended for developers)

git clone https://github.com/nbnbhwyy/scE2TM.git
cd scE2TM

Then install dependencies based on your hardware:

Hardware Command
CPU pip install -r requirements-cpu.txt
GPU pip install -r requirements-gpu.txt

πŸš€ Quick Start

Label usage: scEΒ²TM is label-free by default. Cell-type annotations are not used for model training. The optional --use_labels flag, or use_labels=True in the Python API, only enables label-dependent evaluation metrics such as ARI, NMI, and Purity when ground-truth annotations are available.

Data format

$scE^2TM$ expects the following input files in CSV format:

  • Gene expression matrix: cell-by-gene matrix (*_HIGHPRE.csv)
  • Cell type annotations: ground-truth labels (*_cell_anno.csv) (optional; used only for evaluation and not for model training)
  • Foundation-model embeddings: pre-trained cell embeddings (*.csv)

We provide the Wang dataset as a default example to help users understand and debug the code.

Run $scE^2TM$

Basic run

python run.py

On the provided Wang example dataset, the demo typically finishes within 1–2 minutes on an NVIDIA RTX 3090 GPU. We also thank the anonymous reviewer for testing the CPU-only workflow on macOS; based on the reviewer’s report, the same example may require approximately 10–15 minutes on CPU-only systems depending on hardware and software configuration.

Specify dataset and number of topics

python run.py --dataset_name Wang --num_topics 50

Choose GPU device (-1 for CPU)

# Run on GPU 0
python run.py --gpu_id 0

# Run on CPU
python run.py --gpu_id -1

Enable label-dependent evaluation metrics (optional)

By default, scEΒ²TM runs without cell-type annotations. Adding --use_labels only computes label-dependent evaluation metrics and does not use labels during model training.

# Compute ARI, NMI, Purity, and other label-dependent metrics
python run.py --use_labels

Full parameter example

python run.py \
    --dataset_name Wang \
    --num_topics 100 \
    --num_neighbors 15 \
    --num_top_genes 10 \
    --tac_weight 1.0 \
    --gpu_id 0 \
    --data_dir ./data \
    --output_dir ./output \
    --use_labels False

Jupyter demo

jupyter notebook scE2TM_demo_on_Wang_dataset.ipynb

Using the high‑level Python API

The package provides a simple function scE2TM() that returns all results directly.

from scE2TM import scE2TM

# Basic label‑free run (no cell type annotations needed)
results = scE2TM(
    dataset_name='Wang',          # dataset name
    data_dir='./data',            # directory containing CSV files
    output_dir='./output',        # where to save outputs
    num_topics=100,               # number of topics (K)
    num_neighbors=15,             # number of neighbors for graph construction
    weight_loss_ECR=20.0,        # weight for the ECR loss
    epochs=500,                   # total training epochs
    gpu_id=0,                     # GPU device; use -1 for CPU
)

# Label-free training with optional label-dependent evaluation
# Cell-type annotations are used only to compute ARI, NMI, and Purity.
# They are not used during model training.
results = scE2TM(
    dataset_name='Wang',
    use_labels=True,  # compute label-dependent metrics only; labels are not used for training
    num_topics=100,
    num_neighbors=15,
    weight_loss_ECR=20.0,
    epochs=500,
    gpu_id=0,
)

# Access the resulting matrices
beta = results['topic_gene_matrix']        
theta = results['cell_topic_matrix']      # latent cell-topic scores (cells Γ— topics)
topic_emb = results['topic_embeddings']
gene_emb = results['gene_embeddings']

Output files

After successful execution, the following files are saved in output/Wang/:

File Description
Wang.pth Trained model checkpoint
Wang_topic_distribution.csv Cell-topic distribution (theta, Latent cell-topic scores (cells Γ— topics))
Wang_topic_embedding.csv Topic embeddings
Wang_gene_embedding.csv Gene embeddings
Wang_tg.csv Topic-gene matrix (beta)

πŸ“š Tutorials

We provide tutorials in the tutorial/ directory covering both basic usage and the main downstream analyses used in the paper.

Tutorial Description
Prepare_foundation_embeddings_scGPT.ipynb Generate foundation-model embeddings using scGPT for input to scE2TM.
Topic_number_selection_with_stability.ipynb Use topic stability, diversity, coherence, and clustering metrics (ARI/NMI, optional) to guide the choice of the optimal number of topics (K).
Clustering and Interpretable Evaluation.ipynb Evaluate clustering performance and interpretability of the learned topics.
Consistency between rare types and topics.ipynb Evaluate how well learned topics capture rare cell populations and their consistency with rare cell types.
Pathway Enrichment.ipynb Perform pathway enrichment analysis on learned topics.
Topic gene embedding.ipynb Visualize and analyze topic-gene embeddings.
Topic perturbation experiment.ipynb Analyze the biological effects of perturbing topic intensities.

βš–οΈ Baseline Tutorials

We also provide tutorials for several relevant baselines in the baseline/ directory:

  • scVI.ipynb – Variational inference for single-cell data
  • scVI-LD.ipynb – scVI with latent Dirichlet allocation
  • scETM.ipynb – Embedded topic model for single-cell data
  • d-scIGM.ipynb – Deep single-cell interpretable generative model
  • baseline_louvain_clustering.ipynb – Louvain clustering on PCA‑reduced expression data
  • baseline_cNMF.ipynb – Consensus NMF topic modeling
  • baseline_SPECTRA.ipynb – SPECTRA factor analysis

These notebooks are designed to help users reproduce baseline results in a consistent environment. After setting up the main $scE^2TM$ environment, each baseline tutorial explains:

  1. any additional package installation required for that baseline,
  2. how to run the method on the provided example dataset, and
  3. how to obtain outputs for comparison with $scE^2TM$.

In general, the workflow is:

conda activate scE2TM_env

Then open the corresponding notebook in baseline/ and follow the dependency installation and execution instructions provided there.


πŸ“ Repository Structure

scE2TM/
β”œβ”€β”€ baseline/      # Baseline tutorials: scVI, scVI-LD, scETM, d-scIGM
β”œβ”€β”€ configs/       # Configuration files
β”œβ”€β”€ data/          # Example datasets and processed inputs
β”œβ”€β”€ models/        # Core model implementations
β”œβ”€β”€ runners/       # Training / running utilities
β”œβ”€β”€ tutorial/      # scE2TM tutorials and downstream analysis notebooks
β”œβ”€β”€ utils/         # Utility functions
β”œβ”€β”€ run.py         # Main entry point
β”œβ”€β”€ requirements.txt
└── LICENSE.txt

πŸ“„ License

This project is licensed under the MIT License. See LICENSE.txt for details.


πŸ“¬ Contact

For questions or support, please open an issue on GitHub or contact:

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages